當前，多模態模型是 AI 領域研究的最熱門方向之一。多模態大語言模型從大規模預訓練中獲得圖文理解、創作、知識、推理、指令遵循等能力，透過可監督微調激發對應的能力。此外，模型輸出還要具備特定的風格、符合人類偏好、對齊人類價值觀，因此需要引入基於人類的反饋訊號的強化學習（RLHF）方法來進一步最佳化模型效果。

PPO（Proximal Policy Optimization）演算法是 OpenAI 在 RLHF 階段採用的演算法。PPO 演算法中涉及到多個模型的協同訓練和推理，設計和實現一套高效、準確的 RLHF 訓練系統是多模態模型研究領域的關鍵挑戰之一。

在 2024 年的 QCon 上海站上，小紅書資深技術專家、RLHF 自研框架負責人於子淇發表了題為《基於 PPO 的多模態大模型 RLHF 系統的設計與最佳化》的演講。本次演講主要介紹小紅書大模型團隊自研 MLLM RLHF 訓練框架的實現以及效能最佳化，分析了小紅書團隊如何透過訓練和推理的混布排程最佳化等手段實現極致的模型效能，希望能給大家帶來一些幫助。

強化學習（RL）的核心過程是：一個智慧體在時間 t 處於狀態 St，經過動作 A 後其狀態變化為 St+1，同時環境會給智慧體一個獎勵訊號 R。當 RL 應用於 NLP 領域時，模型收到提示後開始輸出 token，每輸出一個 token 就是一個動作，每個 token 後模型會收到一個獎勵模型提供的獎勵訊號對 token 打分，再輸出下一個 token，由此逐漸輸出符合期望的輸出。當獎勵模型的打分策略是人類制定時，此過程就稱為基於人類反饋訊號（HF）的強化學習（RLHF）。

RLHG 的流程主要由獎勵模型 RM 和強化學習 RL 過程組成。第一步，透過 SFT 監督訓練過程初始化 RM 和 RL 模型。接下來分別訓練 RM 和 RL 模型。

RM 模型的資料構造是對 LLM 模型輸入同一個提示取樣多個不同輸出，生成多個 pair 對。之後人類專家會對這些 pair 對進行質量排序，生成資料集，然後提供給模型進行 pair-loss 偏好訓練。之後 RM 模型就可以對輸入的 pair 對進行打分了。

RL 模型訓練有多種方法，其中有代表性的是 PPO 策略梯度演算法，這是一種 On-Policy 演算法，需要線上生成不同的樣本即時更新。OpenAI 作為行業先驅，驗證了 PPO 演算法在 RLHF 領域的有效性，因此行業進行 RLHF 研究時多采用 PPO。

PPO 演算法包含四個模型。首先是演員模型（Actor Model），是訓練的目標語言模型。第二是評論家模型（Critic Model），負責預估模型當前行為的總收益。第三是獎勵模型（Reward Model），計算當前動作的即時收益。第四是參考模型（Reference Model），是演員模型在 Step 0 階段的引數狀態，它會不斷與狀態更新的演員模型對比計算 diff，從而約束模型輸出，防止輸出偏移，遺忘上游任務。

PPO 的訓練流程可以簡單概括為訓推混布，整體可以抽象為兩大步驟，第一步是經驗取樣（rollout），是指演員模型對當前提示作出一次響應（generate 階段），然後根據提示 – 響應計算獎勵分數，從而構造獎勵的訓練資料集（forward 階段）。第二步就是訓練，也就是對演員和評論家模型進行 PPO 訓練迭代。

上述流程存在一些痛點，首先是流程需要多個 dataloader；其次，演員模型包含了生成自迴歸部分，因此需要 serving 部分；第三，該流程包含了四個模型的推理過程，複雜度較高；最後，該流程需要訓練演員和評論家兩個模型，工作量加大。

基於上述背景，小紅書團隊設計了一個 RLHF 框架。

小紅書團隊設計的整體架構汲取了開源領域的經驗，抽象出了只用 Megatron-core 做訓練，並用 ray 做排程的經典方式來降低複雜度。推理方面則採用 vLLM 推理引擎來加速。因為訓練和推理存在 diff，如果將 diff 引入訓練過程，就會給 PPO 過程帶來 bias，為了避免 bias 需要使用 megatron 計算 logp 機率。

整體架構來看，首先模型會透過提示生成響應，之後輸入四個模型進行前向推理，生成結果經過後處理後生成獎勵資料集，再提供給演員和評論家模型進行訓練。訓練過程是 on-policy 的，需要即時取樣當前樣本。

