[Paper] RollMux：相位级多路复用用于去聚合的RL后训练

Source: arXiv - 2512.11306v1

概览

RollMux 解决了随着强化学习（RL）工作负载向分离架构迁移而出现的瓶颈——即 rollout（数据生成）阶段和 training（模型更新）阶段使用不同集群。由于 on‑policy 算法需要这两个阶段严格同步，一个集群常常在另一个忙碌时处于空闲状态，导致昂贵的 GPU 资源被浪费。RollMux 引入了一种跨集群调度框架，能够“填补”这些空闲时段，在生产规模的 GPU 测试平台上实现最高 1.8× 的成本效率提升。

主要贡献

• 协同执行组抽象（Co‑execution Group Abstraction）： 将整体硬件池划分为相互隔离的局部域，使作业能够共享资源而不相互影响内存占用。
• 两层调度器（Two‑Tier Scheduler）：
  • 组间调度器（Inter‑group scheduler） 使用保守的随机规划决定每个 RL 作业（rollout 与 training）在各组中的放置位置。
  • 组内调度器（Intra‑group scheduler） 实现了可证明最优的轮转（round‑robin）方案，在每个组内部最大化 GPU 利用率。
• 热启动上下文切换（Warm‑Star Context Switching）： 强制驻留约束，使大型模型状态保持在主机内存中缓存，实现 rollout 与 training 阶段之间的近乎瞬时切换。
• 生产规模评估（Production‑Scale Evaluation）： 在 656 GPU（328 H20 + 328 H800）集群上展示了相较于普通分离方案提升 1.84× 成本效率、相较于最佳共置基线提升 1.38×，且 100 % 达到服务等级目标（SLO）合规性。

方法论

1. 问题建模（Problem Modeling）： 作者将 RL 流水线建模为两个交替出现、资源占用巨大的阶段（rollout = 内存受限，training = 计算受限），且必须保持同步。
2. 组形成（Group Formation）： 将硬件池划分为协同执行组——每组包含一组 GPU 与相应的主机内存，可为单个作业的整个生命周期保留资源。这样可以隔离庞大的模型状态，避免昂贵的数据迁移。
3. 组间调度（Inter‑Group Scheduling）： 随机规划器评估候选组在每个阶段会产生的预期空闲时间（“气泡”），并将作业分配到能够最小化整体气泡成本的组。规划器是保守的：它倾向于选择能够在工作负载波动下仍保证 SLO 的放置方案。
4. 组内调度（Intra‑Group Scheduling）： 在组内部，RollMux 采用轮转调度交替执行不同作业的 rollout 与 training 任务，有效实现 GPU 的“复用”。作者证明，在固定组大小和驻留约束下，该调度方案能够最大化利用率。
5. 实现与集成（Implementation & Integration）： 该框架可挂接到现有的 RL 编排栈（如 Ray RLlib），利用标准容器运行时，仅需一个轻量级守护进程来强制组边界并执行调度决策。

结果与发现

关键要点

• 通过将一个阶段的空闲“气泡”与另一个阶段的活跃阶段重叠，RollMux 消除了困扰 on‑policy RL 流水线的大部分死等时间。
• 驻留约束使模型保持在主机内存中，将上下文切换开销降低了超过三倍。
• 即使在满负载（全 656 GPU 集群）下，调度器仍能保持确定性的 SLO 保证，这对生产级 RL 服务至关重要。

实际意义

• 降低云费用： 运行大规模 RL（如机器人、推荐系统、自动驾驶仿真）的公司可以在相同 GPU 开支下实现近乎翻倍的吞吐量。
• 简化集群运维： 组抽象让运维团队为每个 RL 作业分配固定的“槽位”，避免了临时的内存固定技巧，降低了 OOM（内存溢出）崩溃的风险。
• 加速实验迭代： 热启动上下文切换意味着开发者可以在不等待漫长数据生成阶段完成的情况下快速进行策略更新，提升研发到生产的周期。
• 兼容性： RollMux 作为流行 RL 框架的插件使用，现有代码库只需进行最小改动——主要是配置组大小和驻留策略。

局限性与未来工作

• 仅针对 On‑Policy： 当前设计假设严格的 rollout‑training 同步；off‑policy 或异步 RL 算法的收益可能有限。
• 静态组大小： 组在作业启动时即确定；尚未支持在作业负载突增时动态扩容。
• 硬件多样性： 评估在同质的 NVIDIA H20/H800 GPU 上完成；异构加速器（TPU、AMD GPU）可能带来新的调度挑战。
• 未来方向： 将随机规划器扩展至异构资源、支持动态组重新划分，并探索其在其他流水线式工作负载（如视频转码、大规模数据预处理）中的适用性。

作者

• 吴天元
• 曹伦希
• 魏一宁
• 高伟
• 赵宇恒
• 安大凯
• 熊少攀
• 吕志强
• 黄菊
• 杨思然
• 于英浩
• 王佳莹
• 曲林
• 王伟

论文信息

• arXiv ID: 2512.11306v1
• 分类: cs.DC
• 发布时间: 2025 年 12 月 12 日
• PDF: Download PDF