ChatGPT是OpenAI開發(fā)的基于大型語言模型(LLM)的聊天機器人，以其令人驚嘆的對話能力而迅速火爆并被廣泛采用。ChatGPT 成功背后得益于大型語言模型生成領域的新訓練范式：RLHF (Reinforcement Learning from Human Feedback)，即以強化學習方式依據人類反饋優(yōu)化語言模型。不同于傳統(tǒng)的深度學習訓練，只涉及到單個模型的迭代和優(yōu)化，以 RLHF 為代表的Alignment訓練范式涉及到多個大模型的計算和數據交互，這對于構建構建一個易用、高效、可擴展的訓練系統(tǒng)提出了諸多的挑戰(zhàn)。
為了解決上述問題，我們提出了一個新的 Alignment 訓練框架 PAI-ChatLearn。ChatLearn 通過對 Alignment 訓練流程進行合理的抽象和解耦，提供靈活的資源分配和模型調度策略。ChatLearn提供了RLHF、DPO、OnlineDPO、GRPO等 Alignment 訓練，同時也支持用戶自定義 model 的執(zhí)行 flow，來實現(xiàn)自定義的訓練流程。相比于當前的 SOTA 系統(tǒng)，ChatLearn 在 7B+7B (Policy+Reward) 規(guī)模性能提升115%，70B+70B 規(guī)模性能提升208%，并且能支持更大規(guī)模的Alignment訓練，例如300B+300B規(guī)模。同時ChatLearn也一直在支持Qwen大模型的Alignment訓練，在Qwen-Chat、Qwen2-Chat、Qwen2-Math上都取得不錯的效果。
PAI-ChatLearn現(xiàn)已全面開源，助力用戶快速、高效的Alignment訓練體驗。借助ChatLearn，用戶可全身心投入于模型設計與效果優(yōu)化，無需分心于底層技術細節(jié)。ChatLearn將承擔起資源調度、數據傳輸、參數同步、分布式運行管理以及確保系統(tǒng)高效穩(wěn)定運作的重任，為用戶提供一站式解決方案。
PAI-ChatLearn背后的技術框架如何設計？性能和效果如何？未來有哪些規(guī)劃？今天一起來深入了解。
PAI-ChatLearn是什么
PAI-ChatLearn是阿里云PAI團隊自研并開源的、靈活易用的、支持大規(guī)模 Alignment高效訓練的框架。
背景
ChatGPT 是由 OpenAI 開發(fā)的基于大型語言模型 (Large Language Model, LLM) 的聊天機器人，以其令人驚嘆的對話能力而迅速火爆并被廣泛采用。ChatGPT 成功背后得益于大型語言模型生成領域的新訓練范式：RLHF (Reinforcement Learning from Human Feedback)，即以強化學習方式依據人類反饋優(yōu)化語言模型。
不同于傳統(tǒng)的深度學習訓練，只涉及到單個模型的迭代和優(yōu)化，以 RLHF 為代表的訓練范式涉及到多個大模型的計算和數據交互，這對于構建構建一個易用、高效的訓練系統(tǒng)提出了諸多的挑戰(zhàn)。
編程接口：如何設計一個通用且靈活的編程接口，讓用戶能專注于單個模型的建模，同時，又可以靈活地控制模型間的交互。
分布式加速引擎：隨著模型規(guī)模的增大，用戶會選擇一些分布式計算和加速的 backend，比如 訓練有Megatron-LM、DeepSpeed 等，推理有vLLM等，如何結合這些加速 backend 來實現(xiàn)高效的多模型計算框架。
并行策略：多個模型可能各有各的計算特點，比如僅推理的模型和訓練的模型在顯存和計算上的特性都不同，每個模型最佳的并行策略也可能不同。因此，框架應該允許不同的模型配置不同的并行策略以發(fā)揮整體的最佳性能。
資源分配：如何靈活地給多個模型分配資源來實現(xiàn)高效的并發(fā)調度和執(zhí)行。
擴展訓練方式：當前Alignment訓練，除了RLHF還有很多變種，例如：DPO/IPO、KTO、ORPO、online DPO、RLAIF等，如何能便捷、快速地擴展訓練流程以適用不同的Alignment訓練。
為了解決上述問題，阿里云PAI團隊提出了一個新的 Alignment 模型訓練框架 ChatLearn。ChatLearn 通過對模型計算邏輯的抽象，解耦了模型和計算 backend、分布式策略的綁定，提供靈活的資源調度機制，可以支持靈活的資源分配和并行調度策略。
PAI-ChatLearn主要特性
ChatLearn的優(yōu)點總結如下:
易用的編程接口: ChatLearn提供通用的編程抽象，用戶只需要封裝幾個函數即可完成模型構造。用戶只需要專注于單模型的編程，系統(tǒng)負責資源調度、數據流傳輸、控制流傳輸、分布式執(zhí)行等。
高可擴展的訓練方式: ChatLearn 提供 RLHF、DPO、OnlineDPO、GRPO 等 Alignment 訓練，同時也支持用戶自定義 model 的執(zhí)行 flow，使定制化訓練流程變得非常便捷。
多種分布式加速引擎: 用戶可以使用不同的計算 backend 進行模型建模，如 Megatron-LM、DeepSpeed、vLLM 等。用戶也可以組合使用不同的 backend，如用 Megatron-LM 來進行加速訓練，用 vLLM 來加速推理。
靈活的并行策略和資源分配: ChatLearn 支持不同模型配置不同的并行策略，可以結合各模型計算、顯存、通信的特點來制定不同的并行策略。同時 ChatLearn 支持靈活的資源調度機制，支持各模型的資源獨占或復用，通過系統(tǒng)調度策略支持高效的串行/并行執(zhí)行和高效的顯存共享。
高性能: 相較于當前的 SOTA 系統(tǒng)，ChatLearn 在 7B+7B (Policy+Reward) 規(guī)模性能提升115%，70B+70B 規(guī)模性能提升 208%。同時，ChatLearn 支持更大規(guī)模的 Alignment 訓練，例如：300B+300B。
PAI-ChatLearn技術架構
PAI-ChatLearn技術架構如上圖：
API：ChatLearn提供了RLHF、DPO、OnlineDPO、GRPO 等 Alignment 訓練，同時也支持用戶自定義 model 的執(zhí)行 flow，來實現(xiàn)自定義的訓練流程。同時ChatLearn提供Module的抽象，用戶通過繼承MegatronModule、DeepSpeedModule、VLLMModule 完成對不同計算backend的封裝。ChatLearn 通過 yaml 文件的形式為 Alignment 訓練，以及不同的模型配置不同的超參數、并行策略等，來實現(xiàn)靈活的模型和并行策略配置。
Scheduler：ChatLearn 提出了 DistActor 的抽象來支持模型的分布式訓練或推理。DistActor 繼承了 Ray actor 的狀態(tài)管理和 worker 間的隔離性，同時突破了 Ray actor 不能跨機的限制。通過 DistActor，ChatLearn 可以支持任意規(guī)模的模型推理和訓練。同時，ChatLearn Scheduler 通過劃分集群 Resource Group 和調度策略，實現(xiàn)硬件感知的親和性調度。ChatLearn 也支持靈活的資源分配，支持模型間的資源復用、獨占或部分復用等策略，在給定資源數的情況下，實現(xiàn)訓練效率的最大化。
Executor：ChatLearn Executor 將 Alignment 訓練流程劃分為三個主要的模塊，Environment、 Trainer和 evaluator。Environment 負責推理模塊模型和數據的并發(fā)執(zhí)行和管理，Trainer 負責相應的訓練模塊，evaluator 負責模型效果評估。Executor 還負責數據傳輸、參數同步。
Backend：得益于 ChatLearn 良好的編程接口抽象，用戶通過簡單的封裝即可接入各種不同 backend 進行計算優(yōu)化和算法優(yōu)化。
Optimization：ChatLearn 也支持各種計算、顯存、通信優(yōu)化，通過各種并行策略組合來加速訓練，通過 paged attention 和 continuous batching 等來加速推理，通過 EMS(Efficient Memory Sharing) 技術來高效復用顯存，減少總資源需求，通過分組廣播技術來支持 Training 和 Inference 模型間高效參數同步，等等。
性能和效果
我們比較了不同參數量規(guī)模模型的 RLHF 訓練吞吐量，采取 N+N 的模型配置，即 Policy 模型和 Reward 模型采用相同大小的參數量。我們和 DeepSpeed-Chat、OpenRLHF 對比了 7B 和 70B 的模型配置，在 8 GPUs 7B+7B 規(guī)模，有 115% 的加速，在 32 GPUs 70B+70B 規(guī)模，有 208% 的加速。規(guī)模越大，加速效果越明顯。同時ChatLearn還能支持更大規(guī)模的 Alignment 訓練，例如：300B+300B 規(guī)模。
ChatLearn支持了Qwen 大模型的Alignment訓練，Qwen2-72B online DPO訓練效果在開源模型中取得領先：
Qwen2-Math-Instruct GRPO訓練效果領先業(yè)界模型：
Roadmap
后續(xù)我們計劃定期發(fā)布Release版本。ChatLearn近期的Roadmap如下：
● 支持Megatron-mcore格式模型；
● 支持MoE模型Alignment訓練；
● 支持更多的模型；
總結
PAI-ChatLearn 是阿里云 PAI 團隊自研的、靈活易用的、支持大規(guī)模 Alignment 高效訓練的框架。ChatLearn 通過對 Alignment 訓練流程進行合理的抽象和解耦，提供靈活的資源分配和并行調度策略。ChatLearn提供了 RLHF、DPO、OnlineDPO、GRPO 等 Alignment 訓練，同時也支持用戶自定義 model 的執(zhí)行 flow，來實現(xiàn)自定義的訓練流程。相比于當前的 SOTA 系統(tǒng)，ChatLearn 在 7B+7B 規(guī)模有 115% 的加速，在 70B+70B規(guī)模有 208% 的加速。同時ChatLearn可以擴展到更大規(guī)模，如：300B+300B(Policy+Reward)。后續(xù)ChatLearn會持續(xù)擴展支持的模型種類，并支持更多的backend進行訓練或推理，同時會持續(xù)優(yōu)化框架性能，簡化使用接口，方便大家進行 Alignment 訓練。歡迎大家來使用、交流和反饋。
開源地址
開源地址：https://github.com/alibaba/ChatLearn
使用文檔：
中文：https://chatlearn.readthedocs.io/zh-cn/latest/
英文：https://chatlearn.readthedocs.io/en/latest/
交流釘群二維碼：
交流釘群號：98090003312
參考文獻
1. https://arxiv.org/pdf/2407.10671
2. https://qwenlm.github.io/blog/qwen2-math/
3. Megatron-LM: https://github.com/NVIDIA/Megatron-LM
4. DeepSpeed-Chat: https://github.com/microsoft/DeepSpeedExamples/tree/master/applications/DeepSpeed-Chat
5. OpenRLHF: https://github.com/OpenRLHF/OpenRLHF
粵公網安備 44030702001206號