Toward RL on Mixed-Precision, Mixed-Accelerator Samplers

发表时间: 2026-10 · Blog post by Nan Jiang (modal.com)

原文: https://modal.com/blog/toward-rl-on-mixed-precision-mixed-accelerator-samplers

机构:Modal Labs


速读

一句话结论 论文提出结合掩码重要性采样、分数中心化与掩码重放 top-p 采样的离策略修正方案,使强化学习得以在跨 6 种加速器芯片及 NVFP4、FP8、BF16 混合精度的异构集群上完全异步稳定训练,在长上下文代码智能体任务上取得了与全同构硬件基线持平的通过率,并将采样成本降低了约 15%。

要解决什么问题 在大模型强化学习训练中,推理采样环节占用了 80%–85% 的加速卡算力。传统强化学习为了保证数值稳定,通常强制采样器与训练器使用相同代际的硬件、相同的 BF16 精度以及完全对齐的计算内核。如果直接在采样端引入服务系统常用的提速降本手段——如 NVFP4 或 FP8 低精度量化、跨架构硬件及异步执行——会导致采样策略与训练策略之间产生显著的训练-推理不匹配。这种不匹配会在梯度更新中引入两大致命偏差:一是样本权重错误,梯度被采样频率而非训练策略概率加权;二是策略漂移,对数概率梯度在采样分布下的期望不再为零,导致模型无论生成质量好坏都会持续朝采样策略的方向漂移。在没有修正的情况下,这两种偏差会相互放大,造成采样器与训练器之间的 KL 散度激增 3 倍以上、梯度范数暴涨 20 倍,最终引发训练崩溃,表现为模型在多轮交互中突然停止调用工具、产生大量空回复或格式错误。

怎么做的 核心思路是放弃在底层追求位级一致的完全在策略模拟,将计算内核、数值精度、异构芯片架构以及完全异步带来的策略陈旧度,统一抽象为采样策略 $q$ 与训练策略 $p_\theta$ 之间的概率分布差异,通过逐 token 级别的离策略数学修正将该差异约束在有界范围内,从而彻底解耦采样系统与训练系统。整套稳定机制由三个协同部件构成:首先是掩码重要性采样(MIS),它用重要性比率重新加权每个 token 的梯度项,并对超出固定范围 $[\beta_{\min}, \beta_{\max}]$ 的异常比率直接截断置零,防止采样端极少出现的稀有 token 产生巨大权重而主导梯度更新:

$$\begin{aligned} w_t = \begin{cases} \frac{p_\theta(x_t \mid x_{<t})}{q(x_t \mid x_{<t})}, & \text{若 } \frac{p_\theta(x_t \mid x_{<t})}{q(x_t \mid x_{<t})} \in [\beta_{\min}, \beta_{\max}] \\ 0, & \text{其他} \end{cases} \end{aligned}$$

然而,单纯掩码截断会打破恒等变换,重新引入策略漂移。因此第二项设计是分数中心化(SC),它在每个生成位置减去该位置在采样分布 $q$ 下对数概率梯度的期望值,其定义式为:

$$\widetilde{\nabla}_\theta \log p_\theta(x_t \mid x_{<t}) = \nabla_\theta \log p_\theta(x_t \mid x_{<t}) - \sum_{x} q(x \mid x_{<t}) \nabla_\theta \log p_\theta(x \mid x_{<t})$$</div> 这一操作消除了常数优势函数项带来的错误参数位移,使更新在期望上仅受 token 间的优势差驱动。由于计算该期望需要全词表的采样概率,采样器在记录高概率候选之余,未记录的尾部概率通过训练器概率缩放补全。为了彻底消除尾部估计误差,第三项设计引入带掩码重放的 top-p 采样,采样器仅在 top-p 候选集合内生成,训练器在相同候选掩码上重新归一化对数概率,强制使两者的动作空间支撑集完全一致。这不仅杜绝了产生极端比率的极低概率词,还让分数中心化的期望计算变为对候选集的精确求和,最终使异构采样在数学期望上逼近真实的在策略更新。

