如 ChatGPT，GPT-4，Claude?這樣的語言模型之所以強大，是因為它們采用了?基于人類反饋的強化學(xué)習(xí)?(Reinforcement Learning from Human Feedback, RLHF) 來使之更符合我們的使用場景。

本博客旨在展示用 RLHF 訓(xùn)練一個 LLaMA 模型，以回答 Stack Exchange 上的問題。具體而言，包含以下幾個方面:

• 有監(jiān)督的微調(diào) (Supervised Fine-tuning，SFT)。

• 獎勵 / 偏好建模 (Reward / preference modeling，RM)。

• 基于人類反饋的強化學(xué)習(xí) (RLHF)。

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

結(jié)合了上述方法，我們發(fā)布了 StackLLaMA 模型，該模型在 ?? Hub 上開源 (訪問鏈接查看 Meta 的原始 LLaMA )，整個 訓(xùn)練的流程 已經(jīng)集成到了 Hugging Face TRL 庫中 。你可以通過下面的 demo 來嘗試該模型。

LLaMA 模型

在實踐 RLHF 時，選取一個合適的模型很重要: RLHF 只是一個讓模型滿足我們交互形式的需求的微調(diào)過程 。所以我們選取了最近上線的 LLaMA 模型。LLaMA 模型是 Mata AI 最近推出的大語言模型。其參數(shù)量大小涵蓋 7B 到 65B，以及訓(xùn)練在 1T 和 1.4T 的 token 上，這讓其很實用。我們這里采用 7B 的模型。(請?zhí)顚?Meta AI 的這份 表單 來下載模型)。

Stack Exchange 數(shù)據(jù)集

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

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

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

高效訓(xùn)練策略

即使是最小 LLaMA 模型的訓(xùn)練，都需要大量內(nèi)存。估算一下: 以 bf16 半精度，每個參數(shù)用 2 個字節(jié) (以 fp32 精度四字節(jié)的標(biāo)準(zhǔn))，訓(xùn)練時需要 8 個字節(jié) (例如 Adam 優(yōu)化器，參見 Tramsformers 的 性能文檔)。可見 7B 參數(shù)量的模型將用 (2+8)* 7B = 70 GB 的內(nèi)存，并且還可能需要更多用于計算諸如注意力分?jǐn)?shù)的中間值。所以很難在一張 80GB 顯存的 A100 上訓(xùn)練。或許你可以使用一些技巧，比如用更高效的半精度訓(xùn)練的優(yōu)化器來壓縮內(nèi)存，但溢出是遲早的。

另外的可能是?參數(shù)高效的微調(diào)(Parameter-Efficient Fine-Tuning, PEFT) 技術(shù)，比如?peft?庫，它可以對使用 8-bit 加載的模型做?低秩優(yōu)化(Low-Rank Adaptation，LoRA)。

線性層的低秩優(yōu)化: 額外參數(shù) (橙色) 被加在 Frozen 層 (藍色)，編碼后的隱藏狀態(tài)與 Frozen 層的隱藏狀態(tài)疊加在一起。

以 8bit 加載模型會大幅降低內(nèi)存占用，因為每個參數(shù)只要一字節(jié) (比如 7B LLaMA 是 7GB 內(nèi)存)。與直接訓(xùn)練原始模型不同，LoRA 在特定層 (一般是注意力層) 添加少量新參數(shù)，大幅降低了需要訓(xùn)練的參數(shù)。

此情此景，一個衡量標(biāo)準(zhǔn)是 1B 的參數(shù)在整個微調(diào)過程中占 ~1.2-1.4GB (和具體 batch size 及序列長度有關(guān))。在參考的博客中具體討論了，這使得低成本下微調(diào)較大參數(shù)規(guī)模的模型成為可能 (比如在一張 A100 上微調(diào) 50-60B 的參數(shù))。

這些技術(shù)能讓微調(diào)大模型的任務(wù)，在消費級設(shè)備和 Google Colab 上執(zhí)行。這里提供一些值得關(guān)注的演示 demo:?facebook/opt-6.7b?(在 float16 精度下 13GB) 和?openai/whisper-large跑在 Google Colab (15GB 顯存) 上。欲了解?peft?的使用，請參見 github 倉庫 或者之前的 博客介紹: 在客戶端訓(xùn)練 20B 參數(shù)量的模型。

