StackLLaMA：使用RLHF訓練LLaMA的實戰指南

ChatGPT、GPT-4和Claude等模型是強大的語言模型，它們透過“人類反饋強化學習”（RLHF）方法進行了微調，以便更好地符合我們對它們的行為預期和使用方式。

在這篇部落格文章中，我們將展示使用RLHF結合以下方法訓練LlaMa模型以回答Stack Exchange問題的全部步驟：

• 監督式微調（SFT）
• 獎勵/偏好建模（RM）
• 來自人類反饋的強化學習（RLHF）

摘自InstructGPT論文：Ouyang, Long, et al. "Training language models to follow instructions with human feedback." arXiv preprint arXiv:2203.02155 (2022)。

透過結合這些方法，我們釋出了StackLLaMA模型。該模型可在🤗 Hub上獲取（原始LLaMA模型請參見Meta的LLaMA釋出），並且整個訓練管道作為Hugging Face TRL庫的一部分提供。為了讓您瞭解該模型的能力，請嘗試下面的演示！

LLaMA模型

進行RLHF時，從一個有能力的模型開始非常重要：RLHF步驟僅僅是一個微調步驟，旨在使模型與我們期望的互動方式和響應方式保持一致。因此，我們選擇使用最近推出且效能優異的LLaMA模型。LLaMA模型是Meta AI開發的最新大型語言模型。它們的大小範圍從7B到65B引數，並在1T到1.4T個token上進行了訓練，這使得它們非常強大。我們使用7B模型作為所有後續步驟的基礎！要訪問該模型，請使用Meta AI的表單。

Stack Exchange資料集

收集人類反饋是一項複雜且昂貴的任務。為了在此示例中引導該過程，同時仍構建一個有用的模型，我們使用了StackExchange資料集。該資料集包含來自StackExchange平臺（包括StackOverflow的程式碼和許多其他主題）的問題及其對應的答案。它對這種用例很有吸引力，因為答案附帶了點贊數和已接受答案的標籤。

我們遵循Askell et al. 2021中描述的方法，併為每個答案分配一個分數：

分數 = log2 (1 + 點贊數) 四捨五入到最近的整數，如果提問者接受了答案則加1（如果點贊數為負數，我們將其分數設為-1）。

對於獎勵模型，我們總是需要每道問題有兩個答案進行比較，我們稍後會看到。有些問題有幾十個答案，導致了許多可能的配對。我們每道問題最多采樣十個答案對，以限制每個問題的資料點數量。最後，我們透過將HTML轉換為Markdown來清理格式，以使模型的輸出更具可讀性。您可以在此處找到資料集和處理筆記本。

高效訓練策略

即使是訓練最小的LLaMA模型也需要巨大的記憶體。簡單計算一下：在bf16中，每個引數使用2位元組（fp32中是4位元組），此外Adam最佳化器還使用了8位元組（更多資訊請參閱Transformers中的效能文件）。因此，一個7B引數的模型僅在記憶體中就需要(2+8)*7B=70GB，並且在計算中間值（如注意力分數）時可能需要更多。所以，即使是單塊80GB的A100顯示卡也無法這樣訓練模型。你可以使用一些技巧，比如更高效的最佳化器或半精度訓練，以在記憶體中擠出更多空間，但遲早會耗盡。

另一個選項是使用引數高效微調（PEFT）技術，例如peft庫，它可以在載入為8位模型上執行低秩適應（LoRA）。

線性層上的低秩適應：額外的引數（橙色）被新增到凍結層（藍色）旁邊，並且由此產生的編碼隱藏狀態與凍結層的隱藏狀態一起新增。

以8位載入模型可以顯著減少記憶體佔用，因為權重每個引數只需要一個位元組（例如，7B Llama在記憶體中為7GB）。LoRA不是直接訓練原始權重，而是在一些特定層（通常是注意力層）之上新增小的介面卡層；因此，可訓練引數的數量大幅減少。

在這種情況下，一條經驗法則是為每十億引數分配約1.2-1.4GB（取決於批次大小和序列長度），以適應整個微調設定。正如上述部落格文章中詳述的那樣，這使得在NVIDIA A100 80GB上以低成本微調更大規模的模型（高達50-60B引數）成為可能。

這些技術使得在消費裝置和Google Colab上微調大型模型成為可能。值得注意的演示是在Google Colab上微調facebook/opt-6.7b（13GB，float16）和openai/whisper-large（15GB GPU RAM）。要了解更多關於使用peft的資訊，請參閱我們的GitHub倉庫或關於在消費硬體上訓練200億引數模型的上一篇部落格文章（https://huggingface.co/blog/trl-peft）。

