導(dǎo)讀

ChatGPT是OpenAI開發(fā)的基于大型語言模型(LLM)的聊天機(jī)器人，以其令人驚嘆的對(duì)話能力而迅速火爆并被廣泛采用。ChatGPT 成功背后得益于大型語言模型生成領(lǐng)域的新訓(xùn)練范式：RLHF (Reinforcement Learning from Human Feedback)，即以強(qiáng)化學(xué)習(xí)方式依據(jù)人類反饋優(yōu)化語言模型。不同于傳統(tǒng)的深度學(xué)習(xí)訓(xùn)練，只涉及到單個(gè)模型的迭代和優(yōu)化，以 RLHF 為代表的Alignment訓(xùn)練范式涉及到多個(gè)大模型的計(jì)算和數(shù)據(jù)交互，這對(duì)于構(gòu)建構(gòu)建一個(gè)易用、高效、可擴(kuò)展的訓(xùn)練系統(tǒng)提出了諸多的挑戰(zhàn)。

為了解決上述問題，我們提出了一個(gè)新的 Alignment 訓(xùn)練框架 PAI-ChatLearn。ChatLearn 通過對(duì) Alignment 訓(xùn)練流程進(jìn)行合理的抽象和解耦，提供靈活的資源分配和模型調(diào)度策略。ChatLearn提供了RLHF、DPO、OnlineDPO、GRPO等 Alignment 訓(xùn)練，同時(shí)也支持用戶自定義 model 的執(zhí)行 flow，來實(shí)現(xiàn)自定義的訓(xùn)練流程。相比于當(dāng)前的 SOTA 系統(tǒng)，ChatLearn 在 7B+7B (Policy+Reward) 規(guī)模性能提升115%，70B+70B 規(guī)模性能提升208%，并且能支持更大規(guī)模的Alignment訓(xùn)練，例如300B+300B規(guī)模。同時(shí)ChatLearn也一直在支持Qwen大模型的Alignment訓(xùn)練，在Qwen-Chat、Qwen2-Chat、Qwen2-Math上都取得不錯(cuò)的效果。

PAI-ChatLearn現(xiàn)已全面開源，助力用戶快速、高效的Alignment訓(xùn)練體驗(yàn)。借助ChatLearn，用戶可全身心投入于模型設(shè)計(jì)與效果優(yōu)化，無需分心于底層技術(shù)細(xì)節(jié)。ChatLearn將承擔(dān)起資源調(diào)度、數(shù)據(jù)傳輸、參數(shù)同步、分布式運(yùn)行管理以及確保系統(tǒng)高效穩(wěn)定運(yùn)作的重任，為用戶提供一站式解決方案。

PAI-ChatLearn背后的技術(shù)框架如何設(shè)計(jì)？性能和效果如何？未來有哪些規(guī)劃？今天一起來深入了解。

PAI-ChatLearn是什么

PAI-ChatLearn是阿里云PAI團(tuán)隊(duì)自研并開源的、靈活易用的、支持大規(guī)模 Alignment高效訓(xùn)練的框架。

背景

ChatGPT 是由 OpenAI 開發(fā)的基于大型語言模型 (Large Language Model, LLM) 的聊天機(jī)器人，以其令人驚嘆的對(duì)話能力而迅速火爆并被廣泛采用。ChatGPT 成功背后得益于大型語言模型生成領(lǐng)域的新訓(xùn)練范式：RLHF (Reinforcement Learning from Human Feedback)，即以強(qiáng)化學(xué)習(xí)方式依據(jù)人類反饋優(yōu)化語言模型。

 

 

不同于傳統(tǒng)的深度學(xué)習(xí)訓(xùn)練，只涉及到單個(gè)模型的迭代和優(yōu)化，以 RLHF 為代表的訓(xùn)練范式涉及到多個(gè)大模型的計(jì)算和數(shù)據(jù)交互，這對(duì)于構(gòu)建構(gòu)建一個(gè)易用、高效的訓(xùn)練系統(tǒng)提出了諸多的挑戰(zhàn)。

編程接口：如何設(shè)計(jì)一個(gè)通用且靈活的編程接口，讓用戶能專注于單個(gè)模型的建模，同時(shí)，又可以靈活地控制模型間的交互。

分布式加速引擎：隨著模型規(guī)模的增大，用戶會(huì)選擇一些分布式計(jì)算和加速的 backend，比如 訓(xùn)練有Megatron-LM、DeepSpeed 等，推理有vLLM等，如何結(jié)合這些加速 backend 來實(shí)現(xiàn)高效的多模型計(jì)算框架。

