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

一、大模型訓練方式演進

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

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

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

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

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

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

2、模型并行

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

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

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

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

3、任務并行

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

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

（1）RLHF

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

訓練步驟：

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

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

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

二、PAI-ChatLearn 框架介紹

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

1、PAI-ChatLearn 框架

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

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

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

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

（1）API 層

API 層提供模型建模的接口

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

（2）引擎層

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

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

3、PAI-ChatLearn 訓練核心流程

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

4、PAI-ChatLearn 訓練配置

運行環(huán)境配置

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

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

模型配置

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

RLHF 訓練配置

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

PAI-ChatLearn 框架提供了清晰的配置流程，用戶只需定義模型并添加相應的配置，即可輕松啟動訓練過程。

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

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

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

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

三、PAI-ChatLearn 使用示例

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

1、開源 transformers 模型示例

(1)鏡像和模型準備

鏡像：

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

模型：

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

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

(2)訓練步驟

步驟一：訓練 SFT 模型

準備 SFT 訓練的數(shù)據(jù)：

格式：{'query'：問題,'response'：回復}。

訓練方式：

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

訓練好保存 SFT 模型

如果有訓練好的SFT 模型，這個步驟可以跳過。

步驟二：訓練獎勵模型(RM)。

準備 RM 訓練的數(shù)據(jù)：

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

訓練方法：

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

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

調(diào)用 ChatLearn 進行 RM 訓練。

保存訓練好的模型。

步驟三：使用強化學習進行微調(diào)。

準備 RL 訓練的數(shù)據(jù)：

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

訓練方法：

定義 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 進行 RLHF 訓練。

保存訓練好的模型。

步驟四：Inference。

離線批量預測。

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

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

在線預測。

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

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

2、Vicuna 13B+13B 訓練效果

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

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

RLHF 訓練使用了 16 GPUs。

RLHF 訓練中 Reward Model 的評分如圖所示。

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

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

3、PAI-ChatLearn 效果

• Qwen-14B 上的效果：

Policy 和 Reward Model 都是 14B；

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

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

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

Q1：Reward 模型中有多種評分機制，如科學準確性、合規(guī)性以及人性化表達等，是否有多維度獎勵模型評測的實際工程經(jīng)驗？

A1：百川 2 技術報告中提到模型中的 HF 訓練，使用了不止一個獎勵模型。如在數(shù)學領域，現(xiàn)有的獎勵模型在數(shù)學評分方面可能不夠精準，效果不理想，他們會專門采用一個單獨的數(shù)學能力更強的模塊來輔助獎勵模型的訓練。

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

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

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

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

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

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

A3：這個是一個標準的 PPO 實現(xiàn)，這些參數(shù)是必要的。

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