DeepSpeed Chat: 一鍵式RLHF訓練，讓你的類ChatGPT千億大模型提速省錢15倍
1. 概述
近日來，ChatGPT及類似模型引發了人工智慧（AI）領域的一場風潮。 這場風潮對數字世界產生了革命性影響。ChatGPT類模型具有驚人的泛用性，能夠執行歸納、程式設計、翻譯等任務，其結果與人類專家相當甚至更優。為了使ChatGPT等模型的訓練和部署更輕鬆，AI 開源社群進行了各種嘗試（例如 ChatLLaMa、Alpaca、Vicuna、Databricks-Dolly等）。
然而，儘管開源社群付出了巨大的努力，目前仍缺乏一個支援端到端的基於人工反饋機制的強化學習（RLHF）的規模化系統，這使得訓練強大的類ChatGPT模型十分困難。例如，使用現有的開源系統訓練一個具有 67 億引數的類ChatGPT模型通常需要昂貴的多卡至多節點的 GPU 叢集，但這些資源對大多數資料科學家或研究者而言難以獲取。同時，即使有了這樣的計算資源，現有的開源系統的訓練效率通常還不到這些機器所能達到的最大效率的5%。簡而言之，即使有了昂貴的多GPU叢集，現有解決方案也無法輕鬆、快速、經濟的訓練具有數千億引數的最先進的類ChatGPT模型。
ChatGPT模型的訓練是基於InstructGPT論文中的RLHF方式。這與常見的大語言模型的預訓練和微調截然不同。這使得現有深度學習系統在訓練類ChatGPT模型時存在種種侷限。因此，為了讓ChatGPT型別的模型更容易被普通資料科學家和研究者使用，並使RLHF訓練真正普及到AI社群，我們釋出了 DeepSpeed-Chat。DeepSpeed-Chat具有以下三大核心功能：
（i）簡化 ChatGPT 型別模型的訓練和強化推理體驗：只需一個指令碼即可實現多個訓練步驟，包括使用 Huggingface 預訓練的模型、使用 DeepSpeed-RLHF 系統執行 InstructGPT 訓練的所有三個步驟、甚至生成你自己的類ChatGPT模型。此外，我們還提供了一個易於使用的推理API，用於使用者在模型訓練後測試對話式互動。
（ii）DeepSpeed-RLHF 模組：DeepSpeed-RLHF 復刻了 InstructGPT 論文中的訓練模式，並確保包括a) 監督微調（SFT），b) 獎勵模型微調和 c) 基於人類反饋的強化學習（RLHF）在內的三個步驟與其一一對應。此外，我們還提供了資料抽象和混合功能，以支援使用者使用多個不同來源的資料來源進行訓練。
（iii）DeepSpeed-RLHF 系統：我們將 DeepSpeed 的訓練（training engine）和推理能力（inference engine) 整合到一個統一的混合引擎（DeepSpeed Hybrid Engine or DeepSpeed-HE）中用於 RLHF 訓練。DeepSpeed-HE 能夠在 RLHF 中無縫地在推理和訓練模式之間切換，使其能夠利用來自 DeepSpeed-Inference 的各種最佳化，如張量平行計算和高效能CUDA運算元進行語言生成，同時對訓練部分還能從 ZeRO- 和 LoRA-based 記憶體最佳化策略中受益。DeepSpeed-HE 還能夠自動在 RLHF 的不同階段進行智慧的記憶體管理和資料快取。
DeepSpeed-RLHF 系統在大規模訓練中具有無與倫比的效率，使複雜的 RLHF 訓練變得快速、經濟並且易於大規模推廣：
高效性和經濟性：DeepSpeed-HE 比現有系統快 15 倍以上，使 RLHF 訓練快速且經濟實惠。例如，DeepSpeed-HE 在 Azure 雲上只需 9 小時即可訓練一個 OPT-13B模型，只需 18 小時即可訓練一個 OPT-30B模型。這兩種訓練分別花費不到 300 美元和 600 美元。
| GPUs | OPT-6.7B | OPT-13B | OPT-30B | OPT-66B |
|---|---|---|---|---|
| 8x A100-40GB | 5.7 hours | 10.8 hours | 1.85 days | NA |
| 8x A100-80GB | 4.1 hours ($132) | 9 hours ($290) | 18 hours ($580) | 2.1 days ($1620) |
表 1. 單節點 8x A100：訓練時長及預估的 Azure 費用。
卓越的擴充套件性：DeepSpeed-HE 能夠支援訓練擁有數千億引數的模型，並在多節點多 GPU 系統上展現出卓越的擴充套件性。因此，即使是一個擁有 130 億引數的模型，也只需 1.25 小時就能完成訓練。而對於龐大的 擁有1750 億引數的模型，使用 DeepSpeed-HE 進行訓練也只需不到一天的時間。
| GPUs | OPT-13B | OPT-30B | OPT-66B | OPT-175B |
|---|---|---|---|---|
| 64x A100-80G | 1.25 hours ($320) | 4 hours ($1024) | 7.5 hours ($1920) | 20 hours ($5120) |
表 2. 多節點 64x A100-80GB：訓練時長及預估的 Azure 費用。
非常重要的細節: 上述兩個表格（即表一和表二）中的資料均針對 RLHF 訓練的第 3 步，基於實際資料集和 DeepSpeed-RLHF 訓練吞吐量的測試。該訓練在總共 1.35 億（135M）個字元（token）上進行一個時期（epoch）的訓練。我們總共有 6750 萬個查詢（query）字元（131.9k 個 query，每個序列長度為 256）和 6750 萬個生成/回答字元（131.9k 個答案，每個序列長度為 256），每步的最大全域性字元批次大小約為 500 萬個字元（1024 個查詢-答案對）。在與 DeepSpeed-RLHF 進行任何成本和端到端時間比較之前，我們建議讀者注意這些設定。想要了解更多詳細資訊，請參閱我們的頁面 benchmark setting。
實現 RLHF 訓練的普及化：僅憑單個 GPU，DeepSpeed-HE 就能支援訓練超過 130 億引數的模型。這使得那些無法使用多 GPU 系統的資料科學家和研究者不僅能夠輕鬆建立輕量級的 RLHF 模型，還能建立大型且功能強大的模型，以應對不同的使用場景。
| V100 32G | A6000 48G | A100 40G | A100 80G | |
|---|---|---|---|---|
| Model Size | OPT-2.7B | OPT-6.7B | OPT-6.7B | OPT-13B |
表 3. DeepSpeed-HE 在不同的 GPU 單卡上支援的最大模型。
接下來，我們將深入瞭解上面介紹的 DeepSpeed-Chat 的三個功能。
2. 簡潔高效且經濟的 ChatGPT 訓練與推理體驗
接下來，我們向你展示如何使用一行命令即可利用 DeepSpeed-RLHF 系統訓練 OPT-13B 和 OPT-66B 模型。即使時間有限，你也可以在大約兩小時內在單個消費級 GPU 上訓練一個 OPT-1.3B 模型。此外，我們還演示瞭如何利用我們的 DeepSpeed-chat RLHF API 來開發你自己的模型：
使用 DeepSpeed-Chat 的 RLHF 示例輕鬆訓練你的第一個 類ChatGPT 模型
a) 僅需一個指令碼，完成 RLHF 訓練的全部三個階段，生成你的第一個 ChatGPT 模型！
以下示例展示瞭如何使用一個指令碼，以預訓練的 OPT-13B 作為 actor 模型，OPT-350M 作為 reward 模型，生成一個最終的 130 億引數的 ChatGPT 型別的模型：
pip install deepspeed>=0.9.0
git clone https://github.com/microsoft/DeepSpeedExamples.git
cd DeepSpeedExamples/applications/DeepSpeed-Chat/
pip install -r requirements.txt
python train.py --actor-model facebook/opt-13b --reward-model facebook/opt-350m --deployment-type single_node
大約半天時間，你的 擁有130 億引數的模型即可完成訓練。下表展示了整個過程在三個步驟中所需的時間：
| Model Sizes | Step 1 | Step 2 | Step 3 | Total |
|---|---|---|---|---|
| Actor: OPT-13B, Reward: OPT-350M | 2.5hr | 0.25hr | 10.8hr | 13.6hr |
表 4. 在單個DGX節點上（8x A100-40G），針對不同的RLHF步驟， 使用DeepSpeed-Chat訓練OPT-13b所需的時間。
b) 使用 DeepSpeed-Chat 的推理 API 測試訓練好的模型
將訓練好的模型匯入提供的簡易 API，你不僅可以與模型進行問答互動，還可以進行對話式交流。
Human: Do you know Microsoft?
Assistant: Yes, I’m a big fan of Microsoft. What about them?
Human: Can you explian it to a 6-year old child? I wonder how I should describe it
Assistant: Sure, I can try. Microsoft is a company that makes computers, and they make a program called “Windows” which
is the operating system that runs on the computer. It’s like the “brain” of the computer. It’s where all the
programs and files are stored. It’s also where you’ll find the “start” button, which is where you’ll start to
use the computer.
想嘗試不同的模型大小和配置？輕鬆實現！
我們瞭解到使用者通常喜歡嘗試不同的模型大小和配置，以滿足他們不同的訓練時間、資源和質量的需求。藉助 DeepSpeed-Chat，你可以輕鬆實現這些目標。例如，如果你想在 GPU 叢集上訓練一個更大、更高質量的模型，用於你的研究或業務，你可以使用相同的指令碼，只需輸入你期望的模型大小（例如 660 億引數）和 GPU 數量（例如 64 個 GPU）：
python train.py --actor-model facebook/opt-66b --reward-model facebook/opt-350m --deployment-type multi_node
在 接下來的9 小時內，你將擁有一個 660 億引數的 ChatGPT 模型，並可以在你喜歡的前端 GUI 中使用：
| Model Sizes | Step 1 | Step 2 | Step 3 | Total |
|---|---|---|---|---|
| Actor: OPT-66B, Reward: OPT-350M | 82 mins | 5 mins | 7.5hr | 9hr |
表 5. 在8個DGX節點上（64x A100-80G），針對不同的RLHF步驟， 使用DeepSpeed-Chat訓練OPT-66b所需的時間。
如果你只有大約 1-2 小時的咖啡或午餐休息時間，你也可以嘗試使用 DeepSpeed-Chat 訓練一個小型模型。例如，我們為單個資料集準備了一個 13 億引數模型的訓練示例，以便在你的消費級 GPU 上測試我們的框架。當你從午餐休息回來時，你的模型已經準備好供你使用！
python train.py --actor-model facebook/opt-1.3b --reward-model facebook/opt-350m --deployment-type single_gpu
| Model Sizes | Step 1 | Step 2 | Step 3 | Total |
|---|---|---|---|---|
| Actor: OPT-1.3B, Reward: OPT-350M | 2900 secs | 670 secs | 1.2hr | 2.2hr |
表 6. 在單個消費級A6000-48G上，針對不同的RLHF步驟， 使用DeepSpeed-Chat訓練OPT-1.3b所需的時間。
利用 DeepSpeed-Chat 的 RLHF API 自定義你自己的 RLHF 訓練流程
DeepSpeed-Chat 允許使用者使用我們靈活的 API（如下所示）構建自己的 RLHF 訓練流程，使用者可以使用這些 API 重建自己的 RLHF 訓練策略。我們希望這些功能可以為研究探索中建立各種 RLHF 演算法提供通用介面和後端。
engine = DeepSpeedRLHFEngine(
actor_model_name_or_path=args.actor_model_name_or_path,
critic_model_name_or_path=args.critic_model_name_or_path,
tokenizer=tokenizer,
num_total_iters=num_total_iters,
args=args)
trainer = DeepSpeedPPOTrainer(engine=engine, args=args)
for prompt_batch in prompt_train_dataloader:
out = trainer.generate_experience(prompt_batch)
actor_loss, critic_loss = trainer.train_rlhf(out)
3. 完整的 RLHF 訓練流程概述
為了實現無縫的訓練體驗，我們遵循 InstructGPT 論文的方法，並在 DeepSpeed-Chat 中整合了一個端到端的訓練流程，如圖 1 所示。
圖 1: DeepSpeed-Chat 的 RLHF 訓練流程圖示，包含了一些可選擇的功能。
我們的流程包括三個主要步驟：
- 步驟1：監督微調（SFT） —— 使用精選的人類回答來微調預訓練的語言模型以應對各種查詢；
- 步驟2：獎勵模型微調 —— 使用一個包含人類對同一查詢的多個答案打分的資料集來訓練一個獨立的（通常比 SFT 小的）獎勵模型（RW）；
- 步驟3：RLHF 訓練 —— 利用 Proximal Policy Optimization（PPO）演算法，根據 RW 模型的獎勵反饋進一步微調 SFT 模型。
在步驟3中，我們提供了兩個額外的功能，以幫助提高模型質量：
- 指數移動平均（EMA） —— 可以選擇基於 EMA 的檢查點進行最終評估
- 混合訓練 —— 將預訓練目標（即下一個單詞預測）與 PPO 目標混合，以防止在像 SQuAD2.0 這樣的公開基準測試中的效能損失
這兩個訓練功能，EMA 和混合訓練，常常被其他的開源框架所忽略，因為它們並不會妨礙訓練的進行。然而，根據 InstructGPT，EMA 通常比傳統的最終訓練模型提供更好的響應質量，而混合訓練可以幫助模型保持預訓練基準解決能力。因此，我們為使用者提供這些功能，以便充分獲得 InstructGPT 中描述的訓練體驗，並爭取更高的模型質量。
除了與 InstructGPT 論文高度一致外，我們還提供了一項方便的功能，以支援研究人員和從業者使用多個資料資源訓練他們自己的 RLHF 模型：
- 資料抽象和混合能力： DeepSpeed-Chat 能夠使用多個不同來源的資料集訓練模型以獲得更好的模型質量。它配備了（1）一個抽象資料集層，以統一不同資料集的格式；以及（2）資料拆分/混合功能，以便多個資料集在 3 個訓練階段中被適當地混合然後拆分。
在我們之前的章節中，你可以看到使用整個 DeepSpeed-Chat 訓練模型在多輪對話中的表現。
4. DeepSpeed Hybrid Engine —— 統一的高效混合引擎，為 RLHF 訓練提供動力並進行最佳化
DeepSpeed-Chat流程的前兩步與大型模型的常規微調相似，得益於基於ZeRO的記憶體管理最佳化和DeepSpeed訓練中的並行策略靈活組合，實現了規模和速度的提升。然而，流程的第三步在效能方面是最具挑戰性的部分。每次迭代都需要高效處理兩個階段：a) 生成回答的推理階段，為訓練提供輸入；b) 更新 actor 和 reward 模型權重的訓練階段，以及它們之間的互動和排程。這引入了兩個主要困難：（1）記憶體成本，因為在第三階段的整個過程中需要執行多個SFT和RW模型；（2）生成回答階段的速度較慢，如果沒有正確加速，將顯著拖慢整個第三階段。此外，我們在第三階段中新增的兩個重要可選功能，包括指數移動平均（EMA）收集和混合訓練，將產生額外的記憶體和訓練成本。
為了應對這些挑戰，我們將DeepSpeed訓練和推理的系統功能整合為一個統一的基礎設施，稱為混合引擎（Hybrid Engine）。它利用原始DeepSpeed引擎進行高速訓練模式，同時輕鬆應用DeepSpeed推理引擎進行生成/評估模式，為第三階段的RLHF訓練提供了一個明顯更快的訓練系統。如圖2所示，DeepSpeed訓練和推理引擎之間的過渡是無縫的：透過為actor模型啟用典型的eval和train模式，當執行推理和訓練流程時，DeepSpeed選擇其不同的最佳化來執行模型更快並提高整個系統吞吐量。
Figure 2. 設計圖解：DeepSpeed Hybrid Engine，用於加速 RLHF 流程中最耗時的部分。
在RLHF訓練的經驗生成階段的推理執行過程中，DeepSpeed混合引擎使用輕量級記憶體管理系統來處理KV快取和中間結果，同時使用高度最佳化的推理CUDA核和張量平行計算。與現有解決方案相比，DeepSpeed-HE顯著提高了吞吐量（每秒token數）。
在訓練執行過程中，混合引擎使用了多種記憶體最佳化技術，如DeepSpeed的ZeRO系列技術和現在流行的LoRA方法。這些技術在混合引擎中可以彼此相容，並可以組合在一起以提供最高訓練效率。
DeepSpeed-HE可以在訓練和推理之間無縫更改模型分割槽，以支援基於張量平行計算的推理和基於ZeRO的分片機制進行訓練。它還會重新配置記憶體系統以在此期間最大化記憶體可用性。DeepSpeed-HE還透過規避記憶體分配瓶頸和支援大批次大小來進一步提高效能。混合引擎整合了DeepSpeed訓練和推理的一系列系統技術，突破了現有RLHF訓練的極限，併為RLHF工作負載提供了無與倫比的規模和系統效率。
5. DeepSpeed RLHF: 透過 Hybrid Engine 實現無與倫比的規模和效率
回顧
如前所述，DeepSpeed-HE 是一個將強大的用於推理和訓練的結合系統，旨在使 DeepSpeed-RLHF 在各種硬體上實現卓越的規模和效率，使 RLHF 訓練快速、經濟並且易於 AI 社群使用。
在效率和經濟性方面，如表 1 所示，DeepSpeed-HE 在 Azure 雲上只需 9 小時即可訓練一個OPT-13B模型，只需 18 小時既可訓練 OPT-30B模型，分別花費不到 300 美元和 600 美元。在速度和可擴充套件性方面，如表 2 所示，即使是 13B 的模型也可以在 1.25 小時內訓練，而龐大的 175B 模型可以在不到一天的時間內使用 64 個 GPU 叢集進行訓練。在 RLHF 的可訪問性和普及化方面，DeepSpeed-HE 可以在單個 GPU 上訓練超過 130 億引數的模型，如表 3 所示。
與現有 RLHF 系統的吞吐量和模型大小可擴充套件性比較
與其他 RLHF 系統（如 Colossal-AI 或由原生 PyTorch 提供支援的 HuggingFace）相比，DeepSpeed-RLHF 在系統效能和模型可擴充套件性方面表現出色：
- 就吞吐量而言，DeepSpeed 在單個 GPU 上的 RLHF 訓練中實現了 10 倍以上的改進（圖 3）。在多 GPU 設定中，它比 Colossal-AI 快 6 - 19 倍，比 HuggingFace DDP 快 1.4 - 10.5 倍（圖 4）。
- 就模型可擴充套件性而言，Colossal-AI 可以在單個 GPU 上執行最大 1.3B 的模型，在單個 A100 40G 節點上執行 6.7B 的模型，而 DeepSpeed-HE 可以在相同的硬體上分別執行 6.5B 和 50B 的模型，實現高達 7.5 倍的提升。
因此，憑藉超過一個數量級的更高吞吐量，與現有的 RLHF 系統（如 Colossal-AI 或 HuggingFace DDP）相比，DeepSpeed-HE 擁有在相同時間預算下訓練更大的 actor 模型的能力，或者以十分之一的成本訓練類似大小的模型的能力。
圖 3. 在單個 NVIDIA A100-40G GPU 上，將 RLHF 訓練的吞吐量與另外兩個系統框架在步驟 3 進行比較。沒有圖示表示 OOM（記憶體不足）的情況
圖 4. 在單個 DGX 節點上，使用 8 個 NVIDIA A100-40G GPU，對訓練流程第 3 步（耗時最長的部分）的不同模型大小進行端到端訓練吞吐量比較。沒有圖示表示 OOM（記憶體不足）的情況。
這種效率的提高是 DeepSpeed-HE 利用 DeepSpeed 推理最佳化在 RLHF 處理過程中加速 RLHF 生成的結果。圖 5 顯示了 RLHF 訓練迭代中 1.3B 引數模型的時間消耗細節：大部分時間用於生成階段。透過利用 DeepSpeed 的高效能推理核心，DeepSpeed-HE 在這個階段可以實現比 HuggingFace 高達 9 倍的吞吐量改進，比 Colossal-AI 高 15 倍，從而實現無與倫比的端到端效率。
圖 5. DeepSpeed Chat 的混合引擎在生成階段的優越加速：在單個 DGX 節點上使用 8 個 A100-40G GPU 訓練 OPT-1.3B actor 模型 + OPT-350M reward 模型的時間/序列分解。
有效吞吐量和可擴充套件性分析
(I) 有效吞吐量分析。 在 RLHF 訓練的第 3 階段，DeepSpeed-HE 的有效吞吐量取決於它在生成和 RL 訓練階段所實現的吞吐量。在我們的 RLHF （詳見 benchmarking setting）中，生成階段佔總計算的約 20%，而 RL 訓練階段佔剩餘的 80%。然而，儘管比例較小，前者可能會佔用大部分的端到端時間，因為它需要為每個生成的字元執行一次 actor 模型，使其受到記憶體頻寬限制，難以實現高吞吐量。相比之下，RL 訓練階段是計算密集型的，僅需執行參考 actor 模型進行幾次前向和後向傳遞，每個樣本都有來自提示和生成的全部 512 個字元，可以實現良好的吞吐量。
圖 6. 在最大效率的情況下，DeepSpeed-HE 針對不同模型大小的RLHF生成、訓練和有效吞吐量。
為了最大化有效吞吐量，DeepSpeed-HE 對兩個階段進行了最佳化。首先，它使用盡可能大的批次大小以在兩個階段上獲得更高的效率。其次，在生成階段，它利用高效能CUDA核心在模型在單個 GPU 上最大化 GPU 記憶體頻寬利用率，並在其他情況下利用張量並行（Tensor Parallelism, 簡寫作TP）進行計算。DeepSpeed-HE進一步在生成階段使用 TP 而不是 ZeRO 以減少 GPU 之間的通訊並保持高 GPU 記憶體頻寬利用率。
圖 6 顯示了 DeepSpeed-HE 在 1.3B 到 175B 的模型大小範圍內可以實現的最佳有效吞吐量（以 TFlops/GPU 表示）。它還分別顯示了在生成和訓練階段實現的吞吐量。DeepSpeed-HE 對 6.7B-66B 範圍內的模型最為高效。超出這個範圍到 175B 時，由於記憶體有限，無法支援更大的批次大小，吞吐量下降，但仍比小型 1.3B 模型的效率高 1.2 倍。當我們將這些巨大的模型擴充套件到更多具有更多記憶體的 GPU 時，這些模型的每個 GPU 吞吐量可能會進一步提高。
此外，我們想指出，如圖 2 所示，我們系統的有效效能比現有系統高 19 倍，這表明它們的執行速度低於峰值的 5%。這說明了最佳化 RLHF 工作負載的挑戰以及我們的系統在面對挑戰時的有效性。
圖 7. 在不同數量的DGX (A100-40/80G GPU) 節點上，進行13B（左）和66B（右）actor 模型 和 350M reward 模型的可擴充套件性訓練。
(II) 可擴充套件性分析。 不同模型大小的最佳有效吞吐量取決於不同的 GPU 數量。部分原因是因為一些較大的模型大小需要更多的記憶體來執行。基於此，我們接下來討論 DeepSpeed-HE 的可擴充套件性特性。
圖 7 顯示 DeepSeed-RLHF 在多達 64 個 GPU的叢集 上實現了良好的整體擴充套件。然而，如果我們仔細觀察，可以發現 DeepSpeed-RLHF 訓練在小規模時實現了超線性擴充套件，隨後在較大規模時實現了接近線性或次線性擴充套件。這是由於記憶體可用性和最大全域性批次大小之間的相互作用。
DeepSpeed-HE 的核心技術基於 ZeRO，用於訓練過程中將模型狀態分割到每個GPU上。這意味著隨著 GPU 數量的增加，每個 GPU 的記憶體消耗會減少，使得 DeepSpeed-HE 能夠在每個 GPU 上支援更大的批次，從而實現超線性擴充套件。然而，在大規模情況下，儘管可用記憶體持續增加，但最大全域性批次大小仍然限制了每個 GPU 的批次大小，導致接近線性或次線性擴充套件。因此，在給定的最大全域性批次大小（例如，我們設定為 1024 個句子，每個句子長度為 512）下，DeepSpeed-HE 在超線性和次線性可擴充套件性之間實現了最佳的吞吐量和成本效益。具體的平衡點主要取決於每個 GPU 上可執行的最大批次大小，而這又受到可用記憶體和全域性批次大小的函式所決定。