現在我們可以將非常大的模型放入單個GPU中，但訓練可能仍然非常緩慢。在這種情況下，最簡單的策略是資料並行：我們將相同的訓練設定複製到獨立的GPU中，並向每個GPU傳遞不同的批次。透過這種方式，您可以並行化模型的前向/後向傳播，並隨著GPU數量的增加而擴充套件。

我們使用transformers.Trainer或accelerate，它們都支援資料並行，無需任何程式碼更改，只需在透過torchrun或accelerate launch呼叫指令碼時傳遞引數即可。以下分別透過accelerate和torchrun在單臺機器上使用8個GPU執行訓練指令碼。

accelerate launch --multi_gpu --num_machines 1  --num_processes 8 my_accelerate_script.py
torchrun --nnodes 1  --nproc_per_node 8 my_torch_script.py

監督式微調

在我們開始訓練獎勵模型和使用強化學習調整模型之前，如果模型在我們感興趣的領域已經表現良好，那將有所幫助。在我們的例子中，我們希望它能回答問題，而對於其他用例，我們可能希望它能遵循指令，在這種情況下，指令調優是一個好主意。實現這一目標最簡單的方法是繼續使用來自該領域或任務的文字進行語言模型訓練，並採用語言建模目標。StackExchange資料集非常龐大（超過1000萬條指令），因此我們可以輕鬆地在其子集上訓練語言模型。

在進行RLHF之前對模型進行微調沒有什麼特別之處——我們只是在這裡應用了預訓練中的因果語言建模目標。為了高效地利用資料，我們使用了一種稱為“打包”（packing）的技術：在批處理中，我們不是每個樣本一個文字然後填充到最長文字或模型最大上下文，而是將許多文字用EOS標記連線起來，並切割上下文大小的塊來填充批處理，無需任何填充。

透過這種方法，訓練效率大大提高，因為模型中傳遞的每個token都經過訓練，而填充token通常會從損失中遮蔽掉。如果您資料不多，並且更擔心偶爾會截斷一些超出上下文的token，您也可以使用傳統的DataLoader。

打包由ConstantLengthDataset處理，然後我們可以在使用peft載入模型後使用Trainer。首先，我們以int8載入模型，準備訓練，然後新增LoRA介面卡。

# load model in 8bit
model = AutoModelForCausalLM.from_pretrained(
        args.model_path,
        load_in_8bit=True,
        device_map={"": Accelerator().local_process_index}
    )
model = prepare_model_for_int8_training(model)

# add LoRA to model
lora_config = LoraConfig(
    r=16,
    lora_alpha=32,
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM",
)

model = get_peft_model(model, lora_config)

我們使用因果語言建模目標訓練模型數千步，並儲存模型。由於我們將再次使用不同目標調整模型，因此我們將介面卡權重與原始模型權重合並。

免責宣告：由於LLaMA的許可，我們僅釋出了本節以及後續章節中模型檢查點的介面卡權重。您可以透過填寫Meta AI的表單申請基礎模型權重的訪問許可權，然後透過執行此指令碼將其轉換為🤗 Transformers格式。請注意，在v4.28釋出之前，您還需要從原始碼安裝🤗 Transformers。

現在我們已經針對任務對模型進行了微調，我們可以訓練一個獎勵模型了。

獎勵建模和人類偏好

原則上，我們可以直接使用人工標註透過RLHF微調模型。然而，這將需要在每次最佳化迭代後將一些樣本傳送給人類進行評分。由於收斂所需的訓練樣本數量以及人類閱讀和標註速度的固有延遲，這既昂貴又緩慢。

一種替代直接反饋的有效技巧是在RL迴圈之前，根據人類標註訓練一個獎勵模型。獎勵模型的目標是模仿人類如何評價文字。構建獎勵模型有幾種可能的策略：最直接的方法是預測標註（例如，評分或“好”/“壞”的二元值）。實際上，更好的方法是預測兩個示例的排序，即獎勵模型接收給定提示$x$的兩個候選$(y_k, y_j)$，並預測人類標註者會認為哪個更好。

這可以轉化為以下損失函式

$\operatorname{loss}(\theta)=- E_{\left(x, y_j, y_k\right) \sim D}\left[\log \left(\sigma\left(r_\theta\left(x, y_j\right)-r_\theta\left(x, y_k\right)\right)\right)\right]$