上述框架的前向推理階段有四個模型，如果採用並行方法會帶來四倍的 GPU 記憶體佔用。因此團隊採用了 offloading 的方法，將結構相同的評論家與獎勵模型，演員與參考模型分為兩組複用，從而將推理模型數量減少到兩個，降低了記憶體佔用。

另一方面，團隊還將演員和評論家模型的訓練過程放在了獨立的叢集上進行非同步訓練，從而提升了 50% 的效能。非同步訓練啟動時，演員叢集上進行 torchrun 啟動主程序訓練，並自動以 mock 方式拉起評論家叢集。這兩個叢集構成 master-worker 架構，叢集之間的資料傳輸會帶來一定開銷。

上述異構組網架構投入使用後，小紅書團隊發現，在資料量、模型引數量等壓力增大後，叢集壓力會大幅上升。以 Llama 3 70B 為例，SFT 單個模型本身的訓練就需要 4 機 32 卡，32k 長上下文時需要開啟 CP2 並行，此時就需要 8 機 64 卡，推理過程需要 128 卡。另一方面，兩個訓練叢集需要同步快取，隨著序列長度增加等因素，同步耗時也會顯著增長。

為解決以上問題，團隊設計了一個同構組網架構。新的架構中，由於評論家與演員模型還能繼續 offload，因此可以將四個模型進一步減少到一個主模型來複用，透過 offload 切換模型，降低推理成本。訓練階段則 offload 了很多中間狀態，儘可能減少資源佔用。新的架構可以充分利用資源，提供較高的效能。

在訓練效能最佳化方面，小紅書團隊引入了一些常規最佳化方法，包括資料載入的預讀取、雙 dataloader；一些常見的並行最佳化策略，如 TP、PP、CP、SP；常見的長文字場景下的視訊記憶體最佳化（recompute 技術）；Dynamic-batch，訓練和推理使用不同 batchsize；負載均衡方面，提示請求會透過 round-robin 方式傳送到多個 vLLM 引擎；vLLM 推理引擎會做不同粒度的切分，雖然降低了單粒度的 running-batch，但能提升併發性。

流水線最佳化方面，團隊發現演員模型取樣階段是無需訓練的，存在較大空窗。針對這一問題的解決思路是訓推混布。由於訓練任務中的簡單任務耗時佔比超過 80%，且訓練任務負載遠高於推理，因此可以透過增大 serving 併發來降低訓練叢集負載。另一方面，生成和前向推理階段可以進行流水線並行。基於上述思路改進了整體 RL 架構，實現了全量 offload 和流水線並行。

新的架構將整體 reward-batch 切分為多個 mini batch，從而 overlap 遮蔽了資料落盤、vllm-restart/offload、forward、資料預處理等上下文開銷。與之前的版本相比，流水線新架構的訓練耗時從 2250s 減少到了 690s。

r1/o1 系列 reasoner 推理模型場景下 decode 長度會大幅增加，比如 128k response 生成會成為瓶頸，此時全量 offload 更適合這種訓推不均衡的場景。

進行效能最佳化時，PPO 部分有一些細節可以進一步挖掘。首先是 Padding-free，這裡透過 start-end 的 offset 去除了訓練階段的 padding，減少了 10% 的耗時並降低了記憶體佔用。

第二個細節是引數同步。演員模型需要把引數同步給推理引擎，這裡使用新的通訊元件降低了通訊耗時。

第三個細節是 CP 並行，這是 LLM 在長文字場景中必須開啟的特性。這裡需要進行 logp、values、reward 的 CP 重排。對各個重排過程最佳化後也可以進一步降低通訊開銷。

最後是 Logp 實現：在 actor_loss 計算階段對 logits 提前重排，計算 CP 的結果之後進行 allgather，得到 total 的結果，再 -1 得到 logp，從而降低通訊量（bxsxh -> bx）。

MLLM 多模態模型的效能最佳化存在一些痛點，主要是引入了影像處理、圖文混合場景，從而導致訓練負載不均衡。不過多模態模型的視覺處理部分計算量雖然巨大，但引數量較小，在 PPO 場景下影像處理能力是存在冗餘的。

基於上述特點引入三個最佳化方法。首先是組網最佳化，在語言處理場景採用 TP、PP、CP 等並行最佳化，而視覺部分可以複用這些切分最佳化。其次是多路複用，這裡對 vLLM、演員和評論家模型全部複用相同的 img_feature。第三是預讀取最佳化，這裡會凍結 visual-model，將 img_feature 的全部計算流程放入 make-experience 階段內部，透過流水線方式 overlap 多模計算需求。透過上述最佳化，MLLM 可以實現與 LLM 相當的效能表現。