效果如何 实验使用 35B 参数的混合专家模型 Qwen3.6-35B-A3B,在 2.7K 个清洗后的仓库级代码任务上进行最高 256K tokens 上下文的多轮智能体强化学习,每次更新采用 128 个 prompt 及每个 prompt 8 条轨迹。训练器基于 Megatron-LM 在 B200 GPU 上以 BF16 精度更新,采样器基于 SGLang 部署在 Modal 无服务器平台上,进行不限陈旧度的完全异步训练。对比基线包括:仅在 B200 上运行 BF16 权重与 BF16 KV cache 的同构采样器路线,以及混合采样器上未做修正的 GRPO(代表无修正路线)、仅加 MIS、仅加 SC、GRPO 搭配 top-p 等不同消融配置。混合采样器跨越 6 种芯片(B200 和 B300 运行 NVFP4 权重,H100 和 H200 运行 FP8 与 BF16,A100 和 RTX PRO 6000 运行 BF16,KV cache 统一为 FP8)。评测在无重叠的 SWE-bench Pro V2 全集及 HARD-51 困难子集上以 BF16 进行。实验显示,所有未做完整修正的基线均发生训练崩溃;而采用 SC + MIS + top-p 的混合采样器在平均 KL 散度达到同构基线 4.7 倍的情况下依然保持极高稳定性,在 HARD-51 上的 pass@4 稳定在 65%–75%,全集上 pass@4 从 89% 提升至约 94%,最终成绩完全持平于全 B200 BF16 同构基线,且采样单步成本降低了约 15%。作者指出的局限包括:方法基于逐 token 的一阶近似,未包含序列级校正,仍继承了前缀由旧策略生成的近似误差;带掩码重放的 top-p 会导致截断范围外的词元无法获得梯度;此外离策略修正仅改变梯度利用方式,无法改变低精度采样器所采样轨迹本身的质量,选择硬件和精度本质上仍是采样吞吐量与样本效率之间的权衡。

主要贡献

在超大模型强化学习(Reinforcement Learning, RL)的端到端系统中,采样阶段(Rollout/Sampling)通常消耗整个集群 80% 至 85% 的 GPU 算力。为了保证训练的数值稳定性,传统方法通常强制要求采样器与训练器运行在完全相同的加速器型号上,甚至采用相同的计算精度(如统一 BF16)。这一约束使得 RL 采样无法享受主流模型推理服务(Serving)领域中成熟的性能与成本优化红利,例如 FP8、MXFP4、NVFP4 等超低精度量化格式、针对特定硬件协同设计的算子(Co-designed Kernels),以及专门针对长上下文 Prefill 阶段定制的计算加速器。

六种加速器上的采样器各自运行特定的精度和算子,向 B200 上的单个 BF16 训练器提供样本。
六种加速器上的采样器各自运行特定的精度和算子,向 B200 上的单个 BF16 训练器提供样本。

针对上述痛点,本研究探讨了是否能够构建一个混合精度、混合加速器的采样系统——各节点采用不同的硬件架构、量化精度和推理算子,同时依然维持策略训练的高性能与强稳定性。该系统主要面临两大核心挑战:

  1. 学习层面的挑战(训练与推理失配):使用不同的低精度与算子会加剧采样策略与训练策略之间的数值偏离(Mismatch),使样本具有更强的离策略(Off-policy)特性,导致训练崩溃。
  2. 系统层面的挑战(集群异构性):传统 RL 基础设施假定集群硬件同构,而混合采样体系需要弹性伸缩、能按需突发且支持多架构自适应的无服务器(Serverless)GPU 基础设施。

本工作的核心贡献与创新点如下:


背景知识与关键观察

训练与推理失配的来源解析

各组采样器的采样器与训练器 KL 散度,在采用 SC + MIS + top-p 训练期间的平均值。
各组采样器的采样器与训练器 KL 散度,在采用 SC + MIS + top-p 训练期间的平均值。

