PAI-ChatLearn 是阿里云機(jī)器學(xué)習(xí)平臺 PAI 團(tuán)隊(duì)自主研發(fā)的，靈活易用、大規(guī)模模型 RLHF 高效訓(xùn)練框架，支持大模型進(jìn)行 SFT（有監(jiān)督指令微調(diào)）、RM（獎(jiǎng)勵(lì)模型）、RLHF（基于人類反饋的強(qiáng)化學(xué)習(xí)）完整訓(xùn)練流程。PAI-ChatLearn 支持訓(xùn)練和推理組合使用不同的 backend，可對各個(gè)模型配置不同并行策略和靈活的資源分配，支持大規(guī)模（175B + 175B）模型的 RLHF 高效訓(xùn)練，性能比業(yè)界框架有較大提升，有助于用戶專注于模型效果調(diào)優(yōu)。

一、大模型訓(xùn)練方式演進(jìn)

隨著大模型的快速發(fā)展，推動(dòng)了模型訓(xùn)練方式（特別是深度學(xué)習(xí)和人工智能領(lǐng)域）不斷演進(jìn)。隨著模型規(guī)模的增長，單個(gè)設(shè)備（如 GPU 或 CPU）的內(nèi)存和計(jì)算能力已經(jīng)不足以處理整個(gè)模型，因此需要采用并行化策略來分配任務(wù)和計(jì)算負(fù)載。下面分別介紹模型訓(xùn)練的數(shù)據(jù)并行、模型并行和任務(wù)并行。

1、數(shù)據(jù)并行（Data Parallelism）

早期的深度學(xué)習(xí)模型相對較小，如經(jīng)典的卷積神經(jīng)網(wǎng)絡(luò)模型 AlexNet 和 ResNet，可在單個(gè) GPU 上順利運(yùn)行。然而，隨著圖像數(shù)據(jù)量的激增，單 GPU 的訓(xùn)練性能逐漸無法滿足需求。為了應(yīng)對這一挑戰(zhàn)，模型訓(xùn)練開始采用數(shù)據(jù)并行策略來加速。

數(shù)據(jù)并行（Data Parallelism）是將訓(xùn)練數(shù)據(jù)集分割成多個(gè)小批次，并在多個(gè)處理器或設(shè)備上同時(shí)訓(xùn)練模型的副本。每個(gè)設(shè)備都有模型的一個(gè)完整副本，并且在每個(gè)訓(xùn)練步驟結(jié)束時(shí)，通過某種形式的通信（如梯度平均或參數(shù)服務(wù)器）來同步模型參數(shù)。

數(shù)據(jù)并行（Data Parallelism）策略本身也經(jīng)歷了一系列的演進(jìn)。例如，Zero Redundancy Optimizer (ZeRO) 和 Fully Sharded Data Parallel (FSDP) 等概念的出現(xiàn)，它們的核心思想是在訓(xùn)練過程中將狀態(tài)信息、梯度或模型參數(shù)分散存儲(chǔ)，而計(jì)算過程仍然遵循數(shù)據(jù)并行的原則。將訓(xùn)練數(shù)據(jù)集分割成多個(gè)小批次，并在多個(gè)處理器或設(shè)備上同時(shí)訓(xùn)練模型的副本。每個(gè)設(shè)備都有模型的一個(gè)完整副本，并且在每個(gè)訓(xùn)練步驟結(jié)束時(shí)，通過某種形式的通信（如梯度平均或參數(shù)服務(wù)器）來同步模型參數(shù)。

這種方法可以顯著減輕每個(gè) GPU 上的顯存壓力，使得我們能夠訓(xùn)練更大的模型。但隨著模型規(guī)模的增加，數(shù)據(jù)并行策略需要更高效的通信機(jī)制和更大的帶寬來處理增加的同步負(fù)載。

2、模型并行

隨著模型規(guī)模的進(jìn)一步擴(kuò)大，純粹的數(shù)據(jù)并行已不足以滿足訓(xùn)練需求。我們面臨的挑戰(zhàn)不僅是顯存容量不足，還包括計(jì)算能力的限制。因此，模型并行（Model Parallelism）成為了一種大規(guī)模深度學(xué)習(xí)模型訓(xùn)練的常用解決方案。

