在一張 24 GB 的消費級顯卡上用 RLHF 微調 20B LLMs

語言: CN / TW / HK
時間 2023-03-29 13:01:34

我們很高興正式發佈 trl 與 peft 的集成，使任何人都可以更輕鬆地使用強化學習進行大型語言模型 (LLM) 微調！在這篇文章中，我們解釋了為什麼這是現有微調方法的有競爭力的替代方案。

請注意， peft 是一種通用工具，可以應用於許多 ML 用例，但它對 RLHF 特別有趣，因為這種方法特別需要內存！

如果你想直接深入研究代碼，請直接在 TRL 的文檔頁面 直接查看示例腳本。

介紹

LLMs & RLHF

LLM 結合 RLHF (人類反饋強化學習) 似乎是構建非常強大的 AI 系統 (例如 ChatGPT) 的下一個首選方法。

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

  1. 在特定領域或指令和人類示範語料庫上微調預訓練的 LLM;
  2. 收集人類標註的數據集，訓練一個獎勵模型;
  3. 使用 RL (例如 PPO)，用此數據集和獎勵模型進一步微調步驟 1 中的 LLM。

ChatGPT 的訓練協議概述，從數據收集到 RL 部分。 資料來源: OpenAI 的 ChatGPT 博文

基礎 LLM 的選擇在這裏是至關重要的。在撰寫本文時，可以“開箱即用”地用於許多任務的“最佳”開源 LLM 是指令微調 LLMs。著名的模型有: BLOOMZ Flan-T5、Flan-UL2 和 OPT-IML。這些模型的缺點是它們的尺寸。要獲得一個像樣的模型，你至少需要玩 10B+ 級別的模型，在全精度情況下這將需要高達 40GB GPU 內存，只是為了將模型裝在單個 GPU 設備上而不進行任何訓練！

什麼是 TRL？

trl 庫的目的是使 RL 的步驟更容易和靈活，讓每個人可以在他們自己的數據集和訓練設置上用 RL 微調 LM。在許多其他應用程序中，你可以使用此算法微調模型以生成 正面電影評論、進行 受控生成 或 降低模型的毒性。

使用 trl 你可以在分佈式管理器或者單個設備上運行最受歡迎的深度強化學習算法之一: PPO。我們利用 Hugging Face 生態系統中的 accelerate 來實現這一點，這樣任何用户都可以將實驗擴大到一個有趣的規模。

使用 RL 微調語言模型大致遵循下面詳述的協議。這需要有 2 個原始模型的副本; 為避免活躍模型與其原始行為/分佈偏離太多，你需要在每個優化步驟中計算參考模型的 logits 。這對優化過程增加了硬約束，因為你始終需要每個 GPU 設備至少有兩個模型副本。如果模型的尺寸變大，在單個 GPU 上安裝設置會變得越來越棘手。

TRL 中 PPO 訓練設置概述

在 trl 中，你還可以在參考模型和活躍模型之間使用共享層以避免整個副本。 模型解毒示例中展示了此功能的具體示例。

大規模訓練

大規模訓練是具有挑戰性的。第一個挑戰是在可用的 GPU 設備上擬合模型，及其優化器狀態。 單個參數佔用的 GPU 內存量取決於其“精度”(或更具體地説是 dtype)。 最常見的 dtype 是 float32 (32 位) 、 float16 和 bfloat16 (16 位)。 最近，“奇異的”精度支持開箱即用的訓練和推理 (具有特定條件和約束)，例如 int8 (8 位)。 簡而言之，要在 GPU 設備上加載一個模型，每十億個參數在 float32 精度上需要 4GB，在 float16 上需要 2GB，在 int8 上需要 1GB。 如果你想了解關於這個話題的更多信息，請查看這篇研究深入的文章: http://huggingface.co/blog/hf-bitsandbytes-integration。

如果您使用 AdamW 優化器，每個參數需要 8 個字節 (例如，如果您的模型有 1B 個參數，則模型的完整 AdamW 優化器將需要 8GB GPU 內存 來源)。

許多技術已經被採用以應對大規模訓練上的挑戰。最熟悉的範式是管道並行、張量並行和數據並行。

圖片來自 這篇博文

通過數據並行性，同一模型並行託管在多台機器上，並且每個實例都被提供不同的數據批次。 這是最直接的並行策略，本質上是複製單 GPU 的情況，並且已經被 trl 支持。 使用管道並行和張量並行，模型本身分佈在機器上: 在管道並行中，模型按層拆分，而張量並行則跨 GPU 拆分張量操作 (例如矩陣乘法)。使用這些模型並行策略，你需要將模型權重分片到許多設備上，這需要你定義跨進程的激活和梯度的通信協議。 這實現起來並不簡單，可能需要採用一些框架，例如 Megatron-DeepSpeed 或 Nemo。其他對擴展訓練至關重要的工具也需要被強調，例如自適應激活檢查點和融合內核。 可以在 擴展閲讀 找到有關並行範式的進一步閲讀。

因此，我們問自己下面一個問題: 僅用數據並行我們可以走多遠？我們能否使用現有的工具在單個設備中適應超大型訓練過程 (包括活躍模型、參考模型和優化器狀態)？ 答案似乎是肯定的。 主要因素是: 適配器和 8 位矩陣乘法！ 讓我們在以下部分中介紹這些主題:

8 位矩陣乘法

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

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

簡而言之，如果使用 8 位矩陣乘法，則可以將全精度模型的大小減小到 4 分之一 (因此，對於半精度模型，可以減小 2 分之一)。

低秩適配和 PEFT