其中$r$是模型的得分，$y_j$是首選候選。

利用StackExchange資料集，我們可以根據分數推斷出兩個答案中哪個更受使用者偏好。有了這些資訊和上述定義的損失，我們就可以透過新增自定義損失函式來修改transformers.Trainer。

class RewardTrainer(Trainer):
    def compute_loss(self, model, inputs, return_outputs=False):
        rewards_j = model(input_ids=inputs["input_ids_j"],  attention_mask=inputs["attention_mask_j"])[0]
        rewards_k = model(input_ids=inputs["input_ids_k"], attention_mask=inputs["attention_mask_k"])[0]
        loss = -nn.functional.logsigmoid(rewards_j - rewards_k).mean()
        if return_outputs:
            return loss, {"rewards_j": rewards_j, "rewards_k": rewards_k}
        return loss

我們利用了10萬對候選子集，並在5萬個保留集上進行評估。在批處理大小為4的情況下，我們使用LoRA peft介面卡對LLaMA模型進行單次epoch訓練，使用Adam最佳化器和BF16精度。我們的LoRA配置是：

peft_config = LoraConfig(
    task_type=TaskType.SEQ_CLS,
    inference_mode=False,
    r=8,
    lora_alpha=32,
    lora_dropout=0.1,
)

訓練透過Weights & Biases進行日誌記錄，在8塊A100 GPU上使用🤗研究叢集花費了幾個小時，模型最終達到了67%的準確率。雖然這聽起來分數不高，但這項任務本身也非常困難，即使對於人類標註者來說也是如此。

如下一節所述，生成的介面卡可以合併到凍結模型中，並儲存以供後續使用。

來自人類反饋的強化學習

有了微調過的語言模型和獎勵模型，我們現在可以執行RL迴圈了。它大致遵循三個步驟：

1. 從提示中生成響應
2. 使用獎勵模型對響應進行評分
3. 使用評分執行強化學習策略最佳化步驟

查詢和響應提示在分詞並傳遞給模型之前，模板化如下：

Question: <Query>

Answer: <Response>

SFT、RM和RLHF階段均使用相同的模板。

使用強化學習訓練語言模型的一個常見問題是，模型可能透過生成完全無意義的內容來利用獎勵模型，這會導致獎勵模型分配高獎勵。為了平衡這一點，我們對獎勵進行懲罰：我們保留一個不進行訓練的模型作為參考，透過計算KL散度來比較新模型生成的內容和參考模型生成的內容。

$\operatorname{R}(x, y)=\operatorname{r}(x, y)- \beta \operatorname{KL}(x, y)$

其中$r$是來自獎勵模型的獎勵，$\operatorname{KL}(x,y)$是當前策略與參考模型之間的KL散度。

我們再次利用peft進行記憶體高效訓練，這在RLHF環境中提供了額外的優勢。在這裡，參考模型和策略共享同一個基礎，即SFT模型，我們在訓練期間將其載入為8位並凍結。我們專門使用PPO最佳化策略的LoRA權重，同時共享基礎模型的權重。

for epoch, batch in tqdm(enumerate(ppo_trainer.dataloader)):
    question_tensors = batch["input_ids"]
        
    # sample from the policy and generate responses
    response_tensors = ppo_trainer.generate(
        question_tensors,
        return_prompt=False,
        length_sampler=output_length_sampler,
        **generation_kwargs,
    )
    batch["response"] = tokenizer.batch_decode(response_tensors, skip_special_tokens=True)

    # Compute sentiment score
    texts = [q + r for q, r in zip(batch["query"], batch["response"])]
    pipe_outputs = sentiment_pipe(texts, **sent_kwargs)
    rewards = [torch.tensor(output[0]["score"] - script_args.reward_baseline) for output in pipe_outputs]

    # Run PPO step
    stats = ppo_trainer.step(question_tensors, response_tensors, rewards)
    # Log stats to WandB
    ppo_trainer.log_stats(stats, batch, rewards)

我們使用🤗研究叢集在3塊A100-80GB GPU上訓練了20小時，但您也可以更快地獲得不錯的結果（例如，在8塊A100 GPU上約20小時後）。所有訓練執行的統計資料都可以在Weights & Biases上找到。

訓練過程中每批次的獎勵。模型效能在大約1000步後趨於平穩。

那麼，訓練後的模型能做什麼呢？我們來看看！

