跳到正文
原文
Hugging Face：Blog（RSS）·· 2023-08-08精选AI 评分75

Hugging Face 教程：用 DPO 微调 Llama 2

Fine-tune Llama 2 with DPO

AI 导读

Hugging Face 发布教程，介绍 TRL 库新增的 DPOTrainer，用 Direct Preference Optimization 方法跳过奖励建模和 RL 步骤，直接在偏好数据上优化语言模型。

推荐理由

原文用 Llama 2 和 stack-exchange 数据演示 DPO 全流程，读者可以照着复现替代传统 RLHF 的对齐训练。

正文 · AI 翻译

引言

基于人类反馈的强化学习（RLHF）已成为 GPT-4 或 Claude 等大型语言模型事实上的最后训练步骤，以确保语言模型的输出符合人类期望，如聊天性或安全特性。然而，它将强化学习的一些复杂性带入了自然语言处理领域：我们需要构建一个良好的奖励函数，训练模型来估计状态的价值，同时要小心不要偏离原始模型太远，以免产生无意义的文本而非合理的文字。这一过程相当复杂，涉及许多复杂的动态部分，要正确完成并不总是那么容易。

Rafailov、Sharma、Mitchell 等人最近发表的论文直接偏好优化提出将现有方法使用的基于强化学习的目标转换为可以通过简单的二元交叉熵损失直接优化的目标，这大大简化了大型语言模型的精炼过程。

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

DPO 与 PPO

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

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

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

如何使用 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 模型上通过 TRL 的 SFTTrainer 对数据的 SFT 划分进行监督微调步骤：

# 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 位配置加载模型，然后通过 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：所选奖励与对应被拒绝奖励之间的平均差值。

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

我们希望随着代码的发布，能降低读者自行在数据集上尝试这种大语言模型对齐方法的门槛，我们迫不及待想看到你们构建的作品！如果你想亲自试用该模型，可以在这里进行：trl-lib/stack-llama。

来源：Hugging Face：Blog（RSS） · huggingface.co

#教程/实践#开源生态#数据/训练#Hugging Face#Meta