一、背景

最近一直在關(guān)注 RL Infra 相關(guān)的工作，尤其是 RL 性能優(yōu)化，后續(xù)會逐漸介紹一下該領(lǐng)域的相關(guān)文章，本文先簡單介紹一下字節(jié)新發(fā)布的 RhymeRL。

對應(yīng)的論文為：[2508.18588] History Rhymes: Accelerating LLM Reinforcement Learning with RhymeRL

二、摘要

RL 成為提升 LLM Reasoning 能力的關(guān)鍵方法，與傳統(tǒng)預(yù)訓(xùn)練不同，RL 包含多個階段：Rollout、Reward、Training，需要多種類型的 Worker 協(xié)同配合；除此之外，為了效率也可能引入異步訓(xùn)練方式，需要考慮效率與效果的 Tradeoff。然而，當(dāng)前 RL 系統(tǒng)面臨 GPU 利用率低的問題，主要原因為：Rollout 在整個 RL 過程中占比很高，并且 Rollout 的過程中很容易出現(xiàn)負載不均，導(dǎo)致 GPU Bubble。而異步訓(xùn)練、截斷等方式雖然可以緩解該問題，但是可能犧牲準(zhǔn)確率。

RhymeRL 中，作者發(fā)現(xiàn)相鄰訓(xùn)練 Epoch 中，Rollout 的 Response 表現(xiàn)出高度相似性?；诖耍O(shè)計了用于加速 Rollout 的方案：

• HistoSpec，基于相似性，用上一個 Epoch 的 Response 作為草稿，使用投機采樣技術(shù)加速 Rollout 生成速度。
• HistoPipe，基于相似性，用上一個 Epoch 的 Response 的長度分布來預(yù)測當(dāng)前 Epoch 的 Response 長度，進而實現(xiàn)負載均衡。

此外，作者也解決了其中的相關(guān)問題，比如動態(tài)調(diào)整投機采樣的草稿長度，以避免引入過多無效計算；以及引入雙層調(diào)度策略，進一步實現(xiàn)負載的均衡。

作者在真實生產(chǎn)環(huán)境中進行評估，證明其具備從數(shù)十到數(shù)千 GPU 的擴展能力。實驗結(jié)果表明，RhymeRL 在保持準(zhǔn)確性的前提下，相較于現(xiàn)有方法實現(xiàn)了 2.6x 的性能提升。

PS：也有一些需要注意的地方：

• 冷啟和樣本復(fù)用的問題。因為基于樣本的歷史 Response 生成草稿和預(yù)測長度，所以在初次 Rollout 時無法使用。極端情況，如果 Epoch 為 1，所有數(shù)據(jù)只會被使用 1 次，則不會有提升。
• 在優(yōu)化 Rollout 的過程中也需要避免陷入 GPU Util 高的誤區(qū)：

通過負載均衡確實可以降低 Bubble，提升 GPU Util；但是也要盡可能避免降低 Batch Size，導(dǎo)致 Memory Bound 問題加劇。

投機采樣確實可以提升算力利用率，不過如果 Token 接受率較低，將存在大量的無效計算。

• Batch Size 較小可能因為顯存不足（長序列尤其明顯）、時延要求較高等原因?qū)е?，而?Rollout 場景對單個樣本的時延要求不高，也就可以嘗試更多降低顯存開銷來提升 Batch Size 的方案，比如量化、PD 分離等。

三、引言

3.1 RL 相關(guān)問題

如下圖 Figure 2 所示，在一個 32B 模型的 RL 訓(xùn)練中，Code 和 Math 數(shù)據(jù)集下 Rollout 階段占比達到 84%-91%（最大序列長度為 16K），并且 Rollout 中存在明顯的 Bubble 問題；隨著最大序列長度的增加，該問題會更加明顯。

如下圖 Figure 3a 所示，隨著訓(xùn)練 Step 的增加，Response 長度會動態(tài)增加；如下圖 Figure 3b 所示，不同樣本的 Response 也會存在高度差異化。這些都進一步加劇負載不均的問題。如下圖 Figure 3c 所示，不同 GPU 的 SM Active 指標(biāo)出現(xiàn)了明顯的差異。

3.2 RL 性能的優(yōu)化

如上所示，RL 系統(tǒng)會存在 GPU 利用率低的問題，主要有 3 個原因：