現(xiàn)在我們能在一張 GPU 上微調(diào)很大的模型了，但訓(xùn)練還是會很慢。此時最簡單的策略便是并行化: 把一個訓(xùn)練同時放到不同的 GPU 上，各 GPU 接受不同的 batch。這樣我們可以并行執(zhí)行前向傳播和后向傳播，通過增加 GPU 的數(shù)量實現(xiàn)并行能力提升。

我們可以選用?trainsformers.Trainer?或?accelerate，因為它們都支持無代碼變更進行數(shù)據(jù)并行化。只需注意調(diào)用?torchrun?或者?accelerate launch?腳本時的參數(shù)即可實現(xiàn)。比如以下就是在一個 8 顯卡的機器上分別用?accelerate launch?和?torchrun的方法:

有監(jiān)督的微調(diào)

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

在用 RLHF 之前的模型微調(diào)沒有特別的，就是一般的面向語言任務(wù)的預(yù)訓(xùn)練模型微調(diào)。為了高效利用數(shù)據(jù)，我們采用了稱之為?打包?的技術(shù): 與 batch 中的每個樣本均由單一文本組成，最后基于最長的文本來 padding (填充)，我們把很多文本拼接起來，用 EOS token 來隔開，然后分割成一些 chunk (切塊) 來做成 batch，避免 padding。

該方法大大提高了效率，因為模型輸入的所有 token 都對 loss 有所訓(xùn)練，而非 padding 作為掩碼被丟棄了。如果你沒有足夠數(shù)據(jù)，并且擔(dān)心隨意地分開 token 會失去上下文語義，你也可以用傳統(tǒng)的數(shù)據(jù)加載器ConstantLengthDataset?解決了?打包技術(shù)，并且我們能在用?peft?加載模型后用?Trainer。首先，我們用?int8?加載模型，準(zhǔn)備訓(xùn)練，然后加入?LoRA?微調(diào)器。

我們根據(jù)相應(yīng)的語言任務(wù)，對模型訓(xùn)練幾千個 step (步)，并保存模型。由于我們將會有其他微調(diào)模型的目的，我們將 LoRA 的微調(diào)器權(quán)重合并到原模型中。

聲明: 因為 LLaMA 的許可證規(guī)定，我們只能發(fā)布微調(diào)器的權(quán)重，你需要填 Meta AI 的 表格 來獲取模型，然后用這個 腳本 來轉(zhuǎn)成 ?? Transformers 格式。注意 ?? Transformers 應(yīng)該從源碼安裝，或者?v4.28?版。

現(xiàn)在我們已經(jīng)微調(diào)好了模型，可以訓(xùn)練獎勵模型了。

獎勵模型和人類偏好

原則上，我們可以直接用人類標(biāo)注來對模型做 RLHF 微調(diào)。然而，這將需要我們給人類發(fā)送一些樣本，在每輪優(yōu)化后計分。這是貴且慢的，因為收斂需要的訓(xùn)練樣本量大，而人類閱讀和標(biāo)注的速度有限。

一個比直接反饋更好的策略是，在進入 RL 循環(huán)之前用人類標(biāo)注集來訓(xùn)練一個獎勵模型。獎勵模型的目的是模擬人類對文本的打分。構(gòu)建獎勵模型有許多能用的策略: 最直接的便是預(yù)測標(biāo)注 (比如根據(jù)好與壞，輸出比分或者布爾值)。最佳實踐是，預(yù)測結(jié)果的排序，即對每個 prompt (輸入文本) 對應(yīng)的兩個結(jié)果?，模型預(yù)測人類標(biāo)注的比分哪個更高。

或者表示為 loss (損失) 函數(shù):

其中 r?是模型對可能的標(biāo)注 yj?的預(yù)測分?jǐn)?shù)。

在 StackExchange 數(shù)據(jù)集上，我們能得到兩個答案的受歡迎程度。有了這個信息和上面的損失函數(shù)，我們就能自定義 loss 來改?transformers.Trainer?了。

我們用數(shù)據(jù)集中的 100000 對，并在 50000 對上評估。在比較小的 batch size，為 4 下，我們用 LoRA 的 ?peft?微調(diào)器來訓(xùn)練 LLaMA 模型，在 BF16 精度下用 Adam 優(yōu)化器。我們的 LoRA 設(shè)置是:

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


