跳到正文
原文
Hugging Face Blog·· 2023-08-08精选AI 评分61

用 TRL 的 DPO 方法微调 Llama 2

Fine-tune Llama 2 with DPO

AI 导读

Hugging Face 的 TRL 库现已支持 Direct Preference Optimization（DPO），可直接用偏好数据优化语言模型，省去奖励模型和 RL 优化两步。

推荐理由

TRL 官方给出 DPO 原理与可复用代码，读者可照着在自有偏好数据上完成对齐训练。

正文 · AI 翻译

简介

人类反馈强化学习（RLHF）已经成为 GPT-4 或 Claude 等 LLMs 事实上的最后训练步骤，以确保语言模型的输出符合人类期望，如聊天性或安全特性。然而，它将 RL 的一些复杂性带入 NLP：我们需要构建一个好的奖励函数，训练模型来估计状态的价值，同时要注意不要偏离原始模型太远，否则会产生无意义的文本而不是合理的文本。这样一个过程相当复杂，需要许多复杂的动态部分，而且并不总是容易做到正确。

Rafailov、Sharma、Mitchell 等人最近的论文 Direct Preference Optimization 提出将现有方法使用的基于 RL 的目标转化为可以通过简单的二元交叉熵损失直接优化的目标，从而大大简化了精炼 LLMs 的过程。

这篇博文介绍了 Direct Preference Optimization（DPO）方法，该方法现已在 TRL 库 中提供，并展示了如何在 stack-exchange preference 数据集上微调最近的 Llama v2 7B 参数模型，该数据集包含来自各个 stack-exchange 门户的问题的排序答案。

DPO 与 PPO

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

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

因此，这个直接似然目标可以在不需要奖励模型或不需要执行可能繁琐的基于 RL 的优化的情况下进行优化。

如何使用 TRL 进行训练

如前所述，通常 RLHF 流水线由这些不同的部分组成：

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

TRL 库为所有这些部分都提供了辅助工具，然而 DPO 训练省去了奖励建模和 RL 的任务（第 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，换句话说，即被选中的奖励高于被拒绝的奖励（或差距大于零）。随后可以在某个评估数据集上计算这些指标。

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

来源：Hugging Face Blog · hf-mirror.com

#教程/实践#开源/仓库#数据/训练#Hugging Face#Meta