• Rollout 階段占比很大，甚至可以達到 84%-91%，但是由于 LLM 的自回歸特性，Inference 的效率可能不高（Batch Size 不夠大時，Decoding 階段是明顯 Memory Bound）。
• 由于同一個 Batch 中的序列長短不一，會存在長尾效應(yīng)，導(dǎo)致出現(xiàn)明顯的負載不均問題。
• Training 和 Rollout 階段之間模型參數(shù)同步的效率可能不高。

也有一些常見的優(yōu)化方案：

• 通過截斷的方式降低長尾問題，但是也需要權(quán)衡精度和速度。
• 通過多 Batch，Continuous Batching 等方式實現(xiàn)更加均衡的調(diào)度。
• 使用一些常用的 LLM Inference 優(yōu)化方案加速，比如采用 FP8 執(zhí)行 Rollout，使用 Prefix Cache 等。
• 通過分布式 P2P 傳輸模型權(quán)重，充分利用節(jié)點互聯(lián)帶寬。

四、方案

4.1 RhymeRL 概覽

如下圖 Figure 5 所示，RhymeRL 延續(xù)了 HybridFlow（Verl）的混合 Controller 以及解耦 Rollout-Train 的架構(gòu)，通過 Rollout Worker 生成 Response，Reward Worker 計算獎勵，Train Worker 執(zhí)行 Policy 優(yōu)化。

為了提升 RL 系統(tǒng)的整體效率，作者采用了流式流水線架構(gòu)。在每個 Step 中：

• 各 Rollout Worker 從 sub-batch 隊列異步獲取 sub-batch 的 Prompt，并通過 Inference Engine 生成 Response（?）。
• 已完成的 Rollout Response 在傳輸?shù)?Train Worker 的 Reply Buffer 前，先由 Reward Worker 打分（?）。
• 當(dāng) Reply Buffer 積累的樣本達到 Batch Size 的要求時（?），Train Worker 執(zhí)行用戶定義的 RL 算法以優(yōu)化模型（?）。
• 完成 Full-Batch 的 Policy 優(yōu)化后，更新后的模型權(quán)重從 Train Worker 傳輸?shù)?Rollout Worker 的 Weight Buffer（主機內(nèi)存中）。在處理每個 sub-batch 前，Rollout Worker 將權(quán)重同步到 GPU（?）。這種權(quán)重更新策略可以消除全局同步開銷，最小化空閑等待時間。

為了實現(xiàn) Rollout 的加速，RhymeRL 做了幾個優(yōu)化：

• 引入投機采樣技術(shù)（HistoSpec），基于獎勵感知后綴樹方法，以最小開銷從歷史數(shù)據(jù)中高效生成草稿。
• 受 AIMD 啟發(fā)，通過動態(tài)調(diào)整投機 Token 的長度，從而在提升計算強度的同時獲得更高接受率。
• 為了實現(xiàn) Rollout Worker 間的負載均衡，RhymeRL 采用分布感知調(diào)度策略（HistoPipe），利用 Step 間的互補性實現(xiàn)整體工作負載平衡，并通過基于遷移的再平衡機制處理異常離群值。

4.2 HistoSpec

4.2.1 動機

RL 訓(xùn)練包含多個 Epoch，（論文里引用 DeepSeekMath，介紹為 50-100，不過隨著數(shù)據(jù)量的增加，比如更多的合成數(shù)據(jù)，Epoch 可能沒這么大）。在此期間，Prompt 會在不同 Epoch 內(nèi)被重復(fù)采樣；主流的 RL 算法也都會采用梯度裁剪來限制模型更新太快。此外，每個 Prompt 都會生成多組 Response（8-64），從而在每個周期內(nèi)實現(xiàn)多樣化 Reasoning 路徑的探索。因此，在相鄰的訓(xùn)練 Epoch 的 Rollout 輸出中，相同 Prompt 所生成的結(jié)果會呈現(xiàn)出高度相似性。

基于以上的相似性，可以將上一個 Epoch 的 Response 作為當(dāng)前 Rollout 投機采樣的草稿，以加速 Rollout 的生成速度。

• 如下圖 Figure 6a 所示，經(jīng)過 8 個 Epoch 后，數(shù)學(xué)任務(wù)中 93% 的 Token 可以通過此流程被成功接受。
• 如下圖 Figure 6b 也展示了 Respond 接受率的分布情況，在 Math-1 中接受率會更高一些。

