在 24GB 消費級 GPU 上用 RLHF 微調 200 億引數大語言模型

我們很高興正式釋出 trl 與 peft 的整合，讓任何人都能更輕鬆地使用強化學習微調大語言模型 (LLM)！在這篇文章中，我們將解釋為什麼這是一種能與現有微調方法相媲美的替代方案。

注意，peft 是一個可以應用於許多機器學習用例的通用工具，但它對於 RLHF 特別有意思，因為這種方法特別消耗記憶體！

如果你想直接深入程式碼，請直接檢視 TRL 文件頁面上的示例指令碼。

引言

大語言模型與 RLHF

大語言模型 (LLM) 與人類反饋強化學習 (RLHF) 相結合，似乎是構建如 ChatGPT 這般強大 AI 系統的下一個首選方法。

使用 RLHF 訓練語言模型通常包括以下三個步驟

1- 在特定領域或指令語料庫以及人類演示上微調預訓練的 LLM

2- 收集一個由人類標註的資料集並訓練一個獎勵模型

3- 使用獎勵模型和這個資料集，透過強化學習（例如 PPO）進一步微調步驟 1 中的 LLM

基礎 LLM 的選擇在這裡至關重要。在撰寫本文時，可以“開箱即用”於許多工的“最佳”開源 LLM 是經過指令微調的 LLM。著名的模型有：BLOOMZ、Flan-T5、Flan-UL2，以及 OPT-IML。這些模型的缺點是它們的尺寸。要獲得一個不錯的模型，至少需要使用 100 億以上引數規模的模型，這在全精度下需要高達 40GB 的 GPU 視訊記憶體，僅僅是為了在單個 GPU 裝置上載入模型，還完全沒有進行任何訓練！

什麼是 TRL？

trl 庫旨在讓強化學習步驟變得更簡單、更靈活，以便任何人都能在自己的自定義資料集和訓練設定上使用強化學習微調他們的語言模型。在眾多應用中，你可以使用這個演算法來微調模型以生成正面的電影評論、進行可控生成或使模型毒性降低。

使用 trl，你可以在分散式環境或單個裝置上執行最流行的深度強化學習演算法之一，PPO！我們利用 Hugging Face 生態系統中的 accelerate 庫來實現這一點，以便任何使用者都可以將實驗擴充套件到相當可觀的規模。

使用強化學習微調語言模型大致遵循下面詳述的協議。這需要原始模型的 2 個副本；為避免活動模型與其原始行為/分佈偏離太遠，你需要在每個最佳化步驟計算參考模型的 logits。這給最佳化過程增加了一個硬性約束，因為每個 GPU 裝置上至少需要兩個模型副本。如果模型尺寸增大，將整個設定裝入單個 GPU 變得越來越棘手。

在 trl 中，你還可以在參考模型和活動模型之間使用共享層，以避免完整的副本。這個功能的一個具體例子在去毒化示例中有所展示。

大規模訓練

大規模訓練可能具有挑戰性。第一個挑戰是將模型及其最佳化器狀態裝入可用的 GPU 裝置。單個引數佔用的 GPU 視訊記憶體量取決於其“精度”（或更具體地說是 dtype）。最常見的 dtype 是 float32（32-bit）、float16 和 bfloat16（16-bit）。最近，“奇異”的精度也得到了開箱即用的支援，用於訓練和推理（在某些條件和約束下），例如 int8（8-bit）。簡而言之，要在 GPU 裝置上載入模型，每十億個引數在 float32 精度下需要 4GB，在 float16 精度下需要 2GB，在 int8 精度下需要 1GB。如果你想了解更多關於這個主題的資訊，可以看看這篇深入探討的博文：https://huggingface.co/blog/hf-bitsandbytes-integration。

如果你使用 AdamW 最佳化器，每個引數需要 8 位元組（例如，如果你的模型有 10 億個引數，模型的完整 AdamW 最佳化器將需要 8GB GPU 視訊記憶體 - 來源）。

