使用 DPO 微调 Llama 2

引言

基于人类反馈的强化学习 (RLHF) 已成为 GPT-4 或 Claude 等大语言模型 (LLM) 训练的最后一步，以确保语言模型的输出符合人类期望，例如健谈或安全特性。然而，它将强化学习 (RL) 的一些复杂性带入了自然语言处理 (NLP)：我们需要构建一个好的奖励函数，训练模型来评估一个状态的价值，同时要小心不要偏离原始模型太远，以至于产生胡言乱语而不是有意义的文本。这样的过程相当复杂，需要许多复杂的活动部件，要做到恰到好处并不总是那么容易。

最近，Rafailov、Sharma、Mitchell 等人发表的论文 直接偏好优化 (Direct Preference Optimization) 提出，将现有方法使用的基于强化学习的目标函数，转换为一个可以通过简单的二元交叉熵损失直接优化的目标函数，从而极大地简化了改进大语言模型的过程。

这篇博文介绍了直接偏好优化 (DPO) 方法，该方法现已在 TRL 库中提供，并展示了如何使用 stack-exchange 偏好数据集微调最近发布的 Llama v2 7B 参数模型。该数据集包含对 Stack Exchange 各种门户网站上问题的排名答案。

DPO 与 PPO

在通过强化学习优化人类衍生偏好的传统模型中，首选方法是使用一个辅助的奖励模型，并通过强化学习的机制微调目标模型，使其最大化这个给定的奖励。直观地说，我们使用奖励模型向正在优化的模型提供反馈，使其更频繁地生成高奖励样本，减少低奖励样本。同时，我们使用一个冻结的参考模型，以确保生成的内容不会偏离太远并继续保持生成的多样性。这通常是通过一个参考模型向完整的奖励最大化目标添加一个 KL 惩罚来实现的，这有助于防止模型学习欺骗或利用奖励模型。

DPO 公式绕过了奖励建模步骤，并直接在偏好数据上优化语言模型，其关键洞见在于：从奖励函数到最优 RL 策略的分析性映射，这使得作者能够将关于奖励和参考模型的 RL 损失转换为直接关于参考模型的损失！这种映射直观地衡量了给定的奖励函数与给定的偏好数据对齐的程度。因此，DPO 从 RLHF 损失的最优解出发，通过变量替换，推导出一个*仅*关于参考模型的损失函数！

因此，这个直接的似然目标函数可以被优化，而无需奖励模型或执行可能繁琐的基于 RL 的优化。

如何使用 TRL 进行训练

如前所述，RLHF 流程通常包括以下几个不同部分：

1. 监督式微调 (SFT) 步骤
2. 使用偏好标签标注数据的过程
3. 在偏好数据上训练奖励模型
4. 以及强化学习优化步骤

TRL 库为所有这些部分提供了辅助工具，然而，DPO 训练省去了奖励建模和强化学习的任务 (步骤 3 和 4)，并直接在带有偏好标注的数据上优化 DPO 对象。

在这方面，我们仍然需要执行步骤 1，但取代步骤 3 和 4 的是，我们需要向 TRL 中的 DPOTrainer 提供来自步骤 2 的偏好数据，该数据具有非常特定的格式，即一个包含以下三个键的字典：

• prompt: 这包括在推理时提供给模型用于文本生成的上下文提示
• chosen: 包含对相应提示的首选生成响应
• rejected: 包含相对于给定提示不被偏好或不应被采样的响应

例如，对于 stack-exchange 偏好对数据集，我们可以通过以下辅助函数映射数据集条目以返回所需的字典，并删除所有原始列。

def return_prompt_and_responses(samples) -> Dict[str, str, str]:
    return {
        "prompt": [
            "Question: " + question + "\n\nAnswer: "
            for question in samples["question"]
        ],
        "chosen": samples["response_j"],   # rated better than k
        "rejected": samples["response_k"], # rated worse than j
    }

dataset = load_dataset(
    "lvwerra/stack-exchange-paired",
    split="train",
    data_dir="data/rl"
)
original_columns = dataset.column_names

dataset.map(
    return_prompt_and_responses,
    batched=True,
    remove_columns=original_columns
)

一旦我们将数据集整理好，DPO 损失本质上是一个监督式损失，它通过一个参考模型获得隐式奖励。因此，在较高层面上，DPOTrainer 需要我们希望优化的基础模型以及一个参考模型。

dpo_trainer = DPOTrainer(
    model,                 # base model from SFT pipeline
    model_ref,             # typically a copy of the SFT trained base model
    beta=0.1,              # temperature hyperparameter of DPO
    train_dataset=dataset, # dataset prepared above
    tokenizer=tokenizer,   # tokenizer
    args=training_args,    # training arguments e.g. batch size, lr, etc.
)