模型并行依據(jù)模型切分方法主要分為兩種形式：流水并行（Pipeline Parallelism）和張量并行（Tensor Parallelism）：

• 張量并行
  將模型的張量（Tensor）切分成多個(gè)片段，每個(gè)計(jì)算單元（如GPU）負(fù)責(zé)存儲(chǔ)和計(jì)算其中的一部分。這種方法可以有效地減少每個(gè)計(jì)算單元上的內(nèi)存占用，從而支持更大模型的訓(xùn)練。
  缺點(diǎn)：可能會(huì)導(dǎo)致大量的跨設(shè)備通信，因此它更適合于具有高帶寬的內(nèi)部機(jī)器通信，如使用 NVLink 的場景。
• 流水并行（Pipeline Parallelism）
  將模型按照層（Layer）進(jìn)行切分，每個(gè)計(jì)算單元負(fù)責(zé)模型中一部分層的計(jì)算。這種方式同樣可以減輕單個(gè)計(jì)算單元的內(nèi)存負(fù)擔(dān)，并且有助于支持更大規(guī)模的模型訓(xùn)練。通信需求相對較小，適合于基礎(chǔ)設(shè)施較弱、帶寬較低的環(huán)境。
  缺點(diǎn)：可能會(huì)引入大量的空閑時(shí)間（Bubbles），需精心設(shè)計(jì)流水線以避免產(chǎn)生大量的等待時(shí)間。
• 其他模型并行變體
  Mixture of Experts (MoE) 中的專家并行（Expert Parallelism）會(huì)將專家模塊分布到不同的計(jì)算單元上。
  序列并行（Sequence Parallelism）會(huì)在訓(xùn)練過程中將激活函數(shù)的輸出（Activation）分散存儲(chǔ)，以此減少顯存占用。

單一的模型并行往往難以高效地支持模型訓(xùn)練。更大的模型（如GPT3-175B）需要混合多種并行策略?；旌喜⑿校℉ybrid Parallelism）結(jié)合了數(shù)據(jù)并行、模型并行和任務(wù)并行的優(yōu)點(diǎn)，以實(shí)現(xiàn)更高效的訓(xùn)練。如，可以在多個(gè)設(shè)備上使用數(shù)據(jù)并行來訓(xùn)練模型的不同副本，同時(shí)在每個(gè)設(shè)備內(nèi)部使用模型并行來處理模型的不同部分。

3、任務(wù)并行

隨著模型的不斷演進(jìn)，模型訓(xùn)練方式出現(xiàn)了一些新的趨勢。任務(wù)并行不是直接針對模型或數(shù)據(jù)的并行化，而是將不同的任務(wù)分配給不同的處理器或設(shè)備。在大模型訓(xùn)練中，任務(wù)并行可以用來同時(shí)執(zhí)行多個(gè)不同的訓(xùn)練過程，例如，同時(shí)進(jìn)行超參數(shù)搜索、模型評估或不同模型的訓(xùn)練。任務(wù)并行可以提高資源利用率和訓(xùn)練效率，但需要有效的任務(wù)調(diào)度和管理機(jī)制。

• 多任務(wù)/多模態(tài)模型訓(xùn)練
  需要在任務(wù)維度進(jìn)行并行訓(xùn)練，如谷歌提出的 Pathways 架構(gòu)，不同任務(wù)之間可能會(huì)共享參數(shù)。又如 GPT-4v，是多模態(tài)模型，能處理語音、文本等不同類型的數(shù)據(jù)。
• 大模型結(jié)合 RLHF 來提升模型效果，如 ChatGPT、GPT4 等都進(jìn)行了 RLHF。
• 從單程序多數(shù)據(jù)（SPMD）模式向多程序多數(shù)據(jù)（MPMD）模式轉(zhuǎn)變。

（1）RLHF

