LLM - Baichuan7B Lora 训练详解

一.引言

二.环境准备

三.模型训练

1.依赖引入与 tokenizer 加载

2.加载 DataSet 与 Model

3.Model 参数配置

4.获取 peft Model

5.构造 Trainer 训练

6.训练完整代码

四.Shell 执行

1.脚本构建

2.训练流程

3.训练结果

五.总结

一.引言

LLM - Baichuan7B Tokenizer 生成训练数据上文我们介绍了如何将 QA 式的样本构造为训练可用的 DataSet，本文我们基于该训练数据进行 Baichuan7B 的 lora 微调。本文也会尽量把每一个 API 每一个参数的含义大致搞懂，以供自己和大家学习了解。

二.环境准备

• 主要依赖

python 3.9.11
numpy==1.23.5
torch==2.0.1
transformers==4.29.1

可以将上述依赖版本放入 requirements.txt。

• 激活并配置环境

conda create -n baichuan python=3.9
conda activate baichuan

pip install -r requirements.txt

• Baichuan7B 模型下载

baichuan-inc/Baichuan-7B at main，主要模型文件大概 14G：

三.模型训练

下面对主要代码与功能进行解释，并在最后给出完整代码。

1.依赖引入与 tokenizer 加载

from transformers.integrations import TensorBoardCallback
from torch.utils.tensorboard import SummaryWriter
from transformers import TrainingArguments
from transformers import Trainer, HfArgumentParser
from transformers import AutoTokenizer, AutoModelForCausalLM
from transformers import DataCollatorForLanguageModeling
import torch
import torch.nn as nn
from peft import get_peft_model, LoraConfig, TaskType
from dataclasses import dataclass, field
import datasets
import os
from pprint import pprint as print

模型加载需要下载对应的 Baichuan7B 模型，通过 AutoTokenizer 类加载预训练模型得到 tokenizer 用于后续 encode 编码：

# 读取预训练模型
model_checkpoint = "/data3/model/baichuan-7B"
tokenizer = AutoTokenizer.from_pretrained(model_checkpoint, trust_remote_code=True)
tokenizer.pad_token = tokenizer.unk_token

这里还规定了 pad_token 为 unk_token，下面简单介绍下 tokenizer 几种常用的 token 形式：

pad-token	填充标记	[PAD]
eos_token	结束标记	[SEP]
unk_token	不在词汇表的单词	[UNK]

根据不同的模型和设定，上面的 token 形式可能出现不同，除此之外还有 pad_first 参数用于设置是否将填充放到文本最前面。

2.加载 DataSet 与 Model

def main():
    # 获取微调参数-输入、输入、Rank 训练参数-import by Transformers
    finetune_args, training_args = HfArgumentParser(
        (FinetuneArguments, TrainingArguments)
    ).parse_args_into_dataclasses()

    # 读取训练数据地址
    dataset = datasets.load_from_disk('data/tokenized_data/'+finetune_args.tokenized_dataset)
    print(f"\n{len(dataset)=}\n")

    # 读取并加载预训练模型，device_map 用于指定 GPU
    model = AutoModelForCausalLM.from_pretrained(
        model_checkpoint, load_in_8bit=False, trust_remote_code=True,
        device_map="auto" # 模型不同层会被自动分配到不同GPU上进行计算
    )
    print(model.hf_device_map)

• 微调训练参数

@dataclass
class FinetuneArguments:
    tokenized_dataset: str = field(default=" ") # tokenized之后的数据集文件夹
    model_path: str = field(default=" ") # 存储模型地址
    lora_rank: int = field(default=8) # lora 微调 emb rank

finetune_args 来自于自定义类，training_args 来自于 Transformer 类，前者主要包含了微调的基础数据路径和模型路径，后者包含一些训练相关参数。

• DataSet 加载

上文中我们 token 后的数据存放在 /data/toenized_data 目录下，如果有修改改为自己的样本目录即可。通过调用 dataset.load_from_disk 方法并传入文件夹路径，可以将之前保存的数据集重新加载到内存中，并返回一个可直接使用的数据集对象。加载后的数据集可以在后续的数据处理、模型训练或其他任务中使用。