4.2.2 基于樹結(jié)構(gòu)的歷史管理

但是基于歷史 Response 進行投機采樣的加速方法也會存在一些挑戰(zhàn)，主要是 2 個方面：

• 過長的輸出序列會導(dǎo)致顯著的前綴匹配開銷，并且歷史 Response 中的分支結(jié)構(gòu)會使草案 Token 選擇更加復(fù)雜。
• 被接受的 Token 長度不可預(yù)測，并且接受率越低，浪費的算力越多；而每次接受的 Token 太少，又可能獲得不了明顯的加速。

草稿生成開銷過大將會削弱投機采樣帶來的加速收益。為此，作者采用以下方案：

• 異步緩存構(gòu)建：RL 的工作模式可以放開對草稿生成的約束，可以將檢索相關(guān)計算轉(zhuǎn)移到索引階段。RhymeRL 采用 History Worker 異步索引歷史序列——當(dāng)生成新的 Response 時，Controller 將其分派給 History Worker 構(gòu)建索引。當(dāng)調(diào)度 Prompt 到 Rollout Worker 時，Controller 通知相應(yīng) History Worker 傳輸相應(yīng)構(gòu)建好的 Cache。
• 基于獎勵感知的樹管理：RhymeRL 采用后綴樹索引緩存 Response，實現(xiàn)對長度為 m 的 Prefix 進行 O(m) 時間復(fù)雜度的匹配。由于每個 Prompt 可能生成多個 Response，單個 Prefix 可能有多個不同的后綴分支，導(dǎo)致 Token 選擇的復(fù)雜化。作者觀察到，RL 算法會優(yōu)化模型使其以更高概覽生成高獎勵的解決方案，因此，作者在每個樹節(jié)點添加優(yōu)先級數(shù)值，其權(quán)重由該分支的獎勵總和決定。對于存在的每個后綴分支，HistoSpec 會選擇優(yōu)先級最高的分支。如下圖 Figure 7 所示，父節(jié)點的優(yōu)先級是所有子節(jié)點優(yōu)先級（獎勵）的總和。

4.2.3 基于樹結(jié)構(gòu)的歷史管理

為了應(yīng)對 Token 接受率與提升計算強度之間的 Tradeoff，作者借鑒了網(wǎng)絡(luò)擁塞控制采用的經(jīng)典方案 AIMD：

• HistoSpec 為每個 Response 設(shè)計動態(tài)草稿長度，初始值為 2 個 Token；當(dāng)所有 Token 都被接受時，將窗口長度增加 2，直到到達上限閾值（默認 32）。但如果有 Token 被拒絕，則直接將窗口重置為 2。
• HistoSpec 還為 Prefix 設(shè)置了動態(tài)長度，初始為 7，若未能找到匹配后綴，逐步縮減到 3。
• HistoSpec 還考慮了 Batch Size 的影響，在大 Batch Size 時（空閑算力少），縮短草稿長度；在小 Batch Size（空閑算力多），增加草稿長度。

4.3 HistoPipe

4.3.1 動機

歷史 Response 除了可以幫助生成草稿 Token，還可以幫助預(yù)測 Response 長度，具體來說，作者觀察到相同 Prompt 在相鄰 Epoch 周期中的 Response 長度分布相似，也就可以利用歷史長度數(shù)據(jù)對 Rollout 的 Prompt 長度進行排序，以便更好的負載均衡。

如下圖 Figure 9 所示，針對 5 個數(shù)據(jù)集在多個 Epoch 上分析了 Response 長度排序。具體而言，將 Response 長度分為 8 組（從低到高）。在 20 個 Epoch 中，數(shù)學(xué)任務(wù)平均僅有 16% 的 Response 會躍升到更高組別（代碼任務(wù)為 28%）。并且，其中數(shù)學(xué)任務(wù) 13% Response 僅在組邊界相鄰位置波動。

4.3.2 分布感知 Hybrid Pipeline

作者提出了分布感知的 Hybrid Pipeline 方案，該方案會在保持算法完整性的同時，通過連續(xù)訓(xùn)練 Step 間的 Worker 互補，構(gòu)建跨 Step 的協(xié)同負載均衡機制。

