使用 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。