世界今日訊！微軟DeepSpeed Chat，人人可快速訓練百億、千億級ChatGPT大模型

機器之心專欄

微軟DeepSpeed組

【資料圖】

昨日，微軟開源的 DeepSpeed Chat 引起了AI社區的廣泛關注。

它讓我們能夠以更低的成本、更快的速度訓練類似于ChatGPT的高質量大模型。

鏈接：https://github.com/microsoft/DeepSpeed/tree/master/blogs/deepspeed-chat

本文對開源的 DeepSpeed Chat 進行了詳細的介紹。該博客由微軟DeepSpeed組官方撰寫并譯制，并授權「開源社」首發于中文社區，機器之心轉載。

概述

近日來，ChatGPT 及類似模型引發了人工智能（AI）領域的一場風潮。這場風潮對數字世界產生了革命性影響。ChatGPT 類模型具有驚人的泛用性，能夠執行歸納、編程、翻譯等任務，其結果與人類專家相當甚至更優。為了使 ChatGPT 等模型的訓練和部署更輕松，AI 開源社區進行了各種嘗試（例如 ChatLLaMa、Alpaca、Vicuna、Databricks-Dolly 等）。

然而，盡管開源社區付出了巨大的努力，目前仍缺乏一個支持端到端的基于人工反饋機制的強化學習（RLHF）的規?；到y，這使得訓練強大的類 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 美元。

表 1. 單節點 8x A100：訓練時長及預估的 Azure 費用。

卓越的擴展性：DeepSpeed-HE 能夠支持訓練擁有數千億參數的模型，并在多節點多 GPU 系統上展現出卓越的擴展性。因此，即使是一個擁有 130 億參數的模型，也只需 1.25 小時就能完成訓練。而對于龐大的 擁有 1750 億參數的模型，使用 DeepSpeed-HE 進行訓練也只需不到一天的時間。

表 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 模型，還能創建大型且功能強大的模型，以應對不同的使用場景。

表 3. DeepSpeed-HE 在不同的 GPU 單卡上支持的最大模型。

接下來，我們將深入了解上面介紹的 DeepSpeed-Chat 的三個功能。

簡潔高效且經濟的 ChatGPT訓練與推理體驗

接下來，我們向你展示如何使用一行命令即可利用 DeepSpeed-RLHF 系統訓練 OPT-1,3B, 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 類型的模型：

git clone https://github.com/microsoft/DeepSpeed.gitcd DeepSpeedpip install .git clone https://github.com/microsoft/DeepSpeedExamples.gitcd DeepSpeedExamples/applications/DeepSpeed-Chat/pip install -r requirements.txtpython train.py --actor-model facebook/opt-13b --reward-model facebook/opt-350m --num-gpus 8

大約半天時間，你擁有 130 億參數的模型即可完成訓練。下表展示了整個過程在三個步驟中所需的時間：

表 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 explain it to a 6-year old child? I wonder how I should describe itAssistant:    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 --num-gpus 64

在接下來的 9 小時內，你將擁有一個 660 億參數的 ChatGPT 模型，并可以在你喜歡的前端 GUI 中使用：

表 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 --num-gpus 1

表 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)

完整的 RLHF 訓練流程概述

為了實現無縫的訓練體驗，我們遵循 InstructGPT 論文的方法，并在 DeepSpeed-Chat 中整合了一個端到端的訓練流程，如圖 1 所示。

圖 1: DeepSpeed-Chat 的 RLHF 訓練流程圖示，包含了一些可選擇的功能。

我們的流程包括三個主要步驟：

在步驟 3 中，我們提供了兩個額外的功能，以幫助提高模型質量：

這兩個訓練功能，EMA 和混合訓練，常常被其他的開源框架所忽略，因為它們并不會妨礙訓練的進行。然而，根據 InstructGPT，EMA 通常比傳統的最終訓練模型提供更好的響應質量，而混合訓練可以幫助模型保持預訓練基準解決能力。因此，我們為用戶提供這些功能，以便充分獲得 InstructGPT 中描述的訓練體驗，并爭取更高的模型質量。

除了與 InstructGPT 論文高度一致外，我們還提供了一項方便的功能，以支持研究人員和從業者使用多個數據資源訓練他們自己的 RLHF 模型：

在我們之前的章節中，你可以看到使用整個 DeepSpeed-Chat 訓練模型在多輪對話中的表現。

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 工作負載提供了無與倫比的規模和系統效率。

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 在系統性能和模型可擴展性方面表現出色：

因此，憑借超過一個數量級的更高吞吐量，與現有的 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 的內存消耗會減少，使得 DeepSpeed-HE 能夠在每個 GPU 上支持更大的批量，從而實現超線性擴展。然而，在大規模情況下，盡管可用內存持續增加，但最大全局批量大小仍然限制了每個 GPU 的批量大小，導致接近線性或次線性擴展。因此，在給定的最大全局批量大?。ɡ?，我們設置為 1024 個句子，每個句子長度為 512）下，DeepSpeed-HE 在超線性和次線性可擴展性之間實現了最佳的吞吐量和成本效益。具體的平衡點主要取決于每個 GPU 上可運行的最大批量大小，而這又受到可用內存和全局批量大小的函數所決定。

發布：現在就嘗試 DeepSpeed Chat 吧！

我們非常高興地宣布，DeepSpeed-Chat 現已開源并向 AI 社區開放。

DeepSpeed-Chat 是更大的 DeepSpeed 生態系統的一部分，包括眾多深度學習系統和建模技術。要了解更多信息，

DeepSpeed 歡迎你的貢獻！我們鼓勵你在 DeepSpeed GitHub 頁面（https://github.com/microsoft/DeepSpeed/）報告問題、貢獻 PR 并參與討論。請參閱我們的貢獻指南（https://github.com/microsoft/DeepSpeed/blob/master/CONTRIBUTING.md）了解更多詳情。我們愿意與大學、研究實驗室、公司等進行合作，共同開展深度學習研究，將 DeepSpeed 應用于賦能現實世界的 AI 模型和應用等。對于此類需求（以及其他不適合在 GitHub 上提出的需求），請直接發送電子郵件至 deepspeed-info@microsoft.com。

責任編輯：Rex_13