为何不追求绝对的同策略(True On-Policy)RL?

GRPO 与 SC + MIS + top-p 在混合采样器和同构采样器上的训练奖励与采样器-训练器 KL 散度。
GRPO 与 SC + MIS + top-p 在混合采样器和同构采样器上的训练奖励与采样器-训练器 KL 散度。

方法细节

离线策略修正的核心数学原理与双重误差

掩码重要性采样(Masked Importance Sampling, MIS)

重要性比率的截断与掩码。虚线表示无界权重。
重要性比率的截断与掩码。虚线表示无界权重。

分数中心化(Score Centering, SC)

无修正(左)情况下,漂移不断将训练器拉向随之移动的采样器;分数中心化(右)消除了这种牵引力。
无修正(左)情况下,漂移不断将训练器拉向随之移动的采样器;分数中心化(右)消除了这种牵引力。

带掩码重放的 Top-p 采样(Top-p Sampling with Mask Replay)

在某一位置训练器的概率分布,展示无掩码重放与有掩码重放的情况。阴影 Token 落在采样器的 top-p 集合之外;虚线轮廓表示重新归一化之前的概率。
在某一位置训练器的概率分布,展示无掩码重放与有掩码重放的情况。阴影 Token 落在采样器的 top-p 集合之外;虚线轮廓表示重新归一化之前的概率。

复合策略梯度的端到端集成公式

综合上述三项修正组件,对于每个轨迹中位置 $t$ 采样的 Token $x_t$(其在采样器 Top-p 候选集内部,对应局部优势估计为 $\hat{A}_t$),训练器实际执行的参数反向传播更新梯度为:

$$g_\theta = \sum_{t} w_{\text{MIS}}(x_t) \cdot \hat{A}_t \cdot \left( \nabla_\theta \log p_\theta^{\text{replay}}(x_t) - \sum_{x \in \text{Top-}p} q(x) \nabla_\theta \log p_\theta^{\text{replay}}(x) \right)$$


其中 $w_{\text{MIS}}(x_t)$ 严格受限于硬截断区间 $[c_{\min}, c_{\max}]$。在该更新结构下:

  1. Top-p 掩码重放确保了两端处于完全相同的离散采样支撑集中;
  2. 内部减除项(SC)完全消解了由截断或分布偏移产生的常数项更新漂移;
  3. 外部重要性比率(MIS)确保在非漂移方向上,梯度更新的方向严格收敛于当前训练策略 $p_\theta$ 的无偏流形。

方法细节引用文献汇总

引用编号与论文信息 引用段落主题 原文描述内容要点
https://arxiv.org/pdf/2402.03300
DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models (2024)
算法演进起点 作为未修正强化学习策略梯度的基线实现(GRPO without any correction for the mismatch),通过组内相对优势归一化执行基准更新。
https://arxiv.org/pdf/2510.18855v2#page=7
Masked Importance Sampling 相关研究 (2025)
权重修正设计(MIS) 阐明重要性加权在逐 Token 下易受极端小概率除数影响,需引入硬掩码将超出区间的 Token 梯度归零以限制更新方差。
https://arxiv.org/pdf/2609.20807
Score Centering in Policy Gradient Methods (2026)
消除梯度漂移设计(SC) 详细论证在异构采样分布 $q$ 下对数概率梯度期望不为零引发的持续更新漂移,提出在各位置减除期望分数项以抹杀该非零常数漂移。
https://arxiv.org/pdf/2512.02556v1#page=8
DeepSeek-V3.2 Technical Report (2025)
动作空间支撑集对齐(Top-p) 指出 Top-p 采样会构建截断的动作空间,破坏两端共有支撑集,必须在训练端对候选集进行掩码重放和重归一化。
https://arxiv.org/pdf/1802.01561
IMPALA: Scalable Distributed Deep-RL with Importance Weighted Actor-Learner Architectures (2018)
离策略修正的体系源流 说明从解耦的采样进程(Actors)向中心学习器(Learner)传输带延迟经验时,采用算法级权重修正替代强同步锁步机制的历史理论根基。
https://arxiv.org/pdf/2512.01374
Non-determinism and Token-level Approximation in RL (2025)
算子失配与逐 Token 近似 证实不同推理与训练算子会导致相同输入输出不同概率,并指明逐 Token 重要性采样比率是对整条轨迹连乘的一阶近似展开。
https://arxiv.org/pdf/2510.11370
Router Inconsistency in Mixture-of-Experts RL (2025)
稀疏网络路由分歧 讨论 MoE 架构在不同精度或扰动下出现的路由器专家分歧(Router Inconsistency),以及强制执行 Rollout Routing Replay 所带来的目标策略有偏性与系统开销。
https://arxiv.org/pdf/2602.15763v2#page=12
Sparse Attention and Key-Value Routing in LLMs (2026)
稀疏索引器不匹配 阐述长上下文稀疏注意力机制中两端对 Key-Value 动态条目的索引分歧,指出全量重放该索引会带来无法承受的系统通信瓶颈。
https://arxiv.org/pdf/2507.18071
GSPO: Group Sequence Policy Optimization (2025)
序列级修正局限分析 对比序列级策略修正算法(Sequence-level Corrections),指出逐 Token 近似虽忽略了前缀与优势累积误差,但实现轻量,且实验证明在大模型长序列下依然具备足够韧性。

