“StackLLaMA”: 用 RLHF 訓練 LLaMA 的手把手教程

HuggingFace發表於2023-05-05

如 ChatGPT，GPT-4，Claude 語言模型 之強大，因為它們採用了 基於人類反饋的強化學習 (Reinforcement Learning from Human Feedback, RLHF) 來使之更符合我們的使用場景。

本部落格旨在展示用 RLHF 訓練一個 LLaMA 模型，以回答 Stack Exchange 上的問題。具體而言，包含以下幾個方面:

  • 有監督的微調 (Supervised Fine-tuning，SFT)。
  • 獎勵 / 偏好建模 (Reward / preference modeling，RM)。
  • 基於人類反饋的強化學習 (RLHF)。

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

結合了上述方法，我們釋出了 StackLLaMA 模型，該模型在 ? Hub 上開源 (訪問連結檢視 Meta 的原始 LLaMA )，整個 訓練的流程 已經整合到了 Hugging Face TRL 庫中 。你可以透過下面的 demo 來嘗試該模型。

LLaMA 模型

在實踐 RLHF 時，選取一個合適的模型很重要: RLHF 只是一個讓模型滿足我們互動形式的需求的微調過程 。所以我們選取了最近上線的 LLaMA 模型。LLaMA 模型是 Mata AI 最近推出的大語言模型。其引數量大小涵蓋 7B 到 65B，以及訓練在 1T 和 1.4T 的 token 上，這讓其很實用。我們這裡採用 7B 的模型。(請填寫 Meta AI 的這份 表單 來下載模型)。

Stack Exchange 資料集

收集人類的反饋資料集是很複雜且昂貴的勞動。為了做到這個，並且還能保證模型的有效性，我們使用 StackExchange 資料集。該資料集涵蓋了 StackExchange 平臺上的問題和答案 (包含 StackOverflow 的程式設計等話題下的)。這很適合我們的實踐，因為其包含了每個答案的贊和踩的數量。

我們按照 Askell et al. 2021 中的方法，給每個答案賦分:

score = log2 (1 + upvotes) rounded to the nearest integer, plus 1 if the questioner accepted the answer (we assign a score of −1 if the number of upvotes is negative).

對獎勵模型，我們將看到每個問題總是需要兩個答案對比。有些問題有很多答案，可以產生很多對，我們只取十個以限制每個問題的資料量。最後，我們把格式從 HTML 轉化到 Markdown 以提高輸出的可讀性。你可以看到資料集和處理過程的 [筆記本]。(https://huggingface.co/datasets/lvwerra/stack-exchange-paired。)

高效訓練策略

即使是最小 LLaMA 模型的訓練，都需要大量記憶體。估算一下: 以 bf16 半精度，每個引數用 2 個位元組 (以 fp32 精度四位元組的標準)，訓練時需要 8 個位元組 (例如 Adam 最佳化器，參見 Tramsformers 的 效能文件)。可見 7B 引數量的模型將用 (2+8)* 7B = 70 GB 的記憶體，並且還可能需要更多用於計算諸如注意力分數的中間值。所以很難在一張 80GB 視訊記憶體的 A100 上訓練。或許你可以使用一些技巧，比如用更高效的半精度訓練的最佳化器來壓縮記憶體，但溢位是遲早的。

另外的可能是 引數高效的微調(Parameter-Efficient Fine-Tuning, PEFT) 技術，比如 peft 庫，它可以對使用 8-bit 載入的模型做 低秩最佳化(Low-Rank Adaptation，LoRA)。

線性層的低秩最佳化: 額外引數 (橙色) 被加在 Frozen 層 (藍色)，編碼後的隱藏狀態與 Frozen 層的隱藏狀態疊加在一起。

以 8bit 載入模型會大幅降低記憶體佔用，因為每個引數只要一位元組 (比如 7B LLaMA 是 7GB 記憶體)。與直接訓練原始模型不同，LoRA 在特定層 (一般是注意力層) 新增少量新引數，大幅降低了需要訓練的引數。

此情此景，一個衡量標準是 1B 的引數在整個微調過程中佔 ~1.2-1.4GB (和具體 batch size 及序列長度有關)。在參考的部落格中具體討論了，這使得低成本下微調較大引數規模的模型成為可能 (比如在一張 A100 上微調 50-60B 的引數)。

這些技術能讓微調大模型的任務，在消費級裝置和 Google Colab 上執行。這裡提供一些值得關注的演示 demo: facebook/opt-6.7b (在 float16 精度下 13GB) 和 openai/whisper-large
跑在 Google Colab (15GB 視訊記憶體) 上。欲瞭解 peft 的使用，請參見 github 倉庫 或者之前的 部落格介紹: 在客戶端訓練 20B 引數量的模型。

現在我們能在一張 GPU 上微調很大的模型了，但訓練還是會很慢。此時最簡單的策略便是並行化: 把一個訓練同時放到不同的 GPU 上，各 GPU 接受不同的 batch。這樣我們可以並行執行前向傳播和後向傳播，透過增加 GPU 的數量實現並行能力提升。

我們可以選用 trainsformers.Trainer 或 accelerate，因為它們都支援無程式碼變更進行資料並行化。只需注意呼叫 torchrun 或者 accelerate launch 指令碼時的引數即可實現。比如以下就是在一個 8 顯示卡的機器上分別用 accelerate launch 和 torchrun的方法:

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

有監督的微調

在訓練獎勵模型和用 RL 之前，模型若是已經在我們感興趣的方面表現好將會很有幫助。在我們的示例中，我們想要其能回答問題，而其他時候，我們可能它能聽指令 (這時對指令執行的微調是理想的)。實現這個最簡單的方法便是面向該語言任務，用該任務和領域的文字，繼續訓練。StackExchange 資料集 含 10M 的指令量，所以我們能用其子集很容易地訓練。

在用 RLHF 之前的模型微調沒有特別的，就是一般的面向語言任務的預訓練模型微調。為了高效利用資料，我們採用了稱之為 打包 的技術: 與 batch 中的每個樣本均由單一文字組成，最後基於最長的文字來 padding (填充)，我們把很多文字拼接起來，用 EOS token 來隔開，然後分割成一些 chunk (切塊) 來做成 batch，避免 padding。

該方法大大提高了效率，因為模型輸入的所有 token 都對 loss 有所訓練，而非 padding 作為掩碼被丟棄了。如果你沒有足夠資料，並且擔心隨意地分開 token 會失去上下文語義，你也可以用傳統的資料載入器
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, config)