RLHF (Reinforcement Learning from Human Feedback)是將初始語言模型的微調(diào)任務(wù)建模為 RL 強(qiáng)化學(xué)習(xí)問題，用人類反饋信號來優(yōu)化，使生成更符合人類的期望，提高模型的性能和可靠性。

訓(xùn)練步驟：

• 步驟一：預(yù)訓(xùn)練語言模型
  預(yù)訓(xùn)練語言模型，并進(jìn)行 SFT(Supervised Fine-Tuning)，在人工精心撰寫的語料上進(jìn)行微調(diào)。
  開源的預(yù)訓(xùn)練模型：Qwen、LLaMA、Baichuan 等。
• 步驟二：訓(xùn)練獎(jiǎng)勵(lì)模型(RM)
  訓(xùn)練獎(jiǎng)勵(lì)模型判斷模型生成的文本是否符合人類偏好。
  開源的獎(jiǎng)勵(lì)模型：Ziya-LLaMA-7B-Reward 等。
• 步驟三：使用強(qiáng)化學(xué)習(xí)進(jìn)行微調(diào)
  結(jié)合人類反饋信號，使用強(qiáng)化學(xué)習(xí)來微調(diào)模型。
  社區(qū)模型：ChatGPT、GPT4、GPT4-Turbo 等。

（2）ChatLearn 框架構(gòu)建必要性

RLHF訓(xùn)練需要Policy、Reward、Value 等 models 同時(shí)訓(xùn)練。已有框架如 Megatron-LM、DeepSpeed 等適用于單模型訓(xùn)練/推理。已有的 RLHF 框架只支持單機(jī)訓(xùn)練或者多機(jī)將模型 place 在相同的設(shè)備上，支持的模型最大規(guī)模受限，不支持 70B+70B 或以上規(guī)模，并且無法支持混合框架的訓(xùn)練/推理。

二、PAI-ChatLearn 框架介紹

本節(jié)主要介紹 PAI-ChatLearn 框架的底層實(shí)現(xiàn)細(xì)節(jié)、接口如何定義以及如何實(shí)現(xiàn)大規(guī)模高效的 RLHF 訓(xùn)練。

1、PAI-ChatLearn 框架

PAI_ChatLearn 是一個(gè)靈活且適用于大規(guī)模 RHF 訓(xùn)練的高效框架，其特點(diǎn)如下：

• 支持 SFT（有監(jiān)督的微調(diào)）、RM（獎(jiǎng)勵(lì)模型）、RLHF 全流程訓(xùn)練
• Policy、Value model 的推理和訓(xùn)練分離，支持高效 weight 同步
  由于訓(xùn)練和推理的特點(diǎn)不同，所需資源和并行策略也不盡相同。
  HF 訓(xùn)練過程較復(fù)雜，（如策略模型（Policy Model）既參與生成（generation）也參與訓(xùn)練（training））,會(huì)對模型進(jìn)行分割，如將推理（inference）和訓(xùn)練的模型分開，采用不同的并行策略，以實(shí)現(xiàn)高效的訓(xùn)練和推理。
  實(shí)現(xiàn)了基于高速通道的參數(shù)同步機(jī)制，以確保模型參數(shù)的最終一致性，實(shí)現(xiàn)更高效的訓(xùn)練。
• 訓(xùn)練、推理支持不同 backend，實(shí)現(xiàn)環(huán)境隔離。如在訓(xùn)練中可以使用 Megatron 作為后端，而在推理中可以使用 VIM 框架，這樣的組合可能會(huì)達(dá)到更好的效果。
• 各模型支持獨(dú)立配置并行策略和資源，分布式 actor 的設(shè)計(jì)。
• 支持任意規(guī)模模型(如 175B+175B)訓(xùn)練
• 單機(jī)視角編程，分布式執(zhí)行，用戶在編寫代碼時(shí)不需要關(guān)心底層模型的分配和執(zhí)行。

下圖是在 PAI_ChatLearn 中實(shí)現(xiàn)的 RLHF 訓(xùn)練流程的圖示。

2、PAI-ChatLearn 架構(gòu)簡介

（1）API 層

API 層提供模型建模的接口