在 2021 年，一篇叫 LoRA: Low-Rank Adaption of Large Language Models 的論文表明，可以通過凍結預訓練權重，並創建查詢和值層的注意力矩陣的低秩版本來對大型語言模型進行微調。這些低秩矩陣的參數遠少於原始模型，因此可以使用更少的 GPU 內存進行微調。 作者證明，低階適配器的微調取得了與微調完整預訓練模型相當的結果。

原始 (凍結的) 預訓練權重 (左) 的輸出激活由一個由權重矩陣 A 和 B 組成的低秩適配器 (右) 增強。

這種技術允許使用一小部分內存來微調 LLM。 然而，也有一些缺點。由於適配器層中的額外矩陣乘法，前向和反向傳遞的速度大約是原來的兩倍。

什麼是 PEFT？

Parameter-Efficient Fine-Tuning (PEFT) 是一個 Hugging Face 的庫，它被創造出來以支持在 LLM 上創建和微調適配器層。 peft 與 🤗 Accelerate 無縫集成，用於利用了 DeepSpeed 和 Big Model Inference 的大規模模型。

此庫支持很多先進的模型，並且有大量的例子，包括:

  • 因果語言建模
  • 條件生成
  • 圖像分類
  • 8 位 int8 訓練
  • Dreambooth 模型的低秩適配
  • 語義分割
  • 序列分類
  • 詞符分類

該庫仍在廣泛和積極的開發中，許多即將推出的功能將在未來幾個月內公佈。

使用低質適配器微調 20B 參數量的模型

現在先決條件已經解決，讓我們一步步過一遍整個管道，並用圖説明如何在單個 24GB GPU 上使用上述工具使用 RL 微調 20B 參數量的 LLM！

第 1 步: 在 8 位精度下加載你的活躍模型

與全精度模型相比，以 8 位精度加載模型最多可節省 4 倍的內存

使用 transformers 減少 LLM 內存的“免費午餐”是使用 LLM.int8 中描述的方法，以 8 位精度加載模型。 這可以通過在調用 from_pretrained 方法時簡單地添加標誌 load_in_8bit=True 來執行 (你可以在 文檔中 閲讀更多相關信息)。

如前一節所述，計算加載模型所需的 GPU 內存量的“技巧”是根據“十億個參數量”進行思考。 由於一個字節需要 8 位，因此全精度模型 (32 位 = 4 字節) 每十億個參數需要 4GB，半精度模型每十億個參數需要 2GB，int8 模型每十億個參數需要 1GB。

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

第 2 步: 使用 peft 添加額外可訓練的適配器

您可以輕鬆地在凍結的 8 位模型上添加適配器，從而通過訓練一小部分參數來減少優化器狀態的內存需求

第二步是在模型中加載適配器並使這些適配器可訓練。 這可以大幅減少活躍模型所需的可訓練權重的數量。 此步驟利用 peft 庫，只需幾行代碼即可執行。 請注意，一旦適配器經過訓練，您就可以輕鬆地將它們推送到 Hub 以供以後使用。

第 3 步: 使用同樣的模型得到參考和活躍 logits

你可以方便地使用 peft 關閉和使能適配器

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

訓練腳本概述

我們現在將描述如何使用 transformers 、 peft 和 trl 訓練 20B 參數量的 gpt-neox 模型。 這個例子的最終目標是微調 LLM 以在內存受限的設置中生成積極的電影評論。類似的步驟可以應用於其他任務，例如對話模型。

整體來看有三個關鍵步驟和訓練腳本:

  1. 腳本 1 - 在凍結的 8 位模型上微調低秩適配器，以便在 imdb 數據集上生成文本。

  2. 腳本 2 - 將適配器層合併到基礎模型的權重中並將它們存儲在 Hub 上。

  3. 腳本 3 - 對低等級適配器進行情感微調以創建正面評價。

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

訓練過程的第一步是對預訓練模型進行微調。 通常這需要幾個高端的 80GB A100 GPU，因此我們選擇訓練低階適配器。 我們將其視為一種因果語言建模設置，並針從 imdb 數據集中訓練了一個 epoch 的示例，該數據集具有電影評論和指明積極還是消極情緒的標籤。

在 imdb 數據集上訓練 gpt-neox-20b 模型一個 epoch 期間的訓練損失

為了利用已經適配了的模型並使用 RL 執行進一步微調，我們首先需要組合自適應權重，這是通過加載預訓練模型和 16 位浮點精度的適配器並且加和權重矩陣 (應用適當的縮放比例) 來實現的。

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

RL 微調 peft 適配過的 20B 參數量的模型以生成積極影評時的獎勵均值。

如果您想查看更多圖表和文本生成，可在 鏈接處 獲取此實驗的完整權重和偏差報告。

結論

我們在 trl 中實現了一項新功能，允許用户利用 peft 和 bitsandbytes 庫以合理的成本使用 RLHF 微調大型語言模型。 我們證明了可以在 24GB 消費級 GPU 上微調 gpt-neo-x (以 bfloat16 精度需要 40GB！)，我們期望社區將廣泛使用此集成來微調利用了 RLHF 的大型模型，並分享出色的工件。

我們已經為接下來的步驟確定了一些有趣的方向，以挑戰這種集成的極限:

  • 這將如何在多 GPU 設置中擴展？ 我們將主要探索這種集成將如何根據 GPU 的數量進行擴展，是否可以開箱即用地應用數據並行，或者是否需要在任何相關庫上採用一些新功能。
  • 我們可以利用哪些工具來提高訓練速度？ 我們觀察到這種集成的主要缺點是整體訓練速度。 在未來，我們將持續探索使訓練更快的可能方向。

參考

英文原文: http://hf.co/blog/trl-peft

作者: Edward Beeching、Younes Belkada、Leandro von Werra、Sourab Mangrulkar、Lewis Tunstall、Kashif Rasul

譯者: AIboy1993 (李旭東)

審校、排版: zhongdongy (阿東)

「」
「」