如下圖 Figure 10 所示，在奇數(shù)（1,3,5）Step 中，Scheduler 按執(zhí)行長度升序排列并分配給 Worker；在偶數(shù) Step 中按降序分配，通過這種交替互補方式，可以有效填充 Bubble，提升整體利用率。（PS：由于第一個 Epoch 缺乏歷史數(shù)據(jù)，因此在第一個 Epoch 不啟用）

然而，在真實的工作負載中仍會面臨一些挑戰(zhàn)：

• 極端情況非常長的 Response 會破壞這種互補性。
• Rollout 的長尾分布使得連續(xù) Step 無法實現(xiàn)完美的工作負載互補，依然會存在 Bubble。

4.3.3 基于遷移的重新負載

為了緩解超長 Rollout Response 對 Hybrid Pipeline 的影響，作者采用了兩種策略：

• Step 內(nèi)遷移——將過長的 Rollout 遷移到同 Step 內(nèi)其他仍在 Rollout 的 Group。
• 跨 Step 遷移——將異常值遷移到后續(xù)的 Step 中（PS：數(shù)學(xué)不等價，是否會影響效果）。

具體而言，每個排序組都設(shè)置一個閾值，當(dāng) Rollout 長度超過閾值時即觸發(fā)遷移。

• 對于 Step 內(nèi)遷移，會在執(zhí)行過程中動態(tài)地將長 Rollout 重新分配給其他組別，并且已經(jīng)生成的 Token 會保留，后續(xù)會重新計算相應(yīng)的 KV Cache（Prefill 階段）。
• 對于跨 Step 遷移，待遷移的樣本會加入下一個 Step 中，同樣會保留相應(yīng)已生成的 Token。

此外，作者也討論了上述閾值的設(shè)定問題，在 5 個數(shù)據(jù)集 20 個 Epoch 的實驗中，僅有 1.49%-3.55% 的 Response 需要遷移，證明該機制產(chǎn)生的開銷在可接受范圍，且對訓(xùn)練精度的影響可忽略不計。

4.3.4 雙層調(diào)度

實現(xiàn)近乎無 Bubble 的 Step 間互補性，需要各執(zhí)行組之間呈現(xiàn)近似線性的執(zhí)行時間分布。然而，在實際應(yīng)用中，長尾序列會導(dǎo)致執(zhí)行時間呈指數(shù)分布，如下圖 Figure 11a 所示，這會導(dǎo)致某些節(jié)點上依然存在 Bubble。為此，HistoPipe 提出雙層調(diào)度機制：

• 第一層：從 Prompt 到排序組，首先將已排序 Prompt 均勻分配到排序組。
• 第二層：將排序組映射到 GPU。若為每個排序組分配等量 GPU，由于長尾序列存在，其執(zhí)行時間會呈指數(shù)分布。因此，HistoPipe 設(shè)計了分布重組策略：不再平均分配 GPU，而是為短 Response 組和中長 Response 組分配較少 GPU，為少數(shù)長 Response 組分配額外 GPU，從而將執(zhí)行時間分布從指數(shù)型調(diào)整為線性，從而進一步減少 Bubble。如下圖 Figure 11b 所示。

五、實驗&結(jié)果

5.1 端到端訓(xùn)練評估

如下圖 Figure 13 所示，與 veRL（v0.4.1）和 AReaL（v0.3.0）相比，RhymeRL 在 8B-32B 模型，8K/16K 訓(xùn)練長度，DAPO/GRPO 算法上都獲得了不錯的加速。相比 veRL，在 8K 上平均加速 1.9x，16K 上平均加速 16K。

如下圖 Figure 14 為各種優(yōu)化手段相對提升的占比：

5.2 HistoSpec 消融實驗

如下圖 Figure 16 所示，隨著 Step 的增加，基于投機采樣的 RhymeRL 能獲得更明顯的加速：

如下圖 Figure 17 所示，隨著訓(xùn)練的持續(xù)進行，采用投機采樣生成的 Token 比例（ Spec. rate）逐漸上升。草稿 Token 中的接受率保持在 65% - 79% 區(qū)間，并且隨著訓(xùn)練進程同步提升：

5.3 HistoPipe 消融實驗

如下圖 Figure 15 所示，作者同樣驗證了 HistoPipe 中優(yōu)化手段的加速效果：

六、參考鏈接

1. ??https://arxiv.org/abs/2508.18588??

本文轉(zhuǎn)載自??AI閑談??，作者：AI閑談