• RLHF Module 是一個(gè)通用的抽象 RLHF 模塊接口。通過這個(gè)接口，用戶只需實(shí)現(xiàn)一些方法的封裝，就可以完成不同模型的構(gòu)建。
• RLHF Config 配置 e2e 訓(xùn)練參數(shù)
• Model Config 配置各個(gè) model 的參數(shù)

（2）引擎層

Engine 層負(fù)責(zé)具體資源分配和調(diào)度，以及訓(xùn)練和推理的實(shí)際執(zhí)行。

• DistActor 用來抽象單個(gè)分布式任務(wù)和資源，負(fù)責(zé)各個(gè) model 的執(zhí)行環(huán)境和資源分配，可實(shí)現(xiàn)靈活擴(kuò)展。
• 訓(xùn)練和推理容許使用不同的 Backend。如推理可以使用 PyTorch 或 vLLM 框架，訓(xùn)練后端可以選擇 Megatron、DeepSpeed 或自研的框架。
• 還進(jìn)行了各種優(yōu)化，包括計(jì)算、通信和顯存優(yōu)化，以確保訓(xùn)練性能。

3、PAI-ChatLearn 訓(xùn)練核心流程

PAI-ChatLearn 訓(xùn)練核心流程如下，首先對 chatlearn 進(jìn)行初始化，接著定義模型（RLHF 訓(xùn)練過程包含 2 個(gè) training 模型，4 個(gè)模型做 inference），然后定義 Engine 和數(shù)據(jù)集，調(diào)用 engine.learn() 開始訓(xùn)練。

4、PAI-ChatLearn 訓(xùn)練配置

運(yùn)行環(huán)境配置

PAI-ChatLearn 訓(xùn)練框架中，用戶通過 YAML 文件來配置整個(gè)運(yùn)行環(huán)境，無論是 PAI 的執(zhí)行環(huán)境 DLC 還是本地 local 環(huán)境，都可靈活配置。

在任務(wù)分發(fā)過程中，采用了一種內(nèi)部優(yōu)化策略，即通過 'exclude' 功能，可以明確指定不需要分發(fā)的組件。從而有效減少不必要的通信開銷，提高整體效率。

模型配置

每個(gè)模型有獨(dú)立的配置選項(xiàng)，允許用戶根據(jù)需要自定義資源分配和策略。在 RLHF 中，支持六種不同的模型配置。

RLHF 訓(xùn)練配置

用戶可以定義諸如實(shí)際運(yùn)行的數(shù)量、訓(xùn)練批次大小、生成任務(wù)的基線大小等參數(shù)。此外，還可以設(shè)置定期保存檢查點(diǎn)的時(shí)間間隔、數(shù)據(jù)存儲(chǔ)路徑以及評估模型性能的時(shí)間點(diǎn)。

PAI-ChatLearn 框架提供了清晰的配置流程，用戶只需定義模型并添加相應(yīng)的配置，即可輕松啟動(dòng)訓(xùn)練過程。

5、PAI-ChatLearn 接口-模型定義

為了加速模型構(gòu)建過程，我們提供了一系列的子類，這些子類旨在簡化和加快開發(fā)流程。無論是基于 PyTorch、Meta 還是 VIM 的框架，這些子類都能夠提供強(qiáng)大的支持。同時(shí)，用戶可以選擇基于這些子類或者直接基于父類來構(gòu)建 HF 模型。

在 PAI-ChatLearn 框架中，模型的封裝變得很簡單。用戶只需關(guān)注模型的設(shè)置和構(gòu)建過程，例如初始化（init）環(huán)節(jié)。封裝層旨在屏蔽不同框架之間的接口差異，使得用戶只需在我們的框架中進(jìn)行一層封裝。例如，如果用戶需要實(shí)現(xiàn)前向傳播（forward steps），無論是僅進(jìn)行推理還是進(jìn)行完整的訓(xùn)練周期，我們的封裝都能夠簡化這一過程。用戶只需將模型中已定義的訓(xùn)練模塊進(jìn)行整合，即可實(shí)現(xiàn)標(biāo)準(zhǔn)模型訓(xùn)練的流程。

