当训练与推理的数值结果不一致,强化学习如何稳定?SkyRL 用「比特级一致」给出了答案
原文: IsoExec: Unified Execution to Eliminate Trainer-Inference Mismatch in SkyRL
SkyRL 团队提出 IsoExec,通过统一执行抽象和比特级一致的内核,消除了强化学习训练与推理引擎之间的数值漂移问题,让算法调试和硬件优化变得简单可靠。
- 强化学习中,训练引擎和推理引擎使用不同的代码和硬件优化,导致计算出的token概率存在微小差异,这种差异会累积并破坏训练稳定性。
- IsoExec 的核心思想是制定一个统一的「执行契约」,并确保所有内核在不同并行模式下都能产生比特级一致的结果。
- 在单个8卡H100节点上测试,IsoExec 将训练与推理的对数概率差异平均降低到百万分之一以下,同时只带来约25%的性能开销。
- 这项工作揭示了强化学习系统工程中一个常被忽视但至关重要的问题:数值一致性是算法稳定性的基石。
起因:一个被忽视的「数值幽灵」
在强化学习(RL)的训练流程中,有一个理论上完美但实践中充满陷阱的假设:用于生成样本(rollout)的推理策略,和用于更新模型参数的训练策略,是同一个。然而,现实中的RL系统通常由两个独立的引擎驱动:比如用vLLM或SGLang做高效推理,用Megatron或FSDP做大规模训练。这两个引擎为了各自优化,内部的计算图、内核实现、并行策略都不同。问题来了:浮点数运算是不满足结合律的,(a+b)+c 和 a+(b+c) 的结果可能在小数点后很多位都不同。这种微小的差异,在每一步的token概率计算中都会累积,最终可能让整个RL训练崩溃。
ByteDance的VeXact研究和Fireworks的案例都直接证明了这一点:训练和推理之间的微小数值差异(KL散度约0.013),就能导致重要性采样失效、奖励信号崩溃。这就像两个本应步调一致的士兵,因为步伐的微小差异,最终走散了。
拆解:IsoExec 如何「锁定」数值
SkyRL团队提出的IsoExec,不是要造一个更大的单一引擎,而是建立一个跨框架的「执行契约」。这个契约有两大支柱:
第一,统一执行契约。它明确规定了所有影响浮点舍入的执行细节,比如内核的实现、归约的顺序。这相当于给所有引擎(vLLM和Megatron)下发了一份必须遵守的「作战手册」,确保它们在计算同一个数学模型时,走完全相同的计算路径。
第二,比特级一致内核。团队重写或适配了关键内核,使其在张量并行、专家并行、序列并行等不同模式下,都能产生完全相同的二进制结果。他们甚至为一种名为Gated DeltaNet的线性注意力变体,开发了Chunkwise-Parallel Recurrent (CPR) 内核,确保在长序列的训练、预填充和递归解码阶段,数值结果都严格对齐。
趋势洞察:从「能跑」到「可靠跑」的系统工程觉醒
IsoExec的出现,标志着AI基础设施领域的一个深层趋势:随着模型和算法复杂度的提升,系统层面的数值确定性 正在从一个「nice-to-have」的特性,变成一项核心要求。过去,大家更关注速度和吞吐量。但现在,当算法创新(如新的RL方法)越来越依赖于精密的信号(如梯度、优势函数),任何系统的数值噪声都可能扭曲这些信号,导致实验无法复现、创新无法落地。
这就像建筑工程中,地基的微小沉降(数值漂移)在摩天大楼(复杂RL训练)上会被放大成严重的结构问题。IsoExec这类工作,就是在为AI的「摩天大楼」夯实数值地基。
实用价值:对开发者意味着什么?
- 算法开发者:你可以更自信地调试新RL算法。当训练出问题时,可以首先排除系统数值不一致的干扰,聚焦算法本身。这将极大降低创新成本。
- 系统/框架开发者:这是一个优秀的工程范式。它展示了如何在不牺牲性能的前提下,通过严格的契约和内核对齐,实现跨系统的确定性。这为未来构建更可靠的分布式训练系统提供了思路。
- 企业用户:如果你在自研或深度定制RL训练流程(比如机器人控制、复杂游戏AI),关注类似IsoExec的解决方案,能显著提升训练的成功率和稳定性,减少因「玄学」问题浪费的算力。
反常识/意外
一个值得注意的点是,IsoExec在实现比特级一致性的同时,只带来了约25%的性能开销。这打破了「追求确定性必然以大幅牺牲性能为代价」的固有认知。通过巧妙的契约设计和内核优化,SkyRL团队证明了可靠性和性能可以兼得。此外,这项工作再次强调,AI的突破不仅来自算法和模型,也来自让这些算法和模型稳定运行的底层系统工程。当大家追逐更大的模型时,像IsoExec这样确保它们「能稳定训练」的工作,可能才是解锁下一轮创新的关键。
原文地址: IsoExec: Unified Execution to Eliminate Trainer-Inference Mismatch in SkyRL
分析由 BitByAI 生成 · 阅读原文