已經採用了許多技術來應對這些大規模挑戰。最熟悉的正規化是流水線並行、張量並行和資料並行。

在資料並行中，相同的模型並行託管在多臺機器上，每個例項被饋送不同的資料批次。這是最直接的並行策略，本質上是複製了單 GPU 的情況，並且已經被 trl 支援。在流水線和張量並行中，模型本身被分佈在多臺機器上：在流水線並行中，模型是按層分割的，而張量並行則將張量操作（例如矩陣乘法）跨 GPU 分割。透過這些模型並行策略，你需要在多個裝置之間分片模型權重，這要求你定義一個跨程序的啟用和梯度的通訊協議。這實現起來並不簡單，可能需要採用一些框架，如 Megatron-DeepSpeed 或 Nemo。同樣重要的是要強調其他對於擴充套件 LLM 訓練至關重要的工具，如自適應啟用檢查點和融合核心。關於並行正規化的進一步閱讀可以在這裡找到。

因此，我們問了自己以下問題：僅使用資料並行，我們能走多遠？我們能否使用現有工具將超大規模的訓練過程（包括活動模型、參考模型和最佳化器狀態）裝入單個裝置？答案似乎是肯定的。主要要素是：介面卡和 8bit 矩陣乘法！讓我們在以下章節中介紹這些主題。

8-bit 矩陣乘法

高效的 8-bit 矩陣乘法是首次在論文 LLM.int8() 中引入的一種方法，旨在解決量化大規模模型時的效能下降問題。所提出的方法將線性層中在底層應用的矩陣乘法分解為兩個階段：將在 float16 中執行的離群隱藏狀態部分，以及在 int8 中執行的“非離群”部分。

簡而言之，如果使用 8-bit 矩陣乘法，可以將全精度模型的尺寸減少 4 倍（因此，對於半精度模型則減少 2 倍）。

低秩自適應與 PEFT

2021年，一篇名為《LoRA: 大語言模型的低秩自適應》的論文表明，可以透過凍結預訓練權重並建立注意力矩陣中查詢和值層的低秩版本來對大語言模型進行微調。這些低秩矩陣的引數比原始模型少得多，從而可以用少得多的 GPU 視訊記憶體進行微調。作者們證明，微調低秩介面卡取得了與微調整個預訓練模型相當的結果。

這項技術使得使用一小部分記憶體需求就能對 LLM 進行微調。然而，也有一些缺點。由於介面卡層中額外的矩陣乘法，前向和後向傳播大約慢兩倍。

什麼是 PEFT？

引數高效微調（PEFT）是一個 Hugging Face 庫，旨在支援在 LLM 上建立和微調適配器層。peft 與 🤗 Accelerate 無縫整合，利用 DeepSpeed 和大模型推理技術處理大規模模型。

該庫支援許多最先進的模型，並有一系列廣泛的示例，包括

• 因果語言建模
• 條件生成
• 影像分類
• 8-bit int8 訓練
• Dreambooth 模型的低秩自適應
• 語義分割
• 序列分類
• 詞元分類

該庫仍在廣泛和積極的開發中，未來幾個月將宣佈許多新功能。

使用低秩介面卡微調 200 億引數模型

現在，先決條件都已介紹完畢，讓我們一步一步地過一遍整個流程，並透過圖表解釋如何使用上面提到的工具，在單個 24GB GPU 上用強化學習微調一個 200 億引數的 LLM！

步驟 1：以 8-bit 精度載入你的活動模型

使用 transformers 實現 LLM 記憶體減少的“免費午餐”方法是，使用 LLM.int8 中描述的方法以 8-bit 精度載入你的模型。這可以透過在呼叫 from_pretrained 方法時簡單地新增標誌 load_in_8bit=True 來實現（你可以在這裡閱讀更多相關資訊）。

如前一節所述，一個計算載入模型所需 GPU 視訊記憶體量的“技巧”是以“十億引數”為單位來思考。由於一個位元組需要 8 位元，因此全精度模型（32bit = 4bytes）每十億引數需要 4GB，半精度模型每十億引數需要 2GB，而 int8 模型每十億引數需要 1GB。