此外，框架提供了底層接口的靈活選擇，以適應(yīng)不同的執(zhí)行需求。用戶可以選擇 Megatron 作為執(zhí)行引擎，VIM 作為后端，或者采用我們團(tuán)隊(duì)自研的高性能計(jì)算機(jī) HC 作為后端資源。這種自由切換的能力確保了用戶可以根據(jù)具體需求和偏好選擇最合適的工具和資源。

三、PAI-ChatLearn 使用示例

本節(jié)基于開源模型示例講解模型 RLHF 訓(xùn)練的具體流程及訓(xùn)練效果。

1、開源 transformers 模型示例

(1)鏡像和模型準(zhǔn)備

鏡像：

PAI-DLC 平臺上提供已經(jīng)準(zhǔn)備好的鏡像或參考 ChatLearn 提供的 Dockerfile 自定義鏡像

模型：

下載開源模型，例如從 Huggingface 或 ModeScope 上下載。

如果是 HuggingFace transformers 的模型，調(diào)用 ChatLearn 提供的工具將 transformers 模型轉(zhuǎn)成 Megatron-LM 格式，支持修改模型的并行策略。

(2)訓(xùn)練步驟

步驟一：訓(xùn)練 SFT 模型

準(zhǔn)備 SFT 訓(xùn)練的數(shù)據(jù)：

格式：{'query'：問題,'response'：回復(fù)}。

訓(xùn)練方式：

轉(zhuǎn)換好的模型，調(diào)用 ChatLearn 進(jìn)行 SFT 訓(xùn)練。

訓(xùn)練好保存 SFT 模型

如果有訓(xùn)練好的SFT 模型，這個(gè)步驟可以跳過。

步驟二：訓(xùn)練獎(jiǎng)勵(lì)模型(RM)。

準(zhǔn)備 RM 訓(xùn)練的數(shù)據(jù)：

格式：{'query'：問題,'response'：[回復(fù)1, 回復(fù)2, .....],'score'：[score1,score2,.....]}。

訓(xùn)練方法：

定義 RM，一般主體結(jié)構(gòu)復(fù)用 SFT，增加一個(gè) head 層。

使用訓(xùn)練好的 SFT 模型參數(shù)初始化 RM。

調(diào)用 ChatLearn 進(jìn)行 RM 訓(xùn)練。

保存訓(xùn)練好的模型。

步驟三：使用強(qiáng)化學(xué)習(xí)進(jìn)行微調(diào)。

準(zhǔn)備 RL 訓(xùn)練的數(shù)據(jù)：

格式：{"prompt"：問題}。

訓(xùn)練方法：

定義 Policy/Reference Model，結(jié)構(gòu)和 SFT 模型一致。

使用 SFT 模型參數(shù)初始化 Policy/Reference Model。

定義 Reward/Value model，結(jié)構(gòu)和 Reward Model 一致。

使用 RM 模型參數(shù)初始化 Reward/Value model。

調(diào)用 ChatLearn 進(jìn)行 RLHF 訓(xùn)練。

保存訓(xùn)練好的模型。

步驟四：Inference。

離線批量預(yù)測。

準(zhǔn)備好需要預(yù)測的數(shù)據(jù)集合。

調(diào)用 ChatLearn 加載訓(xùn)練好的 RLHF 模型進(jìn)行離線批量預(yù)測。

在線預(yù)測。

調(diào)用 ChatLearn 工具將訓(xùn)練好的 RLHF 模型轉(zhuǎn)成 Transformer 格式。

調(diào)用 PAI-EAS 或 vLLM 部署模型進(jìn)行在線預(yù)測。

2、Vicuna 13B+13B 訓(xùn)練效果

Base 模型使用 Vicuna 13B，RM 也采用 13B 規(guī)模。

在 HH(helpful&harmless)數(shù)據(jù)集進(jìn)行 SFT、RM、RLHF 訓(xùn)練。

RLHF 訓(xùn)練使用了 16 GPUs。

RLHF 訓(xùn)練中 Reward Model 的評分如圖所示。

