pubanswer

使用 ORPO 微调 Llama 3

巅峰小怪兽 • 2024-06-26

ORPO 是一种新兴的微调技术，它将传统的监督微调和偏好对齐阶段结合成一个过程。这减少了训练所需的计算资源和时间。此外，实证结果表明，ORPO 在各种模型尺寸和基准上优于其他对齐方法。

本文中，我们将使用 TRL 库对新的 Llama 3 8B 模型进行 ORPO 微调。代码可以在 Google Colab 和 GitHub 的 LLM 课程中找到。

⚖️ ORPO

指令微调和偏好对齐是将大型语言模型（LLMs）调整到特定任务的关键技术。传统上，这涉及一个多阶段的过程：首先是在指令上进行监督微调（SFT），以使模型适应目标领域，然后是使用人类反馈的强化学习（RLHF）或直接偏好优化（DPO）等偏好对齐方法，以增加生成首选响应而非被拒绝响应的可能性。

然而，研究人员发现这种方法有一个局限性。尽管 SFT 有效地使模型适应了所需领域，但它无意中也增加了生成不期望答案的概率。因此，需要偏好对齐阶段来扩大首选输出和被拒绝输出之间的差距。

ORPO 通过将指令微调和偏好对齐结合成一个单一的训练过程，提供了一个优雅的解决方案。ORPO 修改了标准的语言建模目标，将负对数似然损失与一个赔率（OR）项结合起来。这个 OR 损失在弱惩罚被拒绝响应的同时，强烈奖励首选响应，从而使模型能够同时学习目标任务并与人类偏好对齐。

\mathscr{L}_{ORPO} = \mathbb{E}_{(x, y_{w}, y_l)}[\mathscr{L}_{SFT} + \lambda \cdot \mathscr{L}_{OR}]

ORPO 已经在主要的微调库中实现，如 TRL、Axolotl 和 LLaMA-Factory。下一节我们将看到如何在 TRL 中使用它。

💻 使用 ORPO 微调 Llama 3

Llama 3 是由 Meta 开发的最新 LLM 家族。该模型在一个包含 15 万亿个标记的数据集上进行了训练（相比之下，Llama 2 的标记为 2T）。已经发布了两个模型尺寸：一个是 700 亿参数模型，另一个是较小的 80 亿参数模型。700B 模型已经展示了令人印象深刻的性能，在 MMLU 基准上得分为 82，在 HumanEval 基准上得分为 81.7。

Llama 3 模型还将上下文长度增加到 8,192 个标记（Llama 2 为 4,096 个标记），并可能使用 RoPE 扩展到 32k。此外，模型使用了一个新的分词器，具有 128K 标记词汇表，将编码文本所需的标记数量减少了 15%。这种词汇表也解释了从 7B 到 8B 参数的提升。

ORPO 需要一个偏好数据集，包括一个提示、一个选择的答案和一个被拒绝的答案。在这个例子中，我们将使用 mlabonne/orpo-dpo-mix-40k，这是以下高质量 DPO 数据集的组合：

  • argilla/distilabel-capybara-dpo-7k-binarized: 高评分选择答案 >=5（2,882 样本）
  • argilla/distilabel-intel-orca-dpo-pairs: 高评分选择答案 >=9，不在 GSM8K 中（2,299 样本）
  • argilla/ultrafeedback-binarized-preferences-cleaned: 高评分选择答案 >=5（22,799 样本）
  • argilla/distilabel-math-preference-dpo: 高评分选择答案 >=9（2,181 样本）
  • unalignment/toxic-dpo-v0.2（541 样本）
  • M4-ai/prm_dpo_pairs_cleaned（7,958 样本）
  • jondurbin/truthy-dpo-v0.1（1,016 样本）

感谢 argilla、unalignment、M4-ai 和 jondurbin 提供源数据集。

按照惯例，让我们先安装所需的库：

pip install -U transformers datasets accelerate peft trl bitsandbytes wandb

