字节 RhythmRL：基于投机采样+长度预测的 RL 加速

amei2000go

发布于 2025-9-22 07:05

浏览

0收藏

一、背景

最近一直在关注 RL Infra 相关的工作，尤其是 RL 性能优化，后续会逐渐介绍一下该领域的相关文章，本文先简单介绍一下字节新发布的 RhymeRL。

对应的论文为：[2508.18588] History Rhymes: Accelerating LLM Reinforcement Learning with RhymeRL

二、摘要

RL 成为提升 LLM Reasoning 能力的关键方法，与传统预训练不同，RL 包含多个阶段：Rollout、Reward、Training，需要多种类型的 Worker 协同配合；除此之外，为了效率也可能引入异步训练方式，需要考虑效率与效果的 Tradeoff。然而，当前 RL 系统面临 GPU 利用率低的问题，主要原因为：Rollout 在整个 RL 过程中占比很高，并且 Rollout 的过程中很容易出现负载不均，导致 GPU Bubble。而异步训练、截断等方式虽然可以缓解该问题，但是可能牺牲准确率。

RhymeRL 中，作者发现相邻训练 Epoch 中，Rollout 的 Response 表现出高度相似性。基于此，设计了用于加速 Rollout 的方案：

HistoSpec，基于相似性，用上一个 Epoch 的 Response 作为草稿，使用投机采样技术加速 Rollout 生成速度。
HistoPipe，基于相似性，用上一个 Epoch 的 Response 的长度分布来预测当前 Epoch 的 Response 长度，进而实现负载均衡。

此外，作者也解决了其中的相关问题，比如动态调整投机采样的草稿长度，以避免引入过多无效计算；以及引入双层调度策略，进一步实现负载的均衡。

作者在真实生产环境中进行评估，证明其具备从数十到数千 GPU 的扩展能力。实验结果表明，RhymeRL 在保持准确性的前提下，相较于现有方法实现了 2.6x 的性能提升。

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

冷启和样本复用的问题。因为基于样本的历史 Response 生成草稿和预测长度，所以在初次 Rollout 时无法使用。极端情况，如果 Epoch 为 1，所有数据只会被使用 1 次，则不会有提升。
在优化 Rollout 的过程中也需要避免陷入 GPU Util 高的误区：

通过负载均衡确实可以降低 Bubble，提升 GPU Util；但是也要尽可能避免降低 Batch Size，导致 Memory Bound 问题加剧。

投机采样确实可以提升算力利用率，不过如果 Token 接受率较低，将存在大量的无效计算。

Batch Size 较小可能因为显存不足（长序列尤其明显）、时延要求较高等原因导致，而在 Rollout 场景对单个样本的时延要求不高，也就可以尝试更多降低显存开销来提升 Batch Size 的方案，比如量化、PD 分离等。

三、引言

3.1 RL 相关问题

如下图 Figure 2 所示，在一个 32B 模型的 RL 训练中，Code 和 Math 数据集下 Rollout 阶段占比达到 84%-91%（最大序列长度为 16K），并且 Rollout 中存在明显的 Bubble 问题；随着最大序列长度的增加，该问题会更加明显。

字节 RhythmRL：基于投机采样+长度预测的 RL 加速-AI.x社区

如下图 Figure 3a 所示，随着训练 Step 的增加，Response 长度会动态增加；如下图 Figure 3b 所示，不同样本的 Response 也会存在高度差异化。这些都进一步加剧负载不均的问题。如下图 Figure 3c 所示，不同 GPU 的 SM Active 指标出现了明显的差异。

字节 RhythmRL：基于投机采样+长度预测的 RL 加速-AI.x社区

3.2 RL 性能的优化

如上所示，RL 系统会存在 GPU 利用率低的问题，主要有 3 个原因：

Rollout 阶段占比很大，甚至可以达到 84%-91%，但是由于 LLM 的自回归特性，Inference 的效率可能不高（Batch Size 不够大时，Decoding 阶段是明显 Memory Bound）。
由于同一个 Batch 中的序列长短不一，会存在长尾效应，导致出现明显的负载不均问题。
Training 和 Rollout 阶段之间模型参数同步的效率可能不高。

也有一些常见的优化方案：

通过截断的方式降低长尾问题，但是也需要权衡精度和速度。
通过多 Batch，Continuous Batching 等方式实现更加均衡的调度。
使用一些常用的 LLM Inference 优化方案加速，比如采用 FP8 执行 Rollout，使用 Prefix Cache 等。
通过分布式 P2P 传输模型权重，充分利用节点互联带宽。

四、方案

4.1 RhymeRL 概览