• Model 加载

通过 transformers 提供的 AutoModelForCausalLM.from_pretrained API 加载预训练模型。

AutoModelForCausalLM 引自 Transformers ，用于自动化选择和加载适合特定任务的预训练语言模型 [LM] 的工具。该类是由 Hugging Face 提供的 "Auto" 系列类之一，旨在简化使用和处理不同类型的预训练模型，例如生成式语言建模 [Causal Language Modeling]。

通过使用 AutoModelForCausalLM 类，无需手动选择和加载特定的语言模型，而是仅提供预训练模型的标识符或名称，库将根据您提供的提示自动选择最佳模型，并加载相应的权重和配置。除了我们自己下载的模型外，一些早期的 LM 我们可以直接加载，以 GPT-2 为例：

from transformers import AutoModelForCausalLM, AutoTokenizer

model_name = "gpt2"  # 指定预训练模型的标识符或名称
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

# 使用加载的模型进行后续的生成式语言建模任务

• load_in_8bit

load_in_8bit 顾名思义该参数是指在加载预训练模型时，将权重参数转换为使用 8 位精度进行存储和计算。当设置 load_in_8bit=True 时，预训练模型的权重参数会以更低的精度 [8位] 进行存储，从而减少了模型所需的内存空间。这可以帮助降低模型加载时间、减少内存开销，并使模型能够在较低的设备或资源受限的环境中运行。

尽管 8 位精度相对于默认的 32 位精度会引入一定程度的数值损失，但在大多数情况下，这种损失是可以接受的，并且可以获得显著的性能提升。需要注意的是，load_in_8bit 参数的可用性可能取决于所选的预训练模型和库的版本，例如 Baichuan7B 就支持该参数。

• Lora Rank

lorarank 即在原始预训练模型的基础上增加旁路，旁路训练参数的低维维度。

关于 Lora 的基础知识我主要参考了 LORA：大模型轻量级微调一文。

3.Model 参数配置

• 模型启用梯度检查点功能

    model.gradient_checkpointing_enable()

在深度神经网络 [deep neural network] 训练时，计算梯度是一项极其耗费计算资源的操作。梯度检查点技术通过分解长序列的计算图为多个子图，并在每个子图上进行反向传播和梯度计算，从而减少了计算梯度所需的内存和计算量。

优劣：使用梯度检查点可以节省内存，但代价是向后传递速度较慢。

• 模型启用计算输入梯度的功能

model.enable_input_require_grads()

在深度神经网络 [deep neural network] 训练时，需要对每个参数或权重 [parameter/weight] 计算其对损失函数 [loss function] 的梯度 [gradient]，从而进行反向传播 [back propagation] 和优化[optimization]。默认情况下不会计算输入数据 [input data] 的梯度，即使它们在计算中起到了关键的作用。但是，在某些应用场景中，例如图像生成 [image generation]、注意力机制 [attention mechanism] 等，需要计算输入数据的梯度。此时，可以通过启用计算输入梯度的功能，对输入数据进行求导并利用其梯度信息进行优化。

作用: 启用该功能这对于在保持模型权重固定的同时微调适配器权重非常有用。

• 模型并行

model.model_parallel = False

设置为 True 之后，可能会启动模型并行化，且关闭数据并行，让一个模型分块在多块 GPU 上。一般设置为 False 主要基于以下几个原因：

- 硬件资源有限

- 任务规模较小

- 库或者框架不支持

这里我们设置为 False 主要因为穷。

• 获取 Linear 层

model.lm_head = CastOutputToFloat(model.lm_head)

lm_head 可以获取其线性变换层或者头部网络的参数。对于一个语言模型来说，其任务通常是通过输入一个文本序列，输出该序列的下一个单词或者完成一个文本生成任务。在实现上，通常会采用一个多层的神经网络结构，其中最后一层就是线性变换层或者头部网络，用来将上一层的特征表示映射为最终的输出结果。这些层的参数包括权重矩阵和偏置向量，可以通过 model.lm_head 属性来获取。