安装完成后，我们可以导入必要的库并登录 W&B（可选）：

import gc
import os
import torch
import wandb
from datasets import load_dataset
from google.colab import userdata
from peft import LoraConfig, PeftModel, prepare_model_for_kbit_training
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    BitsAndBytesConfig,
    TrainingArguments,
    pipeline,
)
from trl import ORPOConfig, ORPOTrainer, setup_chat_format

wb_token = userdata.get('wandb')
wandb.login(key=wb_token)

如果您有一个最近的 GPU，您还可以使用 Flash Attention 库来替换默认的 eager attention 实现，以获得更高效的性能。

if torch.cuda.get_device_capability()[0] >= 8:
    !pip install -qqq flash-attn
    attn_implementation = "flash_attention_2"
    torch_dtype = torch.bfloat16
else:
    attn_implementation = "eager"
    torch_dtype = torch.float16

接下来，我们将加载 Llama 3 8B 模型，并使用 bitsandbytes 将其设置为 4-bit 精度。然后，我们使用 PEFT 设置 LoRA 配置用于 QLoRA。我还使用了方便的 setup_chat_format() 函数来修改模型和分词器以支持 ChatML。它自动应用这个聊天模板，添加特殊标记，并调整模型的嵌入层大小以匹配新词汇表大小。

注意，您需要提交请求以访问 meta-llama/Meta-Llama-3-8B 并登录到您的 Hugging Face 帐户。或者，您可以加载未门控的模型副本，如 NousResearch/Meta--Llama-3-8B。

# Model
base_model = "meta-llama/Meta-Llama-3-8B"
new_model = "OrpoLlama-3-8B"

# QLoRA config
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch_dtype,
    bnb_4bit_use_double_quant=True,
)

# LoRA config
peft_config = LoraConfig(
    r=16,
    lora_alpha=32,
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM",
    target_modules=['up_proj', 'down_proj', 'gate_proj', 'k_proj', 'q_proj', 'v_proj', 'o_proj']
)

# Load tokenizer
tokenizer = AutoTokenizer.from_pretrained(base_model)

# Load model
model = AutoModelForCausalLM.from_pretrained(
    base_model,
    quantization_config=bnb_config,
    device_map="auto",
    attn_implementation=attn_implementation
)
model, tokenizer = setup_chat_format(model, tokenizer)
model = prepare_model_for_kbit_training(model)

现在模型已经准备好进行训练，我们可以处理数据集。我们加载 mlabonne/orpo-dpo-mix-40k 并使用 apply_chat_template() 函数将 "chosen" 和 "rejected" 列转换为 ChatML 格式。注意，我这里只使用了 1,000 个样本，而不是整个数据集，因为运行时间会过长。

dataset_name = "mlabonne/orpo-dpo-mix-40k"
dataset = load_dataset(dataset_name, split="all")
dataset = dataset.shuffle(seed=42).select(range(1000))

def format_chat_template(row):
    row["chosen"] = tokenizer.apply_chat_template(row["chosen"], tokenize=False)
    row["rejected"] = tokenizer.apply_chat_template(row["rejected"], tokenize=False)
    return row

dataset = dataset.map(
    format_chat_template,
    num_proc= os.cpu_count(),
)
dataset = dataset.train_test_split(test_size=0.01)

首先，我们需要设置一些超参数：

  • learning_rate: ORPO 使用非常低的学习率相比传统 SFT 或甚至 DPO。这个值来自原始论文，大致对应 SFT 学习率 1e-5 和 DPO 学习率 5e-6。对于真正的微调，我建议将其增加到大约 1e-6。
  • beta: 它是论文中的 $\lambda$ 参数，默认值为 0.1。原始论文中的附录展示了如何通过消融研究选择该值。

其他参数，如 max_length 和 batch size 被设置为使用尽可能多的 VRAM（在此配置中约为 20 GB）。理想情况下，我们会训练模型 3-5 个周期，但这里我们只进行 1 个周期。

