用DPO对Llama 2进行微调

DPO微调Llama 2

介绍

人类反馈强化学习（RLHF）已经成为GPT-4或Claude等LLM的最后训练步骤，以确保语言模型的输出与人类的期望（如健谈或安全功能）保持一致。然而，这将一些强化学习的复杂性引入到自然语言处理中：我们需要构建一个良好的奖励函数，训练模型来估计状态的价值，并且同时要小心不要过分远离原始模型，以免产生胡言乱语而不是有意义的文本。这个过程非常复杂，需要许多复杂的组件，其中并不总是容易做到事事正确。

最近的一篇论文《直接偏好优化》（Direct Preference Optimization）由Rafailov，Sharma，Mitchell等人提出，将现有方法使用的基于强化学习的目标转化为可以直接通过简单的二元交叉熵损失进行优化的目标，从而极大地简化了LLM的改进过程。

本博客文章介绍了直接偏好优化（DPO）方法，该方法现在已经在TRL库中提供，并展示了如何在Stack Exchange偏好数据集上对最新的Llama v2 7B参数模型进行微调，该数据集包含了各种Stack Exchange门户网站上的问题的排名答案。

DPO与PPO

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

• “NVIDIA H100 Tensor Core GPU 现已正式发布，用于新的 Microsoft Azure 虚拟机系列”
• “在NVIDIA Studio中进行内容创作，得益于新的专业GPU、AI工具、Omniverse和OpenUSD合作功能的增强”
• 扩展版：NVIDIA扩展Maxine以用于视频编辑，展示3D虚拟会议研究

DPO公式绕过了奖励建模步骤，并直接通过关键洞察力在偏好数据上优化语言模型，即从奖励函数到最优强化学习策略的解析映射，使作者能够将奖励模型和参考模型上的RL损失转化为仅基于参考模型的损失！这种映射直观地衡量了给定奖励函数与给定偏好数据的一致性。因此，DPO从RLHF损失的最优解开始，通过变量的改变得出了仅基于参考模型的损失！

因此，这个直接似然目标可以在不需要奖励模型或进行可能棘手的基于强化学习的优化的情况下进行优化。

如何使用TRL进行训练

如前所述，通常RLHF流程由以下不同的部分组成：

1. 有监督微调（SFT）步骤
2. 使用偏好标签对数据进行注释的过程
3. 在偏好数据上训练奖励模型
4. 和RL优化步骤

TRL库提供了所有这些部分的辅助函数，然而DPO训练省去了奖励建模和RL（步骤3和4）的任务，直接在注释的偏好数据上优化DPO对象。

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

• prompt，这是在推理时给模型的上下文提示，用于生成文本
• chosen，包含对应提示的首选生成响应
• rejected，包含不首选或不应作为给定提示的抽样响应的响应

例如，对于Stack Exchange偏好配对数据集，我们可以通过以下辅助函数将数据集条目映射为返回所需字典，并删除所有原始列：

def return_prompt_and_responses(samples) -> Dict[str, str, str]:
    return {
        "prompt": [
            "问题：" + question + "\n\n答案："
            for question in samples["question"]
        ],
        "chosen": samples["response_j"],   # 优于k的评级
        "rejected": samples["response_k"], # 差于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,                 # SFT流程中的基础模型
    model_ref,             # 通常是SFT训练的基础模型的副本
    beta=0.1,              # DPO的温度超参数
    train_dataset=dataset, # 上面准备的数据集
    tokenizer=tokenizer,   # 分词器
    args=training_args,    # 训练参数，例如批大小、学习率等
)

其中，beta超参数是DPO损失的温度参数，通常在0.1到0.5的范围内。它控制我们在多大程度上关注参考模型，即当beta变得越小时，我们越忽略参考模型。一旦初始化了训练器，我们可以通过简单调用以下代码来对数据集进行训练：

dpo_trainer.train()

使用Llama v2进行实验

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

有监督微调

如上所述，该过程涉及使用QLoRA对SFT数据中的7B Llama v2模型进行有监督微调步骤，通过TRL的SFTTrainer实现：

# 在4位量化中加载基础模型
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

# 在量化的基础模型上添加LoRA层
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参数
)
trainer.train()

DPO训练

一旦SFT完成，我们可以保存结果模型并进行DPO训练。通常，我们将使用上一步SFT中保存的模型作为DPO的基础模型和参考模型。然后，我们可以使用这些模型在上面显示的stack-exchange偏好数据上使用DPO目标来训练模型。由于这些模型是通过LoRa适配器训练的，因此我们通过Peft的AutoPeftModelForCausalLM辅助函数加载模型：

model = AutoPeftModelForCausalLM.from_pretrained(
    script_args.model_name_or_path, # 保存的SFT模型的位置
    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,  # 与主模型相同的模型
    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位配置加载模型，然后通过QLora方法进行训练，通过peft_config参数。训练器还将根据评估数据集评估训练进展，并报告一些关键指标，如隐式奖励，可以通过WandB等方式记录并显示。然后，我们可以将最终训练好的模型推送到HuggingFace Hub。

结论

用于SFT和DPO的训练脚本的完整源代码可以在examples/stack_llama_2目录中找到，并且带有合并适配器的训练模型可以在HF Hub上找到。

可以在此处找到DPO训练运行的WandB日志，其中在训练和评估过程中，DPOTrainer记录以下奖励指标：

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

直观地说，在训练过程中，我们希望边界增加，准确率达到1.0，或者换句话说，所选奖励高于拒绝奖励（或边界大于零）。然后，可以在某个评估数据集上计算这些指标。

我们希望通过发布代码，降低读者尝试在自己的数据集上对齐大型语言模型的门槛，并且我们迫不及待地想要看到您的构建！如果您想自己尝试该模型，可以在此处进行：trl-lib/stack-llama。