如下图 Figure 5 所示，RhymeRL 延续了 HybridFlow（Verl）的混合 Controller 以及解耦 Rollout-Train 的架构，通过 Rollout Worker 生成 Response，Reward Worker 计算奖励，Train Worker 执行 Policy 优化。

为了提升 RL 系统的整体效率，作者采用了流式流水线架构。在每个 Step 中：

各 Rollout Worker 从 sub-batch 队列异步获取 sub-batch 的 Prompt，并通过 Inference Engine 生成 Response（❷）。
已完成的 Rollout Response 在传输到 Train Worker 的 Reply Buffer 前，先由 Reward Worker 打分（❸）。
当 Reply Buffer 积累的样本达到 Batch Size 的要求时（❹），Train Worker 执行用户定义的 RL 算法以优化模型（❺）。
完成 Full-Batch 的 Policy 优化后，更新后的模型权重从 Train Worker 传输到 Rollout Worker 的 Weight Buffer（主机内存中）。在处理每个 sub-batch 前，Rollout Worker 将权重同步到 GPU（❻）。这种权重更新策略可以消除全局同步开销，最小化空闲等待时间。

为了实现 Rollout 的加速，RhymeRL 做了几个优化：

引入投机采样技术（HistoSpec），基于奖励感知后缀树方法，以最小开销从历史数据中高效生成草稿。
受 AIMD 启发，通过动态调整投机 Token 的长度，从而在提升计算强度的同时获得更高接受率。
为了实现 Rollout Worker 间的负载均衡，RhymeRL 采用分布感知调度策略（HistoPipe），利用 Step 间的互补性实现整体工作负载平衡，并通过基于迁移的再平衡机制处理异常离群值。

字节 RhythmRL：基于投机采样+长度预测的 RL 加速-AI.x社区

4.2 HistoSpec

4.2.1 动机

RL 训练包含多个 Epoch，（论文里引用 DeepSeekMath，介绍为 50-100，不过随着数据量的增加，比如更多的合成数据，Epoch 可能没这么大）。在此期间，Prompt 会在不同 Epoch 内被重复采样；主流的 RL 算法也都会采用梯度裁剪来限制模型更新太快。此外，每个 Prompt 都会生成多组 Response（8-64），从而在每个周期内实现多样化 Reasoning 路径的探索。因此，在相邻的训练 Epoch 的 Rollout 输出中，相同 Prompt 所生成的结果会呈现出高度相似性。

基于以上的相似性，可以将上一个 Epoch 的 Response 作为当前 Rollout 投机采样的草稿，以加速 Rollout 的生成速度。

如下图 Figure 6a 所示，经过 8 个 Epoch 后，数学任务中 93% 的 Token 可以通过此流程被成功接受。
如下图 Figure 6b 也展示了 Respond 接受率的分布情况，在 Math-1 中接受率会更高一些。

字节 RhythmRL：基于投机采样+长度预测的 RL 加速-AI.x社区

4.2.2 基于树结构的历史管理

但是基于历史 Response 进行投机采样的加速方法也会存在一些挑战，主要是 2 个方面：

过长的输出序列会导致显著的前缀匹配开销，并且历史 Response 中的分支结构会使草案 Token 选择更加复杂。
被接受的 Token 长度不可预测，并且接受率越低，浪费的算力越多；而每次接受的 Token 太少，又可能获得不了明显的加速。

草稿生成开销过大将会削弱投机采样带来的加速收益。为此，作者采用以下方案：

异步缓存构建：RL 的工作模式可以放开对草稿生成的约束，可以将检索相关计算转移到索引阶段。RhymeRL 采用 History Worker 异步索引历史序列——当生成新的 Response 时，Controller 将其分派给 History Worker 构建索引。当调度 Prompt 到 Rollout Worker 时，Controller 通知相应 History Worker 传输相应构建好的 Cache。
基于奖励感知的树管理：RhymeRL 采用后缀树索引缓存 Response，实现对长度为 m 的 Prefix 进行 O(m) 时间复杂度的匹配。由于每个 Prompt 可能生成多个 Response，单个 Prefix 可能有多个不同的后缀分支，导致 Token 选择的复杂化。作者观察到，RL 算法会优化模型使其以更高概览生成高奖励的解决方案，因此，作者在每个树节点添加优先级数值，其权重由该分支的奖励总和决定。对于存在的每个后缀分支，HistoSpec 会选择优先级最高的分支。如下图 Figure 7 所示，父节点的优先级是所有子节点优先级（奖励）的总和。

字节 RhythmRL：基于投机采样+长度预测的 RL 加速-AI.x社区