最后，我们可以使用 ORPOTrainer 来训练模型，它作为一个包装器。

orpo_args = ORPOConfig(
    learning_rate=8e-6,
    beta=0.1,
    lr_scheduler_type="linear",
    max_length=1024,
    max_prompt_length=512,
    per_device_train_batch_size=2,
    per_device_eval_batch_size=2,
    gradient_accumulation_steps=4,
    optim="paged_adamw_8bit",
    num_train_epochs=1,
    evaluation_strategy="steps",
    eval_steps=0.2,
    logging_steps=1,
    warmup_steps=10,
    report_to="wandb",
    output_dir="./results/",
)

trainer = ORPOTrainer(
    model=model,
    args=orpo_args,
    train_dataset=dataset["train"],
    eval_dataset=dataset["test"],
    peft_config=peft_config,
    tokenizer=tokenizer,
)
trainer.train()
trainer.save_model(new_model)

在 L4 GPU 上训练这 1,000 个样本大约花了两个小时。让我们检查 W&B 图表：

虽然损失下降，但选择和被拒绝答案之间的差异不明显：平均边际和准确性仅略高于零和0.5。

在原始论文中，作者在 Anthropic/hh-rlhf 数据集（161k 样本）上训练模型达10个周期，比我们的快速运行时间长得多。他们还尝试了 Llama 3，并友好地与我分享了他们的日志（感谢 Jiwoo Hong）。

为了结束本教程，让我们将 QLoRA 适配器与基础模型合并并推送到 Hugging Face Hub。

# Flush memory
del trainer, model
gc.collect()
torch.cuda.empty_cache()

# Reload tokenizer and model
tokenizer = AutoTokenizer.from_pretrained(base_model)
model = AutoModelForCausalLM.from_pretrained(
    base_model,
    low_cpu_mem_usage=True,
    return_dict=True,
    torch_dtype=torch.float16,
    device_map="auto",
)
model, tokenizer = setup_chat_format(model, tokenizer)

# Merge adapter with base model
model = PeftModel.from_pretrained(model, new_model)
model = model.merge_and_unload()
model.push_to_hub(new_model, use_temp_dir=False)
tokenizer.push_to_hub(new_model, use_temp_dir=False)

恭喜，我们完成了 Llama 3 的快速微调：mlabonne/OrpoLlama-3-8B。您可以使用 Hugging Face Space 与其互动（这里有一个 notebook 来制作您自己的）。虽然模型训练不足，如 W&B 曲线所示，但我在 Nous 的基准套件上使用 LLM AutoEval 进行了评估。

我们的 ORPO 微调实际上相当不错，并且在每个基准上都提高了基础模型的性能。这令人鼓舞，可能意味着在整个40k样本上进行微调会产生很好的结果。

这是开源社区的一个激动人心的时刻，越来越多高质量的开源权重模型正在发布。闭源和开源权重模型之间的差距正在逐渐缩小，而微调是为您的用例获得最佳性能的重要工具。

结论

在本文中，我们介绍了 ORPO 算法，并解释了它如何将 SFT 和偏好对齐阶段统一为一个过程。然后，我们使用 TRL 在自定义偏好数据集上对 Llama 3 8B 模型进行了微调。最终模型显示出令人鼓舞的结果，突显了 ORPO 作为新微调范式的潜力。

希望这对您有所帮助，并且我建议运行 Colab notebook 来微调您自己的 Llama 3 模型。在未来的文章中，我们将看到如何创建高质量的数据集——这一点通常被忽视。

相关阅读

ORPO算法：融合监督微调和偏好对齐的新型微调技术
使用PyTorch和Hugging Face工具包微调Llama 3模型以优化医疗对话系统
高效微调 Llama 3：使用 PyTorch FSDP 和 Q-Lora
ORPO 算法：结合模拟退火与粒子群优化的组合优化解决方案
Llama 3.1: 开源AI新纪元