class CastOutputToFloat(nn.Sequential):
    def forward(self, x):
        return super().forward(x).to(torch.float32)

CastOutputToFloat 类为自定义类，用于定义 lm_head 前向传播时的精度。

4.获取 peft Model

peft_config = LoraConfig(
        task_type=TaskType.CAUSAL_LM, # 因果语言模型可以被看作是自回归语言模型的一种扩展，在每个时间步上，它不仅考虑之前的文本序列，还>考虑了未来可能出现的单词信息，从而更好地捕捉语句的因果结构。
        inference_mode=False,
        r=finetune_args.lora_rank,
        lora_alpha=32,
        lora_dropout=0.1,
        target_modules = ["W_pack"] # 把model打印出来，找跟attention相关的模块
    )

# 获取 Lora 微调模型
model = get_peft_model(model, peft_config)

通过上面一系列操作我们已经成功加载好预训练模型与对应 Tokenizer，接下来通过 peft 库实现 Lora 模型的获取。PEFT 是一个面向高性能计算 [HPC] 系统的功耗估计和建模框架，可以在不需要硬件测量的情况下，通过基于任务图的分析方法，对 HPC 应用程序的功耗进行精确建模和估计。

具体来说，PEFT 可以生成应用程序的任务图，并根据该图推断出任务之间的依赖关系和通信开销等信息，从而对应用程序的实际执行过程进行模拟和预测，最终得到应用程序的功耗估计结果。

• LoraConfig

TaskType.CAUSAL_LM 表示构建一个因果语言模型 [causal language model] 的任务类型，可以用于指定不同的模型参数和超参数等设置。

lora 微调需要设置两个参数一个是 r，即中间层神经元的个数，另一个 lora_alpha 是一个 scale 参数，lora_dropout 即为 dropout 比例。

• 模型加载与读取

通过 get_peft_model API 与定义好的 LoraConfig 即可获取对应的 peft model。

model.print_trainable_parameters()

该方法可以打印训练参数信息。模型微调完毕后，额外加载 lora 参数即可实现模型加载：

model.load_state_dict(torch.load(peft_path), strict=False)

5.构造 Trainer 训练

    # 初始化 trainer, 此处报错: NotImplementedError: Cannot copy out of meta tensor; no data!
    trainer = ModifiedTrainer(
        model=model,
        train_dataset=dataset,
        args=training_args,
        callbacks=[TensorBoardCallback(writer)], # 回调将训练情况写入 tensorboard
        data_collator=data_collator,
    )
    trainer.train()

• ModifiedTrainer

class ModifiedTrainer(Trainer):
    # 根据输入计算 Loss
    def compute_loss(self, model, inputs, return_outputs=False):
        return model(
            input_ids=inputs["input_ids"],
            labels=inputs["labels"],
        ).loss

    # 保存模型
    def save_model(self, output_dir=None, _internal_call=False):
        from transformers.trainer import TRAINING_ARGS_NAME

        os.makedirs(output_dir, exist_ok=True)
        torch.save(self.args, os.path.join(output_dir, TRAINING_ARGS_NAME))
        saved_params = {
            k: v.to("cpu") for k, v in self.model.named_parameters() if v.requires_grad
        }
        torch.save(saved_params, os.path.join(output_dir, "adapter_model.bin"))

Trainer 是 Hugging Face Transformers 库中的一个类，它提供了一个简便且强大的训练循环来训练和评估模型。其主要用于监督学习任务，如文本分类、命名实体识别、文本生成等。它帮助自动处理训练过程中的批量数据加载、优化器设置、模型更新等常见任务，以及记录和输出训练指标和日志。本例中 compute_loss 方法指定计算 loss 的 input 与 label、 save_model 方法首先获取需要更新的参数，随后将参与更新的参数存储至 output 生成 .bin 文件，训练结束后读取 .bin 即可获取我们 lora 后的新模型。

下面是 Trainer 的常用 demo：

from transformers import Trainer, TrainingArguments