在 MT-Bench 上使用 GPT-4 API 測評 SFT 和 RLHF 模型的效果，平均得分提升 11%。

• ChatLearn 性能：
  A800 上測試不同規(guī)模 RLHF 訓(xùn)練 ChatLearn 的性能。
  Policy 和 Reward Model 使用相同規(guī)模配置。
  相比于 DeepSpeed-chat，7B+7B~30B+30B 規(guī)模比 DeepSpeed-chat 快 48%~82%。
  66B + 66B 時(shí) 32 GPUs 上 DeepSpeed-chat 會(huì) OOM，ChatLearn 可以正常訓(xùn)練。
  ChatLearn 可以支持 175B+175B 等更大規(guī)模模型訓(xùn)練。

3、PAI-ChatLearn 效果

• Qwen-14B 上的效果：

Policy 和 Reward Model 都是 14B；

基礎(chǔ)評測上效果如上圖所示，優(yōu)于同規(guī)模的其他開源模型；

構(gòu)造一個(gè)包含廣泛主題的測試數(shù)據(jù)集進(jìn)行人工評測，RLHF 訓(xùn)練的模型效果明顯優(yōu)于 SFT 模型。測試題包括知識類、語言理解類，創(chuàng)作類、數(shù)學(xué)類，編碼類等能力。

四、問答環(huán)節(jié)

Q1：Reward 模型中有多種評分機(jī)制，如科學(xué)準(zhǔn)確性、合規(guī)性以及人性化表達(dá)等，是否有多維度獎(jiǎng)勵(lì)模型評測的實(shí)際工程經(jīng)驗(yàn)？

A1：百川 2 技術(shù)報(bào)告中提到模型中的 HF 訓(xùn)練，使用了不止一個(gè)獎(jiǎng)勵(lì)模型。如在數(shù)學(xué)領(lǐng)域，現(xiàn)有的獎(jiǎng)勵(lì)模型在數(shù)學(xué)評分方面可能不夠精準(zhǔn)，效果不理想，他們會(huì)專門采用一個(gè)單獨(dú)的數(shù)學(xué)能力更強(qiáng)的模塊來輔助獎(jiǎng)勵(lì)模型的訓(xùn)練。

一些開源框架，由于是平鋪設(shè)計(jì)，會(huì)遇到顯存不足等各種挑戰(zhàn)，難以擴(kuò)展。

而 PAI-ChatLearn 框架，其采用分離式架構(gòu)設(shè)計(jì)，方便了在不同階段進(jìn)行各種擴(kuò)展。只需在獎(jiǎng)勵(lì)模型部分添加一系列處理邏輯，即可實(shí)現(xiàn)多個(gè)獎(jiǎng)勵(lì)模型的結(jié)合。由于每個(gè)模型都被封裝成一個(gè)獨(dú)立的分布式 actor，所以只需添加一個(gè)新的分布式 actor，并將其結(jié)果串聯(lián)起來，最后用KL散度或其他方法進(jìn)行計(jì)算即可。

Q2：在使用 PAI 開發(fā)套件中，遇到從 Hugging Face 轉(zhuǎn) Megatron 模型進(jìn)行 TP或 PP 切分后保存于磁盤會(huì)變大很多倍的問題，請問是什么原因，如何解決？

A2：對模型進(jìn)行 TP 張量并行或者 PP 流水并行的切分是不會(huì)改變模型大小的。

老版本的 Megatron 上有個(gè) bug，進(jìn)行轉(zhuǎn)換后會(huì)保存多份，最新版本已經(jīng)沒有這個(gè)問題了。

Q3：PAI_ChatLearn 中實(shí)現(xiàn)的 RLHF 訓(xùn)練流程圖中有很多變量，是否可以任意進(jìn)行配置？

A3：這個(gè)是一個(gè)標(biāo)準(zhǔn)的 PPO 實(shí)現(xiàn)，這些參數(shù)是必要的。

PAI-ChatLearn 會(huì)把這個(gè) inference 和 training 分開，增加了參數(shù)同步，也即最上面的線的部分。