其中 beta 超参数是 DPO 损失的温度参数，通常在 0.1 到 0.5 的范围内。这控制了我们对参考模型的关注程度，即当 beta 变小时，我们越是忽略参考模型。一旦我们的训练器初始化完成，我们就可以通过简单地调用以下命令，使用给定的 training_args 在数据集上对其进行训练。

dpo_trainer.train()

使用 Llama v2 进行实验

在 TRL 中实现 DPO 训练器的好处是，可以利用 TRL 及其依赖库 (如 Peft 和 Accelerate) 提供的所有额外功能来训练大型 LLM。借助这些库，我们甚至能够使用由 bitsandbytes 库提供的 QLoRA 技术来训练 Llama v2 模型。

监督式微调

如上所述，该过程涉及使用 QLoRA 在 7B Llama v2 模型上进行监督式微调，该微调在数据的 SFT 分割上通过 TRL 的 SFTTrainer 进行。

# load the base model in 4-bit quantization
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16,
)

base_model = AutoModelForCausalLM.from_pretrained(
    script_args.model_name,        # "meta-llama/Llama-2-7b-hf"
    quantization_config=bnb_config,
    device_map={"": 0},
    trust_remote_code=True,
    use_auth_token=True,
)
base_model.config.use_cache = False

# add LoRA layers on top of the quantized base model
peft_config = LoraConfig(
    r=script_args.lora_r,
    lora_alpha=script_args.lora_alpha,
    lora_dropout=script_args.lora_dropout,
    target_modules=["q_proj", "v_proj"],
    bias="none",
    task_type="CAUSAL_LM",
)
...
trainer = SFTTrainer(
    model=base_model,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    peft_config=peft_config,
    packing=True,
    max_seq_length=None,
    tokenizer=tokenizer,
    args=training_args,         # HF Trainer arguments
)
trainer.train()

DPO 训练

SFT 完成后，我们可以保存生成的模型，然后进行 DPO 训练。通常的做法是，我们将上一步 SFT 保存的模型同时用作 DPO 的基础模型和参考模型。然后，我们可以使用这些模型，在上面展示的 stack-exchange 偏好数据上，通过 DPO 目标来训练模型。由于模型是通过 LoRa 适配器训练的，我们使用 Peft 的 AutoPeftModelForCausalLM 辅助工具来加载模型。

model = AutoPeftModelForCausalLM.from_pretrained(
    script_args.model_name_or_path, # location of saved SFT model
    low_cpu_mem_usage=True,
    torch_dtype=torch.float16,
    load_in_4bit=True,
    is_trainable=True,
)
model_ref = AutoPeftModelForCausalLM.from_pretrained(
    script_args.model_name_or_path,  # same model as the main one
    low_cpu_mem_usage=True,
    torch_dtype=torch.float16,
    load_in_4bit=True,
)
...
dpo_trainer = DPOTrainer(
    model,
    model_ref,
    args=training_args,
    beta=script_args.beta,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    tokenizer=tokenizer,
    peft_config=peft_config,
)
dpo_trainer.train()
dpo_trainer.save_model()

因此，如所见，我们在 4-bit 配置中加载模型，然后通过 peft_config 参数使用 QLora 方法对其进行训练。训练器还将在训练期间评估相对于评估数据集的进度，并报告一些关键指标，如隐式奖励，这些指标可以通过 WandB 等工具记录和显示。然后，我们可以将最终训练好的模型推送到 HuggingFace Hub。

结论

SFT 和 DPO 训练脚本的完整源代码可在以下 examples/stack_llama_2 目录中找到，合并了适配器的训练模型可在 HF Hub 这里 找到。

DPO 训练运行的 WandB 日志可以在这里找到，在训练和评估期间，DPOTrainer 会记录以下奖励指标：

• rewards/chosen：策略模型和参考模型对所选响应的对数概率之差的平均值，按 beta 缩放
• rewards/rejected：策略模型和参考模型对被拒响应的对数概率之差的平均值，按 beta 缩放
• rewards/accuracies：所选奖励大于相应被拒奖励的频率的平均值
• rewards/margins：所选奖励与相应被拒奖励之差的平均值。

直观上，在训练过程中，我们希望边际值 (margins) 增加，准确率 (accuracies) 趋向于 1.0，换句话说，就是所选奖励高于被拒奖励 (或者说边际值大于零)。这些指标随后可以在某个评估数据集上计算。

我们希望通过代码的发布，能为各位读者降低门槛，在自己的数据集上尝试这种对齐大型语言模型的方法，我们迫不及待地想看到你们的创造！如果你想亲自尝试这个模型，可以在这里进行：trl-lib/stack-llama。