我們根據相應的語言任務，對模型訓練幾千個 step (步)，並儲存模型。由於我們將會有其他微調模型的目的，我們將 LoRA 的微調器權重合併到原模型中。

宣告: 因為 LLaMA 的許可證規定，我們只能釋出微調器的權重，你需要填 Meta AI 的 表格 來獲取模型，然後用這個 指令碼 來轉成 ? Transformers 格式。注意 ? Transformers 應該從原始碼安裝，或者 v4.28 版。

現在我們已經微調好了模型，可以訓練獎勵模型了。

獎勵模型和人類偏好

原則上，我們可以直接用人類標註來對模型做 RLHF 微調。然而，這將需要我們給人類傳送一些樣本，在每輪最佳化後計分。這是貴且慢的，因為收斂需要的訓練樣本量大，而人類閱讀和標註的速度有限。

一個比直接反饋更好的策略是，在進入 RL 迴圈之前用人類標註集來訓練一個獎勵模型。獎勵模型的目的是模擬人類對文字的打分。構建獎勵模型有許多能用的策略: 最直接的便是預測標註 (比如根據好與壞，輸出比分或者布林值)。最佳實踐是，預測結果的排序，即對每個 prompt (輸入文字) 對應的兩個結果 \( (y_k, y_j) \)模型預測人類標註的比分哪個更高。

或者表示為 loss (損失) 函式:

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

其中\( r \)是模型對可能的標註\( y_j \)的預測分數。

在 StackExchange 資料集上，我們能得到兩個答案的受歡迎程度。有了這個資訊和上面的損失函式，我們就能自定義 loss 來改 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

我們用資料集中的 100000 對，並在 50000 對上評估。在比較小的 batch size，為 4 下，我們用 LoRA 的 peft 微調器來訓練 LLaMA 模型，在 BF16 精度下用 Adam 最佳化器。我們的 LoRA 設定是:

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

訓練用 Weights & Biases 來記日誌，並在 ? 訓練叢集上，用 8 卡 A-100，要數小時，最後準確率為 67%。儘管看上去可能低了，但想想這個任務的難度。

如下文要細說的，訓練結果將作為固定引數，以供下游使用。

基於人類反饋的強化學習

現在我們手頭有了微調的語言模型和獎勵模型，可以開始執行 RL 迴圈了: 這個過程大致分為三步

  1. 生成對 prompt (輸入文字) 的反饋。
  2. 用獎勵模型來對反饋評分。
  3. 對評分，進行一輪策略最佳化的強化學習。

在被 token 化並輸入獎勵模型前，提問和回答的 prompt 模版如下:

Question: <Query>
Answer: <Response>

在有監督訓練 (SFT)，獎勵模型訓練 (RM) 和 RLHF 的階段都用此模版。