# 定义训练参数和配置
training_args = TrainingArguments(
    output_dir='./results',          # 指定训练结果保存的目录
    num_train_epochs=3,              # 训练的轮次（epochs）
    per_device_train_batch_size=16,  # 每个设备的训练批次大小
    per_device_eval_batch_size=8,    # 每个设备的评估批次大小
    warmup_steps=500,                # 学习率预热步数
    weight_decay=0.01,               # 权重衰减参数
    logging_dir='./logs',            # 日志保存的目录
)

# 创建 Trainer 对象并进行训练
trainer = Trainer(
    model=model,                     # 要训练的模型
    args=training_args,              # 训练参数和配置
    train_dataset=train_dataset,     # 训练数据集
    eval_dataset=eval_dataset        # 评估数据集
)

trainer.train()  # 开始训练

# 在评估集上进行评估
trainer.evaluate()

# 保存训练好的模型
trainer.save_model("./saved_model")

• callback 回调

writer = SummaryWriter()

在 main 函数的最前面我们初始化 SummaryWriter 类，该类为工具类，用于将条目直接写入 log_dir 中的时间文件以供 Tensorboard 使用，后面通过 callbacks 传入 TensorBoardCallback 用于指标记录供后续展示。任务结束后记得调用 close 方法关闭 writer。

• DataCollatorForLanguageModeling

data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer,mlm=False)

该类也是 Hugging Face Transformers 库中的一个类，用于帮助组织和处理语言建模任务的训练数据。在自然语言处理和语言建模中，语言模型通常是通过对大量文本进行预训练，并使用生成式任务来微调模型。DataCollatorForLanguageModeling 类提供了一个方便的方式来处理和准备用于语言建模任务的训练数据。该类的主要功能包括：

- 数据批次化 Batching: 将输入文本数据按照指定的批次进行大小划分以便批量训练

- 数据填充 Padding: 输入文本长度不一致时可自动填充，以便样本一致从而高效并行计算

- 掩码处理 Masking: 对于基于 mask 的语言任务例如 Bert，可以帮助生成掩码标签

下面展示了如何使用 DataCollect 类准备语言建模任务的训练数据：

from transformers import DataCollatorForLanguageModeling

data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=True)

# 准备输入文本数据
input_text = ["I love natural language processing.", "Transformers library is great!"]

# 根据指定的批次大小进行数据批次化和填充
batch = data_collator(input_text)

# 使用批次进行语言建模任务的训练
model.train_on_batch(batch)

6.训练完整代码

from transformers.integrations import TensorBoardCallback
from torch.utils.tensorboard import SummaryWriter
from transformers import TrainingArguments
from transformers import Trainer, HfArgumentParser
from transformers import AutoTokenizer, AutoModelForCausalLM
from transformers import DataCollatorForLanguageModeling
import torch
import torch.nn as nn
from peft import get_peft_model, LoraConfig, TaskType
from dataclasses import dataclass, field
import datasets
import os
from pprint import pprint as print

import contextlib

autocast = contextlib.nullcontext

# 读取预训练模型
model_checkpoint = "/data/model/baichuan-7b/"
tokenizer = AutoTokenizer.from_pretrained(model_checkpoint, trust_remote_code=True)


# 微调参数
@dataclass
class FinetuneArguments:
    tokenized_dataset: str = field(default=" ")  # tokenized之后的数据集文件夹
    model_path: str = field(default=" ")  # 存储模型地址
    lora_rank: int = field(default=8)  # lora 微调 emb rank


# 输出转换 IntToFloat
class CastOutputToFloat(nn.Sequential):
    def forward(self, x):
        return super().forward(x).to(torch.float32)


# 填充标记 pad_token-填充标记-[PAD] eos_token-结束标记-[SEP] unk_token-用于填充不在词汇表中的关键字-[unk] pad_first–是否将填充放到文本最前面 
tokenizer.pad_token = tokenizer.unk_token
# 数据函数隶属于 Trainer，用 list of elements 来构造一个 batch
data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)