实验环境

数据集与评估基准

模型与训练-采样架构

采样器(Samplers)异构分组配置

采样集群基于 SGLang 推理引擎构建,运行在 Modal 的弹性无服务器(Serverless)GPU 架构之上。采样环境划分为两大对照体系:

每个采样器组及其所处理的服务请求份额。
每个采样器组及其所处理的服务请求份额。
  1. 混合采样系统(Mixed Samplers):跨越六种完全异构的 GPU 硬件平台,划分为 8 个服务执行组,按各自实时吞吐能力通过全局路由器动态分配会话(各组处理份额保持稳定):

    • B200 节点组:权重采用 NVFP4 量化,KV Cache 统一为 FP8。
    • B300 节点组:权重采用 NVFP4 量化,KV Cache 统一为 FP8。
    • H100 节点组 (FP8):权重采用 FP8 量化,KV Cache 统一为 FP8。
    • H100 节点组 (BF16):权重保持 BF16,KV Cache 统一为 FP8。
    • H200 节点组 (FP8):权重采用 FP8 量化,KV Cache 统一为 FP8。
    • H200 节点组 (BF16):权重保持 BF16,KV Cache 统一为 FP8。
    • A100 节点组:权重保持 BF16,KV Cache 统一为 FP8。
    • RTX PRO 6000 节点组:权重保持 BF16,KV Cache 统一为 FP8。
    • 算子特性:所有加速器运行自身架构专用的优化推理算子。
  2. 同构采样系统(Homogeneous Samplers - 基准线):

    • 纯粹运行在 NVIDIA B200 GPU 之上。
    • 权重强制使用 BF16,KV Cache 同样强制使用 BF16。
    • 完全匹配训练器的硬件架构与计算精度,两端仅存在由于异步流水线引发的策略滞后性(Staleness)以及推训算子差异。

实验结果

训练稳定性与消融对比实验

在 SWE-bench Pro V2 及其 HARD-51 子集上的完整训练过程中,各算法配置的收敛性呈现断崖式分化:

在 HARD-51(左)和完整集(右)上训练过程中 SWE-bench Pro V2 pass@4 的变化趋势。
在 HARD-51(左)和完整集(右)上训练过程中 SWE-bench Pro V2 pass@4 的变化趋势。

策略崩溃的核心物理表征与失效模式分析

所有运行的训练奖励、采样器-训练器 KL 散度、轨迹长度以及因重复格式错误而终止的回合比例。
所有运行的训练奖励、采样器-训练器 KL 散度、轨迹长度以及因重复格式错误而终止的回合比例。