用 RL 訓練語言模型出現的常見問題是，模型可能學會胡說八道以糊弄獎勵模型，後者可能給高分。為了權衡，我們對獎勵增加懲罰: 留一份沒有訓練的模型，如何比較兩者輸出的 KL 散度

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

其中 \( r \) 是獎勵模型的結果，\( {KL}(x,y) \) 是當前模型和對比模型的 KL 散度差。

再提一遍，我們用 peft 來實現記憶體高效的訓練，其對 RLHF 階段提供了優勢。這裡參考的模型和訓練的模型用同一個基底，也就是有監督訓練 (SFT) 的結果，它是用 8-bit 來載入，並且自始自終是固定的。我們僅用 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)

我們用 ? 叢集，在 3x8 A100-80GB 的機器上訓練了 20h，但一個差不多的結果很快 (大概，在 8 A100-80GB 上訓練 20h)。所有的訓練過程都在 Weight & Biases 上找到。

每個 batch 的獎勵，對每步的訓練，在 ~1000 步時模型的效果最好。

所以模型訓好了能幹啥嘞 ? 我們拭目以待 !

儘管我們不該太相信其結果，至少目前。但結果已經很好了，甚至附上了 Google 連結。我們來看看訓練時的挑戰。

挑戰，不穩定和突破口

用 RL 訓練 LLM (Large Language Models，大語言模型) 不總是一帆風順的，你看到的本文也是經歷無數實驗，無數失敗和無數調參的。即便如此，該模型也不能說變現完美。這兒，我們分享一些遇到的觀察和問題。

獎勵更高代表更好表現?

天吶，這個實驗肯定表現很好 ! 看獎勵的曲線多甜啊 !

在 RL 中，一般而言，獎勵越高越好。在 RLHF 中，我們用了一個獎勵模型，它不完美，所以留給了 PPO 演算法撿漏的機會。這能導致獎勵突然上升，然而當檢查文字結果時，卻充斥了字元 “`”，因為獎勵模型對含有程式碼 stack exchange 的答案更信任。幸運的是，該問題碰到的很少，應該是採取的 KL 散度的懲罰項起到了作用。

KL 散度總是正的?

如我們前面所提到的，一個 KL 懲罰項被用來保證訓練後的分佈和原始分佈接近。一般地 , KL 散度來度量兩個分佈的相似程度，並且總是正的。然而，在 trl 我們用了一個 KL 的近似，期望值和真的 KL 散度相同。

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

顯然，當訓練中一個 token 比原始模型機率低，這會導致 KL 散度為負，合適的取樣和平均總能得到正的。但是一些取樣的生成策略導致了不勻稱的取樣。比如，當生成被 padding 的序列 batch 時和當設定 EOS token 被壓縮的最小長度是，模型會有很大/很小的機率到負 KL 散度的 token。同時 PPO 演算法是面向獎勵最佳化的，模型就會追逐負的懲罰，導致訓練不穩定。

對生成和取樣，你需要特別小心。我們建議一開始用最簡單的方式，如何在逐漸複雜。

任然存在的問題

任然有很多問題我們不懂，比如下面，loss 間斷地跳躍，導致之後的不穩定

一旦我們解決了這些問題，我們就會上傳變化到 trl 上，以保證社群受益。

總結

在本部落格，我們走過了 RLHF 訓練的整個流程，從準備人類標註的資料集開始，調整語言模型到特定領域，訓練獎勵模型，並最終用 RL 訓練一個模型。

透過使用 peft，任何人都能在一張 GPU 上跑我們的實驗 ! 如果訓練慢了，可以用資料並行化的方法，不需要改任何程式碼，或者用多張 GPU 並行提高訓練速度。

對實際應用，這僅僅是第一步 ! 一旦你有了模型，你就要和其他模型比較優劣。這個可以用一個面向不同模型的排名生成做到，和我們訓練獎勵資料集類似。

一旦你加入了評估的步驟，好玩的就開始了: 你可以在原資料集上反覆煉丹，也可以增加資料集或者對原資料集提純。另外，你可以對獎勵模型和生成試不同大小和結構的模型，這需要時間。

我們在積極提高 TRL 以保證 RLHF 的每一步都可見，並且十分激動能看到人們用它來構建的東西。如果你想有所貢獻，歡迎看我們的 Github Issue。

引用

@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 分享了他對文字生成絕妙的 demo, 我們的 demo 也是基於他的。我們也感謝 Omar Sanseviero 和 Louis Castricato 對我們部落格的草稿提供寶貴詳盡的反饋。

英文原文: https://hf.co/blog/stackllama

作者: Edward Beeching, Kashif Rasul, Younes Belkada, Lewis Tunstall, Leandro von Werra Nazneen Rajani, Nathan Lambert

譯者: Vermillion-Qi(張奇), zhongdongy (阿東)

相關文章