class ModifiedTrainer(Trainer):
    # 根据输入计算 Loss
    def compute_loss(self, model, inputs, return_outputs=False):
        return model(
            input_ids=inputs["input_ids"],
            labels=inputs["labels"],
        ).loss

    # 保存模型
    def save_model(self, output_dir=None, _internal_call=False):
        from transformers.trainer import TRAINING_ARGS_NAME

        os.makedirs(output_dir, exist_ok=True)
        torch.save(self.args, os.path.join(output_dir, TRAINING_ARGS_NAME))
        saved_params = {
            k: v.to("cpu") for k, v in self.model.named_parameters() if v.requires_grad
        }
        torch.save(saved_params, os.path.join(output_dir, "adapter_model.bin"))


def main():
    # 工具类: 用于将条目直接写入 log_dir 中的时间文件以供 Tensorboard 使用
    writer = SummaryWriter()
    # 获取微调参数-输入、输入、Rank 训练参数-import by Transformers
    finetune_args, training_args = HfArgumentParser(
        (FinetuneArguments, TrainingArguments)
    ).parse_args_into_dataclasses()

    # 读取训练数据地址
    dataset = datasets.load_from_disk('data/tokenized_data/' + finetune_args.tokenized_dataset)
    # dataset = dataset.select(range(10000)) # 获取数据集前10000个项目
    print(f"\n{len(dataset)=}\n")

    # 读取并加载预训练模型，device_map 用于指定 GPU
    model = AutoModelForCausalLM.from_pretrained(
        model_checkpoint, load_in_8bit=False, trust_remote_code=True,
        device_map="auto"  # 模型不同层会被自动分配到不同GPU上进行计算
    )
    print(model.hf_device_map)

    # 模型配置
    model.gradient_checkpointing_enable()
    model.enable_input_require_grads()
    model.lm_head = CastOutputToFloat(model.lm_head)

    # TaskType.CAUSAL_LM 表示构建一个因果语言模型（causal language model）的任务类型，可以用于指定不同的模型参数和超参数等设置。
    peft_config = LoraConfig(
        task_type=TaskType.CAUSAL_LM,  # 因果语言模型可以被看作是自回归语言模型的一种扩展，在每个时间步上，它不仅考虑之前的文本序列，还考虑了未来可能出现的单词信息，从而更好地捕捉语句的因果结构。
        inference_mode=False,
        r=finetune_args.lora_rank,
        lora_alpha=32,
        lora_dropout=0.1,
        target_modules=["W_pack"]  # 把model打印出来，找跟attention相关的模块
    )

    # 获取 Lora 微调模型
    model = get_peft_model(model, peft_config)

    # 训练步骤
    model.save_pretrained(
        training_args.output_dir)  # 因为 adapter_config.json 只能通过这个 save_pretrained 来生成，先这里生成一份，好在训练完之前就可以尝试中间的checkpoint

    # 初始化 trainer, 此处报错: NotImplementedError: Cannot copy out of meta tensor; no data! 
    trainer = ModifiedTrainer(
        model=model,
        train_dataset=dataset,
        args=training_args,
        callbacks=[TensorBoardCallback(writer)],  # 回调将训练情况写入 tensorboard
        data_collator=data_collator,
    )
    trainer.train()
    writer.close()
    # save model
    model.save_pretrained(training_args.output_dir)


if __name__ == "__main__":
    main()

四.Shell 执行

1.脚本构建

构造下述 train.sh 的脚本，其中 train 的文件名为我们上文 tokenizer 生成的 DataSet 样本，这里 lora_rank 设定为 4，batch_size 设定为 8，共训练 10 个 epoch，由于训练数据很少，所以只是跑通流程，大家也可以构建更多 QA 样本，采用同样流程训练。

train="simple_token_by_baichuan-7B"
output="simple_token_by_baichuan-7B"

CUDA_VISIBLE_DEVICES=0,1 python baichuan_lora_tuning.py \
     --tokenized_dataset $train \
     --lora_rank 4 \
     --per_device_train_batch_size 8 \
     --gradient_accumulation_steps 1 \
     --num_train_epochs 10 \
     --save_steps 200 \
     --save_total_limit 2 \
     --learning_rate 1e-4 \
     --fp16 \
     --remove_unused_columns false \
     --logging_steps 50 \
     --output_dir weights/$output