進入業務部署階段，小紅書團隊在訓推一致性層面發現了一些痛點。首先是 RM-serving 過程中存在 acc 掉點，其次是推理和評估時的精度與訓練時並不完全一致。分析後發現，rm 的 vhead 是一個 linear，本質上是一個矩陣乘，會累積放大誤差，初始誤差 diff 只有 0.001，累積後可達 0.3。

針對這一問題的解決方案是針對 RM+RL 複用相同的網路結構和訓練框架（megatron-core），同時針對 RL 場景下的 RM-serving 轉變為本地 offload 實現，從而保證精度完全一致。最後，訓練和推理採用 mcore 的計算負載，保證評測任務上訓推 tokenwise 一致。

即便有了前述最佳化過程，演員模型部分的生成耗時依舊相對較高。對此一般使用推理量化方式進行加速，但傳統量化方式有一定精度損失，無法做到 PPO 訓練所需的無損最佳化。

團隊發現投機取樣是一種無損最佳化方式，因此使用了 Medusa 演算法來做最佳化。投機取樣是一種計算換空間的方式，利用 FLOPS 換取 VRAM 頻寬。這裡在 LLM Transformer 級別的最後一層新增一個 Medusa head，其結構與 LLM 同構。訓練時為了確保精度需要凍結 LLM 主幹。這裡的 Running-batch 提升到最大，從而實現最大吞吐量。

另一個問題是 backbone 與 mhead 的匹配度不足導致接收率下降，降低加速比。這裡引入了伴生訓練的方式來解決問題。伴生訓練要求精度無損，所以 backbone 要即時更新，mhead 更新則可以延後（因為不影響精度），從而可以實現訓練 overlap。此外，訓練時 mhead 可以少量微調，其資料量只需要 rollout-batchsize 即可。透過這些方式，生成速度得到了 50% 提升（注意：投機取樣會引入 verify 的額外計算量，適合推理併發過大時導致 batchsize_per_engine 過小沒有充分利用計算資源的場景）。

在一個基於開源模型的案例中，使用新的 RLHF 框架後，模型可以正確遵循原本無法正確識別的提示。例如提示詢問“午餐吃什麼”，原本的模型回答給出了早餐的選項，改進後可以正確回答午餐的選項。

另一個案例則顯示出了模型答案內容豐富性的改進，針對提示問題，原始的模型答案只有簡單的是 / 否結果，新的答案則給出了詳細的解釋。

進一步的評測顯示，該 RLHF 框架對開源模型進行改進後，模型的推理、創作、問答、數學、對話和程式碼能力分別提升了 15%、7%、7%、11%、4%、3%，綜合提升 6%。

該 RLHF 框架還可以提示模型的 PRM 效果。PRM 是指過程獎勵，PRM 效果提升後，模型回答問題時的可解釋性會更好。更進一步的評估顯示，PRM 可以改善模型的激勵粒度，帶來 5% 的模型能力上限提升。

在 PPO 調參方面，團隊也積累了一些經驗：

團隊未來計劃進行更深層次的探索。首先是效能最佳化，團隊自主研發了 MARIX 專案，該專案是全流程訓推一體框架，未來將進一步壓縮生成耗時、不同序列的負載均衡，做深度流水線排程最佳化，等等。其次是演算法探索，這方面的目標是去掉 RLHF 的 HF 人工獎勵部分，結合 RL-COT 打造更深層的推理能力，實現真正的 RL scaling-law。

於子淇，小紅書資深技術專家，RLHF 自研框架負責人，主要從事 RLHF 系統從 0 到 1 的構建以及訓推一體效能最佳化，幫助團隊拿到端到端的模型效果收益；曾擔任阿里雲高階開發工程師，卓越工程師，技術佈道師，參與 AIACC（ DawnBench 榜單第一）加速庫開發，NCCL 通訊庫最佳化等分散式訓練的效能最佳化工作。

在 AI 大模型重塑軟體開發的時代，我們如何把握變革？如何突破技術邊界？4 月 10-12 日，QCon 全球軟體開發大會· 北京站 邀你共赴 3 天沉浸式學習之約，跳出「技術繭房」，探索前沿科技的無限可能。

本次大會將匯聚頂尖技術專家、創新實踐者，共同探討多行業 AI 落地應用，分享一手實踐經驗，深度參與 DeepSeek 主題圓桌，洞見未來趨勢。