訓(xùn)練用 Weights & Biases 來記日志，并在 ?? 訓(xùn)練集群上，用 8 卡 A-100，要數(shù)小時，最后準(zhǔn)確率為 **67%**。盡管看上去可能低了，但想想這個任務(wù)的難度。

如下文要細說的，訓(xùn)練結(jié)果將作為固定參數(shù)，以供下游使用。

基于人類反饋的強化學(xué)習(xí)

現(xiàn)在我們手頭有了微調(diào)的語言模型和獎勵模型，可以開始執(zhí)行 RL 循環(huán)了: 這個過程大致分為三步

1. 生成對 prompt (輸入文本) 的反饋。

2. 用獎勵模型來對反饋評分。

3. 對評分，進行一輪策略優(yōu)化的強化學(xué)習(xí)。

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

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

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

其中? r 是獎勵模型的結(jié)果， KL(x,y) 是當(dāng)前模型和對比模型的 KL 散度差。

再提一遍，我們用?peft?來實現(xiàn)內(nèi)存高效的訓(xùn)練，其對 RLHF 階段提供了優(yōu)勢。這里參考的模型和訓(xùn)練的模型用同一個基底，也就是有監(jiān)督訓(xùn)練 (SFT) 的結(jié)果，它是用 8-bit 來加載，并且自始自終是固定的。我們僅用 PPO 方法優(yōu)化最終模型的 LoRA 權(quán)重，同時全部共享一個基底模型。

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

每個 batch 的獎勵，對每步的訓(xùn)練，在 ?~1000 步時模型的效果最好。

所以模型訓(xùn)好了能干啥嘞 ? 我們拭目以待 !

盡管我們不該太相信其結(jié)果，至少目前。但結(jié)果已經(jīng)很好了，甚至附上了 Google 鏈接。我們來看看訓(xùn)練時的挑戰(zhàn)。

挑戰(zhàn)，不穩(wěn)定和突破口

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

獎勵更高代表更好表現(xiàn)?

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

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

KL 散度總是正的?

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

顯然，當(dāng)訓(xùn)練中一個 token 比原始模型概率低，這會導(dǎo)致 KL 散度為負(fù)，合適的取樣和平均總能得到正的。但是一些采樣的生成策略導(dǎo)致了不勻稱的采樣。比如，當(dāng)生成被 padding 的序列 batch 時和當(dāng)設(shè)置 EOS token 被壓縮的最小長度是，模型會有很大/很小的概率到負(fù) KL 散度的 token。同時 PPO 算法是面向獎勵優(yōu)化的，模型就會追逐負(fù)的懲罰，導(dǎo)致訓(xùn)練不穩(wěn)定。

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

任然存在的問題

任然有很多問題我們不懂，比如下面，loss 間斷地跳躍，導(dǎo)致之后的不穩(wěn)定

一旦我們解決了這些問題，我們就會上傳變化到?trl?上，以保證社區(qū)受益。

總結(jié)

在本博客，我們走過了 RLHF 訓(xùn)練的整個流程，從準(zhǔn)備人類標(biāo)注的數(shù)據(jù)集開始，調(diào)整語言模型到特定領(lǐng)域，訓(xùn)練獎勵模型，并最終用 RL 訓(xùn)練一個模型。

通過使用?peft，任何人都能在一張 GPU 上跑我們的實驗 ! 如果訓(xùn)練慢了，可以用數(shù)據(jù)并行化的方法，不需要改任何代碼，或者用多張 GPU 并行提高訓(xùn)練速度。

對實際應(yīng)用，這僅僅是第一步 ! 一旦你有了模型，你就要和其他模型比較優(yōu)劣。這個可以用一個面向不同模型的排名生成做到，和我們訓(xùn)練獎勵數(shù)據(jù)集類似。

一旦你加入了評估的步驟，好玩的就開始了: 你可以在原數(shù)據(jù)集上反復(fù)煉丹，也可以增加數(shù)據(jù)集或者對原數(shù)據(jù)集提純。另外，你可以對獎勵模型和生成試不同大小和結(jié)構(gòu)的模型，這需要時間。

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

引用

感謝

我們感謝 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 (阿東)