并行策略：多個(gè)模型可能各有各的計(jì)算特點(diǎn)，比如僅推理的模型和訓(xùn)練的模型在顯存和計(jì)算上的特性都不同，每個(gè)模型最佳的并行策略也可能不同。因此，框架應(yīng)該允許不同的模型配置不同的并行策略以發(fā)揮整體的最佳性能。

資源分配：如何靈活地給多個(gè)模型分配資源來實(shí)現(xiàn)高效的并發(fā)調(diào)度和執(zhí)行。

擴(kuò)展訓(xùn)練方式：當(dāng)前Alignment訓(xùn)練，除了RLHF還有很多變種，例如：DPO/IPO、KTO、ORPO、online DPO、RLAIF等，如何能便捷、快速地?cái)U(kuò)展訓(xùn)練流程以適用不同的Alignment訓(xùn)練。

為了解決上述問題，阿里云PAI團(tuán)隊(duì)提出了一個(gè)新的 Alignment 模型訓(xùn)練框架 ChatLearn。ChatLearn 通過對(duì)模型計(jì)算邏輯的抽象，解耦了模型和計(jì)算 backend、分布式策略的綁定，提供靈活的資源調(diào)度機(jī)制，可以支持靈活的資源分配和并行調(diào)度策略。

PAI-ChatLearn主要特性

ChatLearn的優(yōu)點(diǎn)總結(jié)如下:

易用的編程接口: ChatLearn提供通用的編程抽象，用戶只需要封裝幾個(gè)函數(shù)即可完成模型構(gòu)造。用戶只需要專注于單模型的編程，系統(tǒng)負(fù)責(zé)資源調(diào)度、數(shù)據(jù)流傳輸、控制流傳輸、分布式執(zhí)行等。

高可擴(kuò)展的訓(xùn)練方式: ChatLearn 提供 RLHF、DPO、OnlineDPO、GRPO 等 Alignment 訓(xùn)練，同時(shí)也支持用戶自定義 model 的執(zhí)行 flow，使定制化訓(xùn)練流程變得非常便捷。

多種分布式加速引擎: 用戶可以使用不同的計(jì)算 backend 進(jìn)行模型建模，如 Megatron-LM、DeepSpeed、vLLM 等。用戶也可以組合使用不同的 backend，如用 Megatron-LM 來進(jìn)行加速訓(xùn)練，用 vLLM 來加速推理。

靈活的并行策略和資源分配: ChatLearn 支持不同模型配置不同的并行策略，可以結(jié)合各模型計(jì)算、顯存、通信的特點(diǎn)來制定不同的并行策略。同時(shí) ChatLearn 支持靈活的資源調(diào)度機(jī)制，支持各模型的資源獨(dú)占或復(fù)用，通過系統(tǒng)調(diào)度策略支持高效的串行/并行執(zhí)行和高效的顯存共享。

高性能: 相較于當(dāng)前的 SOTA 系統(tǒng)，ChatLearn 在 7B+7B (Policy+Reward) 規(guī)模性能提升115%，70B+70B 規(guī)模性能提升 208%。同時(shí)，ChatLearn 支持更大規(guī)模的 Alignment 訓(xùn)練，例如：300B+300B。

PAI-ChatLearn技術(shù)架構(gòu)

 

PAI-ChatLearn技術(shù)架構(gòu)如上圖：

API：ChatLearn提供了RLHF、DPO、OnlineDPO、GRPO 等 Alignment 訓(xùn)練，同時(shí)也支持用戶自定義 model 的執(zhí)行 flow，來實(shí)現(xiàn)自定義的訓(xùn)練流程。同時(shí)ChatLearn提供Module的抽象，用戶通過繼承MegatronModule、DeepSpeedModule、VLLMModule 完成對(duì)不同計(jì)算backend的封裝。ChatLearn 通過 yaml 文件的形式為 Alignment 訓(xùn)練，以及不同的模型配置不同的超參數(shù)、并行策略等，來實(shí)現(xiàn)靈活的模型和并行策略配置。

Scheduler：ChatLearn 提出了 DistActor 的抽象來支持模型的分布式訓(xùn)練或推理。DistActor 繼承了 Ray actor 的狀態(tài)管理和 worker 間的隔離性，同時(shí)突破了 Ray actor 不能跨機(jī)的限制。通過 DistActor，ChatLearn 可以支持任意規(guī)模的模型推理和訓(xùn)練。同時(shí)，ChatLearn Scheduler 通過劃分集群 Resource Group 和調(diào)度策略，實(shí)現(xiàn)硬件感知的親和性調(diào)度。ChatLearn 也支持靈活的資源分配，支持模型間的資源復(fù)用、獨(dú)占或部分復(fù)用等策略，在給定資源數(shù)的情況下，實(shí)現(xiàn)訓(xùn)練效率的最大化。

