微调 LLM——以 Baichuan-7B 为例

发布于 2023-07-20  140 次阅读


参考视频:https://www.bilibili.com/video/BV1yu411L7JN/

仓库:https://github.com/EvilPsyCHo/train_custom_LLM

输入构建

单轮对话

  1. 直接拼接数据集中的输入和输出。缺点:不便于模型判断什么时候终止对话。
  2. 在尾部添加 eos_token。缺点:模型不好判断哪个是输入哪个是输出。
  3. 标识出输入输出,比如添加 Instruction:Output:

多轮对话

在输入中加入轮次信息。例如,对于每一轮历史对话,加上 [Round x]\n\n,当前轮次也要加上,例如:

def build_inputs(self, tokenizer, query: str, history: List[Tuple[str, str]] = None):
    prompt = ""
    for i, (old_query, response) in enumerate(history):
        prompt += "[Round {}]\n\n问:{}\n\n答:{}\n\n".format(i + 1, old_query, response)
    prompt += "[Round {}]\n\n问:{}\n\n答:".format(len(history) + 1, query)
    inputs = tokenizer([prompt], return_tensors="pt")
    inputs = inputs.to(self.device)
    return inputs

Label Mask

我们希望模型去学习如何回答,而不用预测我们额外添加的信息,比如输入输出标识和轮次信息。

因此我们可以使用 Label Mask 去屏蔽这部分的 loss。

模型

模型加载(量化或非量化)

  • baichuan-7B显存开销 float16 (13GB), 8bit (7.1GB), 4bit (4.1GB)
  • 模型定义变化 torch.nn.Linear => bitsandbytes.nn.Linear8bitLt => bitsandbytes.nn.Linear4bit

模型准备

  • 冻结参数,Lora 训练不改变原模型参数,因此需要冻结。
  • 调整输出层精度,对 loss 的计算有好处。
  • 如果使用了 gradient_checkpoint,梯度将被写到显存以外的地方来节省显存,如此 lora 拿不到这部分梯度,因此需要额外设置 enable_input_require_grads

以上准备可以很方便地由 peft.prepare_model_for_kbit_training 实现。

Lora 配置

  • 找到需要使用 lora 训练的 layer module names,有两个方法:

    1. 遍历模型的所有 module 和对应的 name,如果是目标 class,则记录下来并将其返回,注意,如果 lm_head (输出层)也在其中,需要将其剔除。示例:
      
      def find_all_linear_names(model):
      #cls = bnb.nn.Linear8bitLt
      cls = bnb.nn.Linear4bit
      lora_module_names = set()
      for name, module in model.named_modules():
      if isinstance(module, cls):
          names = name.split('.')
          lora_module_names.add(names[0] if len(names) == 1 else names[-1])

    if 'lm_head' in lora_module_names: # needed for 16-bit
    lora_module_names.remove('lm_head')
    return list(lora_module_names)

    
    2. 直接根据模型信息,将目标层作为参数传入(LoraConfig.target_modules)。
  • lora 配置
  • get_peft_model 得到 lora 挂载模型。

Label shift and Loss

  • 一般使用 CrossEntropyLoss。使用 ignore_index=tokenizer.pad_token_id 来屏蔽填充项的损失。
  • 在生成时,输入和标签是一致的,而模型是根据上一个(实际上是之前的序列) token 生成下一个 token,因此生成出来的序列相比 lable 有偏移,所以要修正。例如,输入和标签均为 [1, 2, 3, 4] ,那么生成将是 [2, 3, 4, 5] ,所以我们去掉标签的第一位,以及生成的最后一位,这样两者都变成了 [2, 3, 4],可以计算 Loss。

训练

  • transfomers 内置的 Trainer 类已经够用,实践中基本只需重写其 compute_loss 方法。
  • TrainingArgumentsTrainer 配套使用,其定义了绝大部分训练所需要的参数。