一、背景

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

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

二、摘要

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

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

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

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

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

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

• 冷啟和樣本復(fù)用的問(wèn)題。因?yàn)榛跇颖镜臍v史 Response 生成草稿和預(yù)測(cè)長(zhǎng)度，所以在初次 Rollout 時(shí)無(wú)法使用。極端情況，如果 Epoch 為 1，所有數(shù)據(jù)只會(huì)被使用 1 次，則不會(huì)有提升。
• 在優(yōu)化 Rollout 的過(guò)程中也需要避免陷入 GPU Util 高的誤區(qū)：

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

投機(jī)采樣確實(shí)可以提升算力利用率，不過(guò)如果 Token 接受率較低，將存在大量的無(wú)效計(jì)算。

• Batch Size 較小可能因?yàn)轱@存不足（長(zhǎng)序列尤其明顯）、時(shí)延要求較高等原因?qū)е拢?Rollout 場(chǎng)景對(duì)單個(gè)樣本的時(shí)延要求不高，也就可以嘗試更多降低顯存開銷來(lái)提升 Batch Size 的方案，比如量化、PD 分離等。

三、引言

3.1 RL 相關(guān)問(wèn)題

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

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

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

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

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

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

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

四、方案

4.1 RhymeRL 概覽

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

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

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

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

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

4.2 HistoSpec

4.2.1 動(dòng)機(jī)

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

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

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

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

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

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

草稿生成開銷過(guò)大將會(huì)削弱投機(jī)采樣帶來(lái)的加速收益。為此，作者采用以下方案：

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

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

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

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

4.3 HistoPipe

4.3.1 動(dòng)機(jī)

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

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

4.3.2 分布感知 Hybrid Pipeline

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

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

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

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

4.3.3 基于遷移的重新負(fù)載

為了緩解超長(zhǎng) Rollout Response 對(duì) Hybrid Pipeline 的影響，作者采用了兩種策略：

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

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

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

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

4.3.4 雙層調(diào)度

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

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

五、實(shí)驗(yàn)&結(jié)果

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

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

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

5.2 HistoSpec 消融實(shí)驗(yàn)

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

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

5.3 HistoPipe 消融實(shí)驗(yàn)

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

六、參考鏈接

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

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