雖然我們目前不應該完全相信它在LLaMA方面的建議，但這個答案看起來很連貫，甚至提供了一個谷歌連結。接下來，我們來看看一些訓練挑戰。

挑戰、不穩定性及變通方法

使用RL訓練LLM並非一帆風順。我們今天演示的模型是許多實驗、失敗執行和超引數搜尋的結果。即便如此，該模型也遠非完美。在這裡，我們將分享我們在這個示例製作過程中遇到的一些觀察和令人頭痛的問題。

更高的獎勵意味著更好的效能，對嗎？

哇，這次執行一定很棒，看看這甜美的獎勵！

通常在強化學習中，你希望獲得最高的獎勵。在RLHF中，我們使用一個不完美的獎勵模型，而PPO演算法一旦有機會就會利用這些不完美之處。這可能表現為獎勵的突然增加，然而當我們檢視策略生成的文字時，它們大多包含“```”字串的重複，因為獎勵模型發現包含程式碼塊的Stack Exchange答案通常比沒有程式碼塊的答案排名更高。幸運的是，這個問題很少出現，而且通常KL懲罰應該可以抵消這種利用行為。

KL值總是正數，不是嗎？

正如我們之前提到的，KL懲罰項用於使模型的輸出保持接近基礎策略。一般來說，KL散度衡量兩個分佈之間的距離，並且總是一個正值。然而，在trl中，我們使用KL的估計值，該估計值在期望上等於真實的KL散度。

$KL_{pen}(x,y) = \log \left(\pi_\phi^{\mathrm{RL}}(y \mid x) / \pi^{\mathrm{SFT}}(y \mid x)\right)$

顯然，當從策略中取樣到的token的機率低於SFT模型時，這將導致負的KL懲罰，但平均而言，它將是正的，否則您將無法正確地從策略中取樣。然而，一些生成策略可以強制生成某些token，或者某些token可以被抑制。例如，在批次生成時，完成的序列會被填充，並且在設定最小長度時，EOS token會被抑制。模型可以為這些token分配非常高或低的機率，這會導致負KL值。由於PPO演算法最佳化獎勵，它會追逐這些負懲罰，從而導致不穩定性。

在生成響應時需要小心，我們建議在採用更復雜的生成方法之前，始終先使用簡單的取樣策略。

持續存在的問題

仍有許多問題需要我們更好地理解和解決。例如，損失有時會出現尖峰，這可能導致進一步的不穩定性。

隨著我們識別和解決這些問題，我們將把更改上游到trl，以確保社群能夠從中受益。

結論

在這篇文章中，我們介紹了RLHF的整個訓練週期，從準備帶有“人類標註”的資料集開始，接著將語言模型適應到特定領域，訓練獎勵模型，最後使用RL訓練模型。

透過使用peft，任何人都可以使用單個GPU執行我們的示例！如果訓練太慢，您可以使用資料並行，無需更改程式碼，只需新增更多GPU即可擴充套件訓練。

對於實際應用，這只是第一步！一旦您訓練好一個模型，您必須對其進行評估，並將其與其他模型進行比較，以瞭解其效能如何。這可以透過對不同模型版本的生成結果進行排名來完成，類似於我們構建獎勵資料集的方式。

一旦您添加了評估步驟，樂趣就開始了：您可以開始迭代資料集和模型訓練設定，看看是否有改進模型的方法。您可以新增其他資料集，或者對現有資料集應用更好的過濾器。另一方面，您可以嘗試不同大小和架構的獎勵模型，或者訓練更長時間。

我們正在積極改進TRL，以使RLHF中涉及的所有步驟更易於訪問，並很高興看到人們用它構建的東西！如果您有興趣貢獻，請檢視GitHub上的問題。

引用

@misc {beeching2023stackllama,
    author       = { Edward Beeching and
                     Younes Belkada and
                     Kashif Rasul and
                     Lewis Tunstall and
                     Leandro von Werra and
                     Nazneen Rajani and
                     Nathan Lambert
                   },
    title        = { StackLLaMA: An RL Fine-tuned LLaMA Model for Stack Exchange Question and Answering },
    year         = 2023,
    url          = { https://huggingface.co/blog/stackllama },
    doi          = { 10.57967/hf/0513 },
    publisher    = { Hugging Face Blog }
}

致謝

我們感謝Philipp Schmid分享了他精彩的流式文字生成演示，我們的演示就是基於此。我們還要感謝Omar Sanseviero和Louis Castricato對部落格文章草稿提供了寶貴而詳細的反饋。