跳到主要内容
返回时间线
arXiv来源发表:

Nereus 让 RL 后训练在运行中自适应并行:真实数据轨迹上平均步延迟降 27.7%,8B PPO 吞吐最高达 OpenRLHF 的 7.27 倍

核心概要

Nereus 是一个面向大语言模型强化学习后训练的成本感知运行时,它把每个模型-阶段副本抽象为 Elastic Model Unit,用全局转移图编排 Split、Merge、Extend、Destroy 四种原语,在运行中按“收益能否偿还转移成本”的准则在线调整 DP/TP/PP 执行计划;在真实数据构建的轨迹上,在线 TP/PP 自适应相对初始固定 TP/PP 布局把平均步延迟降低 27.7%,1,000 步扩展到 1,024 GPU 的运行中六次转移仅占总运行时间的 0.079%,8B PPO 端到端吞吐相对 OpenRLHF 提升 2.14–7.27 倍、相对 Verl 提升 1.10–1.47 倍。

AI-generated editorial illustration: Nereus: Adaptive Parallelism for LLM Post-Training

深度剖析

提出低开销、成本感知的自适应策略:控制器监控序列长度、可用 GPU、峰值内存与实测计算/通信效率,规划器在内存可行约束下选出预测稳态步延迟最低的全局计划,再按“节省能否在触发信号下次越阈前偿还转移成本”的回收期准则决定是否执行转移。 已有 RL 后训练框架在启动后固定分配与并行选择,弹性训练系统通常只管理单个模型,检查点式重分片需要重启并重建状态;Nereus 把回收期准入原则(源自 Pollux、Sia 的调度思路)应用到单个耦合的 RL 作业上,并区分紧急路径(当前计划不可行时豁免盈利性)与机会路径。 在 Cluster #2 上,成本模型的训练阶段延迟项在 72 个跨 8–256 GPU、四个序列长度桶的样本外计划上平均绝对百分比误差 4.96%、最大误差 14.61%;决策开销从 32 GPU 的 0.17 ms 到 1,024 GPU 的 338 ms,而基于 SCIP 的求解器从 3.4 ms 增长到 511 s,重规划峰值内存低于 100 MB。

提出与依赖对齐的状态抽象 Elastic Model Unit(EMU):一个 EMU 是一个模型-阶段副本,把紧耦合的 TP/PP 布局与状态封装在单元内部,把松耦合的 DP 副本暴露在外,DP 度即该模型-阶段活跃 EMU 的数量;在线自适应被归约为 Split、Merge、Extend、Destroy 四种原语。 粗粒度边界(检查点/整作业)会搬移无关状态,细粒度分片边界把 TP/PP 重分片决策暴露给规划器;EMU 选择模型-阶段这一中间粒度,使规划器以整单元推理、由原语处理底层分片,同时保持训练语义。 在 1632 GPU 的同一 8B actor/critic 负载上,资源扩展成本为检查点/重启 836.74 s、分片级 66.43 s、EMU 6.52 s;Extend 在 48 到 3264 GPU 上耗时 2.45–8.48 s,比 Oobleck 快 3.8–9.9 倍、比 Gemini 快 6.7–10.8 倍、比 Tenplex 快 9.3–16.2 倍、比 UCP 快 115.7–284.1 倍;70B 模型 64 GPU 上重分片成本相对 MCP 降低 99.1%、相对 Tenplex 降低 96.7%。

提出安全的并发转移编排:把全局计划变化编译为 EMU 原语的有向无环图,用“先释放后获取”和阶段顺序两条确定性规则插入跨模型-阶段资源依赖边,避免共享 GPU 上的死锁,并让就绪原语跨 GPU 组与流并发执行。 此前按组件独立迁移的做法不处理共享 GPU 的瞬时归属冲突;Nereus 只在瞬时 GPU 重叠阻塞获取时添加跨阶段边,并在被阻塞时选择可执行的 Destroy 释放所需 GPU,从而保持 DAG 无环、无循环等待。 三类转移各 100 次随机启动时序与 GPU 分配的试验中,Nereus 插入 1–2 条资源依赖边并全部成功,DynaRL 式逐组件迁移成功率仅 34–62%,失败均为更长超时也无法解决的共享 GPU 死锁;转移 DAG 规划在 1,024 GPU 上耗时 1.22 ms,而 SCIP 需 120.10 s,转移完成时间相对 SCIP 的差距在各规模内不超过 6.9%。

端到端验证:在三个跨厂商、跨互连、跨规模的 GPU 集群上,Nereus 提升 8B PPO 吞吐并保持训练行为,且收益延伸到 ReMax、GRPO 与异步执行。 相对 OpenRLHF 与 Verl 的固定启动配置,Nereus 在同一 GPU 预算内随序列长度与阶段瓶颈变化自适应;相对 DynaRL 式准入,它额外要求节省偿还转移成本。 8B PPO 吞吐相对 OpenRLHF 提升 2.14–7.27 倍(中位 3.99)、相对 Verl 提升 1.10–1.47 倍(中位 1.21);18 个测量设置中选中计划全部落在经验最优的 5% 以内,61.1% 的设置精确匹配;三条留出轨迹上平均 858.7 s/步,对比 DynaRL 式准入 928.3 s/步;ReMax 与 GRPO 上延迟最多降低 13.3% 与 15.8%,异步设置相对 Laminar 最多降低 37.3%;前 50 步奖励与 PPO KL 估计与 Verl 接近,墙钟时间减少 13.9%。

启示与展望

这项工作面向在 GPU 集群上运行 RL 后训练(PPO、ReMax、GRPO 等)的工程与系统团队,适用于资源供给、序列长度、内存压力与阶段瓶颈会在一次运行内变化的场景;自适应发生在安全边界上,即同步执行的 RL 步完成处与异步执行的权重同步处。它使运行中的作业可以在不重建完整作业状态的前提下改变 DP/TP/PP 布局,并在节点突然失效时从完好单元重新规划,同时与容错系统互补。作者指出,模型-阶段边界与模块化运行时为进一步扩展到上下文/专家并行、外部工具服务自动扩缩以及更多 RL 框架提供了路径。

成本模型的效率按操作类别(如 GEMM、注意力、某类链路上的 all-reduce)存储,未见类别以保守值起步再逐步修正,因此新模型或新硬件上的早期决策质量仍是一个开放问题;计划空间限定为模型-阶段内同构的 TP/PP 布局与 2 的幂次并行度,作者也把上下文/专家并行列为后续扩展方向。评估以 Llama-3.1-8B 的 PPO 为主,更大模型与更多算法组合下的表现仍需更多验证;此外,本证据包为论文全文,图表以文字描述形式出现,具体曲线细节无法在此逐点核对。

来源