通过监控训练过程中的动力学状态,论文总结了策略走向崩溃的三大确定性先兆:

  1. 训练奖励断崖下跌:环境交互 Reward 迅速归零;
  2. 采样-训练分布距离剧烈失控:采样器与训练器之间的 KL 散度激增 3 倍以上;
  3. 格式失效与工具调用中止:绝大多数交互回合因为重复的格式错误而异常早报终止。

相比之下,序列轨迹长度(Trajectory Length)并不构成崩溃的判别依据:部分崩溃模型因无法输出有效指令而在极短轮次内提前挂掉,而另一部分崩溃模型则陷入无意义的循环生成而跑满步数。

健康回合与崩溃回合的对比,每个单元格代表一轮交互。
健康回合与崩溃回合的对比,每个单元格代表一轮交互。

崩溃 Agent 的微观行为模式:深入单回合日志表明,崩溃的模型极少输出底层语法错误的 Bash 指令。相反,模型通常在完成部分探索后突然停止调用工具接口。由于环境判定规则为"连续 3 轮无工具调用则强制结算",Agent 随即被迫中断。

混合采样器与同构采样器的性能平价与成本收益

混合采样器与同构采样器在采用 SC + MIS + top-p 时的训练奖励、采样器-训练器 KL、梯度范数和熵。
混合采样器与同构采样器在采用 SC + MIS + top-p 时的训练奖励、采样器-训练器 KL、梯度范数和熵。

在均配置 SC + MIS + Top-p 离线策略修正的前提下,对混合硬件采样与同构纯 B200 采样进行了直接对决:

在 HARD-51(左)和完整集(右)上采用 SC + MIS + top-p 时,SWE-bench Pro V2 pass@4 与累积采样成本的关系。
在 HARD-51(左)和完整集(右)上采用 SC + MIS + top-p 时,SWE-bench Pro V2 pass@4 与累积采样成本的关系。

修正组件对系统动力学的影响(消融细节)

在混合采样器上每种方法的重要性比率落在 [0.2, 5] 之外的 Token 比例、采样器-训练器 KL、梯度范数和熵。
在混合采样器上每种方法的重要性比率落在 [0.2, 5] 之外的 Token 比例、采样器-训练器 KL、梯度范数和熵。

异构硬件组别的离群 Token 分布特征

在混合采样器上采用 SC + MIS + top-p 时,各组在训练样本和重要性比率落在 [0.2, 5] 之外的 Token 中的份额。
在混合采样器上采用 SC + MIS + top-p 时,各组在训练样本和重要性比率落在 [0.2, 5] 之外的 Token 中的份额。

对各个混合采样组别输出的 Token 质量进行细粒度归因分析表明:


补充细节与局限性

算法与系统层面的局限性(Limitations)

架构解耦对未来基础设施的启示(What's Next)


结论

本研究针对大模型强化学习中采样端霸占 80%–85% 计算资源、且被迫与训练端绑定同构硬件与精度的低效现状,提出并成功验证了混合精度、混合加速器采样架构的可行性。

  1. 破除同构教条:研究打破了"高稳定性必须依赖训练-采样两端数值绝对接近"的传统认知,证明训练崩溃的本质诱因在于更新漂移引发的正反馈无界膨胀,而非单步硬件失配的绝对大小。
  2. 算法协同根治发散:通过将掩码重要性采样(MIS)、分数中心化(SC)以及带掩码重放的 Top-p 采样有机结合,成功解决了异构采样带来的加权错误、梯度漂移以及动作空间不一致问题。在单步 KL 散度达到同构系统 4.7 倍的恶劣异构环境下,实现了全流程无崩溃的长序列完全异步强化学习。
  3. 性能平价与成本优化:在长达 256K 交互序列的智能体编程评测(SWE-bench Pro V2)中,混合采样架构取得了与纯旗舰 B200 BF16 同构基准持平的通过率,同时将采样成本直接削减约 15%,验证了解耦架构兼顾模型性能、系统稳健性与经济性的巨大工程与理论价值。