4.2.3 基于树结构的历史管理

为了应对 Token 接受率与提升计算强度之间的 Tradeoff，作者借鉴了网络拥塞控制采用的经典方案 AIMD：

HistoSpec 为每个 Response 设计动态草稿长度，初始值为 2 个 Token；当所有 Token 都被接受时，将窗口长度增加 2，直到到达上限阈值（默认 32）。但如果有 Token 被拒绝，则直接将窗口重置为 2。
HistoSpec 还为 Prefix 设置了动态长度，初始为 7，若未能找到匹配后缀，逐步缩减到 3。
HistoSpec 还考虑了 Batch Size 的影响，在大 Batch Size 时（空闲算力少），缩短草稿长度；在小 Batch Size（空闲算力多），增加草稿长度。

4.3 HistoPipe

4.3.1 动机

历史 Response 除了可以帮助生成草稿 Token，还可以帮助预测 Response 长度，具体来说，作者观察到相同 Prompt 在相邻 Epoch 周期中的 Response 长度分布相似，也就可以利用历史长度数据对 Rollout 的 Prompt 长度进行排序，以便更好的负载均衡。

如下图 Figure 9 所示，针对 5 个数据集在多个 Epoch 上分析了 Response 长度排序。具体而言，将 Response 长度分为 8 组（从低到高）。在 20 个 Epoch 中，数学任务平均仅有 16% 的 Response 会跃升到更高组别（代码任务为 28%）。并且，其中数学任务 13% Response 仅在组边界相邻位置波动。

字节 RhythmRL：基于投机采样+长度预测的 RL 加速-AI.x社区

4.3.2 分布感知 Hybrid Pipeline

作者提出了分布感知的 Hybrid Pipeline 方案，该方案会在保持算法完整性的同时，通过连续训练 Step 间的 Worker 互补，构建跨 Step 的协同负载均衡机制。

如下图 Figure 10 所示，在奇数（1,3,5）Step 中，Scheduler 按执行长度升序排列并分配给 Worker；在偶数 Step 中按降序分配，通过这种交替互补方式，可以有效填充 Bubble，提升整体利用率。（PS：由于第一个 Epoch 缺乏历史数据，因此在第一个 Epoch 不启用）

字节 RhythmRL：基于投机采样+长度预测的 RL 加速-AI.x社区

然而，在真实的工作负载中仍会面临一些挑战：

极端情况非常长的 Response 会破坏这种互补性。
Rollout 的长尾分布使得连续 Step 无法实现完美的工作负载互补，依然会存在 Bubble。

4.3.3 基于迁移的重新负载

为了缓解超长 Rollout Response 对 Hybrid Pipeline 的影响，作者采用了两种策略：

Step 内迁移——将过长的 Rollout 迁移到同 Step 内其他仍在 Rollout 的 Group。
跨 Step 迁移——将异常值迁移到后续的 Step 中（PS：数学不等价，是否会影响效果）。

具体而言，每个排序组都设置一个阈值，当 Rollout 长度超过阈值时即触发迁移。

对于 Step 内迁移，会在执行过程中动态地将长 Rollout 重新分配给其他组别，并且已经生成的 Token 会保留，后续会重新计算相应的 KV Cache（Prefill 阶段）。
对于跨 Step 迁移，待迁移的样本会加入下一个 Step 中，同样会保留相应已生成的 Token。

此外，作者也讨论了上述阈值的设定问题，在 5 个数据集 20 个 Epoch 的实验中，仅有 1.49%-3.55% 的 Response 需要迁移，证明该机制产生的开销在可接受范围，且对训练精度的影响可忽略不计。

4.3.4 双层调度

实现近乎无 Bubble 的 Step 间互补性，需要各执行组之间呈现近似线性的执行时间分布。然而，在实际应用中，长尾序列会导致执行时间呈指数分布，如下图 Figure 11a 所示，这会导致某些节点上依然存在 Bubble。为此，HistoPipe 提出双层调度机制：

第一层：从 Prompt 到排序组，首先将已排序 Prompt 均匀分配到排序组。
第二层：将排序组映射到 GPU。若为每个排序组分配等量 GPU，由于长尾序列存在，其执行时间会呈指数分布。因此，HistoPipe 设计了分布重组策略：不再平均分配 GPU，而是为短 Response 组和中长 Response 组分配较少 GPU，为少数长 Response 组分配额外 GPU，从而将执行时间分布从指数型调整为线性，从而进一步减少 Bubble。如下图 Figure 11b 所示。

字节 RhythmRL：基于投机采样+长度预测的 RL 加速-AI.x社区