Executor：ChatLearn Executor 將 Alignment 訓(xùn)練流程劃分為三個(gè)主要的模塊，Environment、 Trainer和 evaluator。Environment 負(fù)責(zé)推理模塊模型和數(shù)據(jù)的并發(fā)執(zhí)行和管理，Trainer 負(fù)責(zé)相應(yīng)的訓(xùn)練模塊，evaluator 負(fù)責(zé)模型效果評(píng)估。Executor 還負(fù)責(zé)數(shù)據(jù)傳輸、參數(shù)同步。

Backend：得益于 ChatLearn 良好的編程接口抽象，用戶通過簡(jiǎn)單的封裝即可接入各種不同 backend 進(jìn)行計(jì)算優(yōu)化和算法優(yōu)化。

Optimization：ChatLearn 也支持各種計(jì)算、顯存、通信優(yōu)化，通過各種并行策略組合來加速訓(xùn)練，通過 paged attention 和 continuous batching 等來加速推理，通過 EMS(Efficient Memory Sharing) 技術(shù)來高效復(fù)用顯存，減少總資源需求，通過分組廣播技術(shù)來支持 Training 和 Inference 模型間高效參數(shù)同步，等等。

性能和效果

我們比較了不同參數(shù)量規(guī)模模型的 RLHF 訓(xùn)練吞吐量，采取 N+N 的模型配置，即 Policy 模型和 Reward 模型采用相同大小的參數(shù)量。我們和 DeepSpeed-Chat、OpenRLHF 對(duì)比了 7B 和 70B 的模型配置，在 8 GPUs 7B+7B 規(guī)模，有 115% 的加速，在 32 GPUs 70B+70B 規(guī)模，有 208% 的加速。規(guī)模越大，加速效果越明顯。同時(shí)ChatLearn還能支持更大規(guī)模的 Alignment 訓(xùn)練，例如：300B+300B 規(guī)模。

 

 

ChatLearn支持了Qwen 大模型的Alignment訓(xùn)練，Qwen2-72B online DPO訓(xùn)練效果在開源模型中取得領(lǐng)先：

 

 

Qwen2-Math-Instruct GRPO訓(xùn)練效果領(lǐng)先業(yè)界模型：

 

 

Roadmap

后續(xù)我們計(jì)劃定期發(fā)布Release版本。ChatLearn近期的Roadmap如下：

● 支持Megatron-mcore格式模型；

● 支持MoE模型Alignment訓(xùn)練；

● 支持更多的模型；

總結(jié)

PAI-ChatLearn 是阿里云 PAI 團(tuán)隊(duì)自研的、靈活易用的、支持大規(guī)模 Alignment 高效訓(xùn)練的框架。ChatLearn 通過對(duì) Alignment 訓(xùn)練流程進(jìn)行合理的抽象和解耦，提供靈活的資源分配和并行調(diào)度策略。ChatLearn提供了 RLHF、DPO、OnlineDPO、GRPO 等 Alignment 訓(xùn)練，同時(shí)也支持用戶自定義 model 的執(zhí)行 flow，來實(shí)現(xiàn)自定義的訓(xùn)練流程。相比于當(dāng)前的 SOTA 系統(tǒng)，ChatLearn 在 7B+7B 規(guī)模有 115% 的加速，在 70B+70B規(guī)模有 208% 的加速。同時(shí)ChatLearn可以擴(kuò)展到更大規(guī)模，如：300B+300B(Policy+Reward)。后續(xù)ChatLearn會(huì)持續(xù)擴(kuò)展支持的模型種類，并支持更多的backend進(jìn)行訓(xùn)練或推理，同時(shí)會(huì)持續(xù)優(yōu)化框架性能，簡(jiǎn)化使用接口，方便大家進(jìn)行 Alignment 訓(xùn)練。歡迎大家來使用、交流和反饋。

 

開源地址

開源地址：https://github.com/alibaba/ChatLearn

 

使用文檔：

中文：https://chatlearn.readthedocs.io/zh-cn/latest/

英文：https://chatlearn.readthedocs.io/en/latest/

 

交流釘群二維碼：

 

交流釘群號(hào)：98090003312

 

參考文獻(xiàn)

 

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