所以首先，我們只加載 8-bit 的活動模型。讓我們看看第二步需要做什麼！

步驟 2：使用 peft 新增額外的可訓練介面卡

第二步是在模型內部載入介面卡並使這些介面卡可訓練。這使得活動模型所需的可訓練權重數量大幅減少。這一步利用了 peft 庫，可以用幾行程式碼完成。請注意，一旦介面卡訓練完成，你可以輕鬆地將它們推送到 Hub 以供日後使用。

步驟 3：使用同一模型獲取參考和活動 logits

由於介面卡可以被停用，我們可以使用同一個模型來獲取 PPO 的參考和活動 logits，而無需建立同一個模型的兩個副本！這利用了 peft 庫中的一個功能，即 disable_adapters 上下文管理器。

訓練指令碼概述：

我們現在將描述我們如何使用 transformers、peft 和 trl 訓練一個 200 億引數的 gpt-neox 模型。這個例子的最終目標是在記憶體受限的情況下微調一個 LLM 來生成正面的電影評論。類似的步驟可以應用於其他任務，如對話模型。

總共有三個關鍵步驟和訓練指令碼

1. 指令碼 - 在凍結的 8-bit 模型上微調一個低秩介面卡，用於在 imdb 資料集上進行文字生成。
2. 指令碼 - 將介面卡層合併到基礎模型的權重中，並將它們儲存在 Hub 上。
3. 指令碼 - 對低秩介面卡進行情感微調，以建立正面評論。

我們在 24GB NVIDIA 4090 GPU 上測試了這些步驟。雖然在 24GB GPU 上可以完成整個訓練執行，但完整的訓練執行是在 🤗 研究叢集上的單個 A100 上進行的。

訓練過程的第一步是在預訓練模型上進行微調。通常這需要多個高階的 80GB A100 GPU，所以我們選擇訓練一個低秩介面卡。我們將其視為因果語言建模設定，並在 imdb 資料集的示例上訓練一個 epoch，該資料集包含電影評論和表明它們是正面還是負面情緒的標籤。

為了使用適配過的模型並用強化學習進行進一步微調，我們首先需要合併適配的權重，這是透過以 16-bit 浮點數載入預訓練模型和介面卡，並與權重矩陣相加（應用了適當的縮放）來實現的。

最後，我們可以在凍結的、經過 imdb 微調的模型之上，再微調另一個低秩介面卡。我們使用一個 imdb 情感分類器來為 RL 演算法提供獎勵。

此實驗的完整 Weights and Biases 報告可在此處獲取：這裡，如果你想檢視更多圖表和文字生成結果。

結論

我們在 trl 中實現了一項新功能，允許使用者透過利用 peft 和 bitsandbytes 庫，以合理的成本使用 RLHF 微調大語言模型。我們證明了在 24GB 消費級 GPU 上微調 gpt-neo-x（在 bfloat16 下為 40GB！）是可能的，我們期待社群能廣泛使用這一整合來微調更大的模型，利用 RLHF 並分享優秀的成果。

我們已經確定了下一步推動該整合極限的一些有趣方向

• 這在多 GPU 設定中將如何擴充套件？我們主要將探索這種整合在 GPU 數量方面的擴充套件情況，是否可以開箱即用地應用資料並行，或者是否需要在任何相關庫中採用新功能。
• 我們可以利用哪些工具來提高訓練速度？我們觀察到，這種整合的主要缺點是整體訓練速度。未來我們希望能探索可能的方向，使訓練速度更快。

參考文獻

• 並行正規化： https://huggingface.co/docs/transformers/v4.17.0/en/parallelism
• transformers 中的 8-bit 整合： https://huggingface.co/blog/hf-bitsandbytes-integration
• LLM.int8 論文： https://arxiv.org/abs/2208.07339
• 梯度檢查點解釋： https://docs.aws.amazon.com/sagemaker/latest/dg/model-parallel-extended-features-pytorch-activation-checkpointing.html