字節(jié) RhythmRL：基于投機采樣+長度預測的 RL 加速

amei2000go

發(fā)布于 2025-9-22 07:05

瀏覽

0收藏

一、背景

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

對應的論文為：[2508.18588] History Rhymes: Accelerating LLM Reinforcement Learning with RhymeRL

二、摘要

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

RhymeRL 中，作者發(fā)現(xiàn)相鄰訓練 Epoch 中，Rollout 的 Response 表現(xiàn)出高度相似性。基于此，設計了用于加速 Rollout 的方案：

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

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

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

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

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

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

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

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

三、引言

3.1 RL 相關問題

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

字節(jié) RhythmRL：基于投機采樣+長度預測的 RL 加速-AI.x社區(qū)

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

字節(jié) RhythmRL：基于投機采樣+長度預測的 RL 加速-AI.x社區(qū)

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

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

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

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

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

四、方案

4.1 RhymeRL 概覽

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

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

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

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

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

字節(jié) RhythmRL：基于投機采樣+長度預測的 RL 加速-AI.x社區(qū)

4.2 HistoSpec

4.2.1 動機

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

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

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

字節(jié) RhythmRL：基于投機采樣+長度預測的 RL 加速-AI.x社區(qū)

4.2.2 基于樹結構的歷史管理

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

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

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

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

字節(jié) RhythmRL：基于投機采樣+長度預測的 RL 加速-AI.x社區(qū)

4.2.3 基于樹結構的歷史管理

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

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

4.3 HistoPipe

4.3.1 動機

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

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

字節(jié) RhythmRL：基于投機采樣+長度預測的 RL 加速-AI.x社區(qū)

4.3.2 分布感知 Hybrid Pipeline

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

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

字節(jié) RhythmRL：基于投機采樣+長度預測的 RL 加速-AI.x社區(qū)

然而，在真實的工作負載中仍會面臨一些挑戰(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ù)學不等價，是否會影響效果）。

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

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

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

4.3.4 雙層調度

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

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

字節(jié) RhythmRL：基于投機采樣+長度預測的 RL 加速-AI.x社區(qū)