Modal 研究混合精度、混合加速器采样器上的强化学习

在大模型强化学习（RL）中，采样器通常与训练器运行在相同的加速器上，并且往往使用相同的精度，从而使两者在数值上保持接近，以保证训练稳定。然而在 RL 之外，服务系统并不受此约束：它们可以自由地使用 FP8、MXFP4 和 NVFP4 等低精度格式、协同设计的 kernel，甚至专为长上下文 prefill 打造的加速器，来针对速度和成本优化生成。

由于采样已经占用了大规模 RL 训练中 80–85% 的 GPU，这里还有很大的优化空间！我们提出的问题是：能否使用混合精度、混合加速器的采样器，各自运行自己的 kernel，同时保持训练稳定且高效。

六种加速器上的采样器，各自拥有自己的精度和 kernel，为一个运行在 B200 上的 BF16 训练器提供数据。

有两个问题横亘在前。

第一个问题是学习。这些优化中的每一项都会加剧训练与推理之间的不匹配，使样本更加 off-policy，这往往会破坏训练的稳定性。Off-policy 校正是自然的解决方案；它让训练器能够从由不同策略生成的数据中学习，例如异步训练中的陈旧样本。由于它不对两个策略为何不同做任何假设，因此应当同样适用于不同的精度或加速器。

第二个问题是系统。RL 基础设施通常围绕同构加速器集群构建。混合采样器需要的恰恰相反：能够按需自动扩缩容、在需要时突发、并使用各自精度和 kernel 运行的加速器。

通过 off-policy 校正和我们的自动扩缩容、无服务器 GPU，我们展示了：通过在 agentic coding 任务上完全异步地训练一个开放权重模型，使用混合精度、混合加速器的采样器可以实现稳定、高效且更便宜的 RL。

我们的设置

我们使用小米随 MiMo-V2.6 发布的 2.7K 仓库任务，在 agentic coding 上训练 Qwen3.6-35B-A3B，上下文长度最高达 256K token，并事先移除了仓库历史以及其他一些任务借此泄露参考修复的痕迹。在每个 episode 中，模型通过多轮对话为某个 issue 创建补丁，并通过 Modal Sandbox 中的 bash 工具提交。如果提交的补丁在全新 sandbox 中通过该任务的测试，则获得奖励 1，否则为 0。我们在与训练集无重叠的 SWE-bench Pro V2 及其困难子集 HARD-51 上分别进行评估。每个 checkpoint 都以 BF16 进行评估。

我们使用 Miles 作为 RL 框架，搭配 SGLang 采样器和 Megatron-LM 训练器，在 B200 GPU 上以 BF16 更新策略。训练完全异步，对陈旧度没有限制，每次更新使用 128 个 prompt，每个 prompt 8 条轨迹。采样器保留各自的加速器、精度和 kernel，训练器不会重放它们的 MoE 路由决策。

采样器运行在以下两种设置之一：

混合采样器涵盖六种加速器：B200 和 B300 承载 NVFP4 权重，H100 和 H200 承载 FP8 和 BF16 权重，A100 和 RTX PRO 6000 承载 BF16 权重。每个采样器运行其加速器自有的内核，并将其 KV 缓存保持为 FP8。路由器将每个智能体会话分配给这八个组之一，分配比例与该组的服务容量成正比，因此每个组在整个训练过程中都服务稳定份额的请求。
同构采样器仅在 B200 上运行，使用 BF16 权重和 BF16 KV 缓存。它们与训练器的加速器和精度相匹配，因此只有内核和陈旧度将它们区分开来。

Each group of samplers and its share of the requests served. — 每组采样器及其所服务请求的份额。

在 Modal 上，每组采样器都按需扩缩容，甚至可以在运行中途进行，并以自己的精度服务训练器的最新权重。

使混合采样器 RL 稳定

我们从 GRPO 出发，不对不匹配做任何修正，然后添加三个部分：

掩码重要性采样（MIS）修正权重。它将每个 token 的项乘以其重要性比率，并掩码比率落在固定范围之外的 token，因此在其保留的 token 上，期望更新遵循训练策略，且每个比率都保持有界。如果没有掩码，采样器很少产生的少数 token 可能携带非常大的比率并主导更新。
分数居中（SC）消除漂移。在每个位置，它减去采样策略下对数概率的期望梯度，因此期望上只有 token 之间优势的差异会推动策略。重要性比率也能消除漂移，但掩码会重新引入漂移。
带掩码重放的 Top-p 采样（top-p）将稀有 token 排除在外，并让两个策略具有相同的支持集。采样器仅从 top-p 候选集中采样每个 token，因此可能携带极端比率的稀有 token 永远不会进入数据。训练器随后在相同的集合上重新归一化自己的概率。没有它的运行则从完整词表采样。

每个被采样的 token 将其优势乘以其对数概率的梯度添加到更新中。综合起来，MIS 通过重要性比率对该项重新加权，SC 使梯度居中，top-p 将两个策略限制在相同的候选集上。

SWE-bench Pro V2 pass@4 over training on HARD-51 (left) and the full set (right). — SWE-bench Pro V2 pass@4 在 HARD-51（左）和完整集（右）上的训练过程。

如果没有修正，训练在两种设置上都会崩溃：在混合采样器上，GRPO 在训练开始后不久就崩溃。在同构采样器上，训练持续更长时间，但最终也会崩溃。每个修正部分单独都能延迟崩溃，但每次这样的运行最终都达到或低于基础模型。只有 SC + MIS，无论有无 top-p，都能稳定训练到最后：pass@4 在大约前 100–140 步上升，然后在 HARD-51 上稳定在约 65–75%。完整集从更高的起点显示出相同的模式，从 89% 到约 94%。

Training reward, sampler–trainer KL, trajectory length and episodes ending on repeated format errors, for every run. — 每次运行的训练奖励、采样器–训练器 KL、轨迹长度以及因重复格式错误而结束的回合。

每次崩溃都有三个共同迹象：训练奖励下降，采样器与训练器之间的 KL 散度增长 3 倍或更多，且大多数回合以重复的格式错误结束。轨迹长度并未呈现出一致的模式：一些崩溃的运行会提前结束回合，而另一些则开始自我重复并运行更长时间。稳定的运行则没有这些迹象。匹配训练器的加速器和精度只能推迟崩溃：在同构采样器上运行 GRPO 时，初始 KL 散度约低 4 倍，但从大约第 100 步开始便遵循相同的路径。

A healthy and a collapsed episode, with one cell per turn. — 一个健康回合与一个崩溃回合，每轮一个单元格。

崩溃的回合很少写出格式错误的命令。相反，在任务进行到一半时，模型停止调用其工具，在连续三次回复都没有工具调用后，智能体便结束该回合。在混合采样器上 GRPO 的第 40 步检查点，62% 的评估回合以这种方式结束，中位数为第 22 轮，而通过的回合则为 57 轮。其中 98% 的最后一条回复为空，而只有 6% 包含格式错误的调用。

以更低成本达到与同构采样器相当的水平

Training reward, sampler–trainer KL, gradient norm and entropy for SC + MIS + top-p on the mixed and the homogeneous samplers. — 在混合采样器和同构采样器上使用 SC + MIS + top-p 时的训练奖励、采样器-训练器 KL 散度、梯度范数和熵。

使用 SC + MIS + top-p 时，混合采样器达到的 SWE-bench Pro V2 分数与同构采样器相当：两个运行的 pass@4 在整个训练过程中保持接近。它们的训练奖励和梯度范数也遵循相似的路径，尽管平均采样器-训练器 KL 散度大 4.7 倍。熵在约第 160 步之前保持一致；之后，同构采样器上的熵上升，而混合采样器上的熵保持平稳。重要的是，校正的工作方式并不特定于此模型或这些加速器，因此我们预计这一模式会更广泛地适用。

SWE-bench Pro V2 pass@4 against cumulative sampling cost for SC + MIS + top-p on HARD-51 (left) and the full set (right). — 在 HARD-51（左）和完整集（右）上使用 SC + MIS + top-p 时，SWE-bench Pro V2 pass@4 与累积采样成本的对比。

将采样器与训练器解耦也使这些结果更易以更低成本实现。在训练的大部分时间里，混合采样器上每次更新的采样器成本比同构采样器低约 15%。请注意，我们的设置是为了验证使用混合精度、混合加速器采样器的性能，而非追求最低成本，因此还有更多节省空间：每个采样器都可以单独调优，并且混合比例可以向最具成本效益的采样器倾斜。

接下来的三个部分将解释这些结果：训练-推理不匹配从何而来，我们为何不使用真正的同策略 RL，以及离策略校正如何工作。

训练-推理不匹配从何而来

一个被采样的 token 可能从采样器获得一个概率，而从训练器获得另一个概率。原因贯穿整个 RL 系统：

来源	示例	影响
内核	分离的推理和训练内核；无批处理不变性	相同输入产生不同输出，即使精度相同
精度	BF16；FP8；MXFP4；NVFP4	格式在精度（尾数位）和动态范围（指数位、块缩放）上不同，因此两侧产生不同的舍入误差；随着更多计算被量化，不匹配可能增大
加速器	架构特定的内核和数字格式：A100 上没有 FP8 张量核心，Hopper 上没有 FP4 张量核心	相同代码在不同加速器上产生不同的不匹配
策略陈旧度	完全异步 RL，其中采样器落后若干训练步运行；部分 rollout；在途权重更新	可能随陈旧度增长的不匹配
路由	Top-k 专家选择；稀疏注意力索引器的 top-k 选择	路由器不一致：两侧为同一 token 选择了不同的专家，或不同的键值条目
Top-p / top-k 采样	token 仅从 top-p 或 top-k 候选中采样	动作空间不匹配，导致训练不稳定
多轮上下文	反分词–再分词往返；聊天模板重新渲染早前的轮次或剥离其推理过程	训练器在模型从未见过的上下文中对 token 打分

混合采样器会同时结合其中大部分来源，其不匹配平均是同质采样器的 4.7 倍。

Sampler–trainer KL for each group of samplers, averaged over training with SC + MIS + top-p. — 每组采样器的采样器–训练器 KL，在训练中取平均，使用 SC + MIS + top-p。

为什么不是真正的同策略 RL？

当每个 token 都从正在训练的策略中采样时，训练就是同策略的。当采样器和训练器是分离的系统时，要做到同策略需要两件事：采样器计算出与训练器逐位相同的概率（即真正的同策略 RL），并且每个样本都来自训练器当前的权重。

逐位相同的采样与训练已被证明可行。开放 RL 框架使两者对小模型完全相同，对 GLM-5 几乎完全相同，而 DeepSeek-V4 构建了批次不变、确定性的内核，使其训练和推理逐位可复现。然而，匹配每一个来源是有代价的：

来源	如何匹配	代价
内核	两侧使用批次不变、确定性的内核	采样更慢，或为每个模型编写内核
精度	两侧使用相同格式	训练器被绑定到采样器的格式，且没有单一低精度计算格式能在所有加速器上运行：A100 没有 FP8 张量核心，Hopper 没有 FP4
加速器	采样和训练在同一代加速器上	排除了所有其他加速器
策略陈旧度	同步 RL，每批一次更新	放弃了异步 RL 的吞吐量
路由	Rollout 路由重放；索引器重放	对 MoE 而言，有偏的目标策略，外加内存和通信开销；对索引器而言，开销大得多
Top-p / top-k 采样	Top-p / top-k 掩码重放	匹配候选集，而非概率；掩码存储与传输
多轮上下文	Token 进 Token 出	若保留早前推理则上下文更长；若剥离则训练计算更多

这些来源大多只改变 token 概率，而校正可以处理这一点。多轮上下文和 top-p / top-k 采样破坏了校正所依赖的共享上下文和共享支撑集，因此我们的运行对它们进行匹配： Token 进 Token 出，以及每当用 top-p 采样时进行掩码重放。对于其余来源，匹配是一种选择。

在实践中，大规模运行会匹配其中一些来源，并尝试校正其余来源。匹配精度、路由和 top-p 掩码的运行仍会限制陈旧度，并对具有极端重要性比率的 token 进行掩码，而没有陈旧度的运行仍会校正训练–推理不匹配。在大规模下，逐位匹配意味着内核要针对特定模型编写。但即便如此，异步运行仍会因陈旧度而处于离策略状态。

通常认为这种偏差越小，训练就越稳定，因此各次运行都会让采样器在数值上尽量贴近训练器。然而，稳定性更多取决于偏差是否持续增长，而非偏差的大小。若不加以修正，偏差会使更新产生偏置，进而可能扩大偏差，直至训练崩溃。较小的偏差只是推迟了这一点：在同构采样器上，GRPO 初始的 KL 比在混合采样器上低约 4 倍，但后来仍然崩溃了。而加入修正后，尽管存在显著偏差，在混合采样器上的训练仍保持稳定，性能与同构采样器相当。

Training reward and sampler–trainer KL for GRPO and SC + MIS + top-p, each on the mixed and the homogeneous samplers. — GRPO 和 SC + MIS + top-p 在混合采样器与同构采样器上的训练奖励及采样器–训练器 KL。

偏差的大小与其增长有不同的成因。大小来自系统：每个采样器运行的加速器、精度和内核。而增长来自更新，这正是修正所解决的问题。因此加速器和精度不需要新的修正方式：大规模运行已经在使用的离策略修正可以扩展到它们。目标是让偏差保持有界，而不是让它保持低值。

所以我们不追求真正的同策略 RL：采样器和训练器是为不同任务构建的，各自都在自己的方向上不断改进。让它们匹配会把二者绑在一起，各自都要放弃一些优化；而有了修正，各自都能保留这些优化。

离策略修正如何工作

从另一个策略生成的数据中学习是 RL 中一个已确立的问题。IMPALA 将生成经验的行动者与学习器分离，并针对策略滞后进行修正，而不是让它们保持同步。采样策略 q 生成 token，训练策略 p_θ 从中学习。在混合采样器下，每个 token 来自某个采样器的某个权重版本，每个采样器都会记录该 token 的概率，因此 q 可能因 token 而异。修正从不需要知道 q 为何与 p_θ 不同：内核、精度、加速器和陈旧度都只表现为不同的 q。

当 token 来自 q 而非 p_θ 时，每个位置上的策略梯度更新会引入两个误差。

漂移。在训练策略下，对数概率的梯度在每个位置上平均为零；而在 q 下，它平均会朝着 q 迈出一步。因此，无论哪些 token 获得了奖励，更新的一部分都会推动训练策略：当期望优势为正时，推向采样策略；当期望优势为负时，远离采样策略。在组内对优势进行归一化并不能消除这一点，因为给定前缀后的期望优势通常不为零。随后采样器加载训练器的新权重，再次引入其误差，于是训练器不断朝着一个随它一起移动的目标前进。

错误的加权。更新根据采样器产生每个 token 的频率来加权，而不是根据训练策略。一个 token 可能在训练策略已使其不太可能出现之后，仍持续获得较大的梯度。

当 q = p_θ 时，这两个误差都会消失，否则就会增大。

重要性采样通过对每个 token 的项按其比率 p_θ / q 加权，同时修正这两种误差，前提是采样器能够采样到训练策略所能采样的每个 token。用整条轨迹的比率乘积来加权是精确的，但其方差会随长度呈指数增长，因此大规模运行会逐 token 地应用该比率，这是一种一阶近似，其误差会随不匹配程度和轨迹长度而增大。即便逐 token 处理，采样器很少产生的 token 也可能带有高达 1 / q 的权重。通过截断或掩码来限制权重可以控制方差，但会重新引入漂移。MIS 采用掩码：它将比率保持在一个固定范围内，并将其在范围外设为零。

Truncating and masking the importance ratio. The dashed line is the unbounded weight. — 截断和掩码重要性比率。虚线为无界权重。

分数居中则改为消除漂移。在每个位置上，它减去 q 下对数概率的期望梯度，因此在期望意义上，恒定的优势不再推动策略。它保留了错误的加权，这正是它与 MIS 搭配的原因：对加权项进行居中可消除漂移，无论权重如何，包括掩码重新引入的漂移。SC + MIS 消除了漂移，并按训练策略对其保留的 token 加权。计算 q 下的期望梯度需要每个 token 的 q。采样器为其最可能的 token 记录 q；对于其他所有 token，q 由训练器的概率近似，并重新缩放到采样器剩余的概率质量上。

无校正（左）时，漂移不断将训练器拉向一个随之移动的采样器；分数居中（右）消除了这种拉力。改编自 Martin Marek 的动画。

Top-p 采样破坏了共享支撑集，而掩码重放恢复了它。采样器仅从 top-p 集合中采样，因此该集合之外的 token 在 q 下概率为零。在训练器中重放该掩码会在同一集合上重新归一化 p_θ。那些可能带有最大比率的最稀有 token 永远不会进入数据，而采样器记录了每个候选的概率，因此居中是精确的而非估计的。

The trainer's probabilities at one position, without and with mask replay. Shaded tokens fall outside the samplers' top-p set; dashed outlines show the probabilities before renormalizing. — 训练器在某一位置上的概率，未使用和使用掩码重放时。阴影 token 落在采样器的 top-p 集合之外；虚线轮廓显示重新归一化之前的概率。

我们的运行结果显示了什么

Share of tokens with an importance ratio outside [0.2, 5], sampler–trainer KL, gradient norm and entropy, for each method on the mixed samplers. — 在混合采样器上，每种方法中重要性比率落在 [0.2, 5] 之外的 token 占比、采样器–训练器 KL、梯度范数和熵。

GRPO 保留了两种误差。训练开始后不久，其梯度范数增长约 20 倍。
GRPO + MIS 修正了加权，但未修正其掩码重新引入的漂移。其梯度范数保持在起始值的几倍以内，但在训练中途，随着其 KL 和比率落在 [0.2, 5] 之外的 token 占比跃升超过 10 倍，其熵崩溃。
SC 消除了漂移，但保留了错误的加权。其熵稳步下降，KL 增长约 3 倍后运行崩溃。
GRPO + top-p 将最稀有的 token 排除在外。当 GRPO 的梯度范数爆炸时，其自身的梯度范数保持小 3–10 倍，但其 KL 稳步攀升，并在之后崩溃。
SC + MIS 消除了漂移，并修正了其保留 token 上的加权。除了一次短暂的梯度尖峰外，它保持稳定，同时其 KL 上升约 2 倍。
SC + MIS + top-p 也恰好居中。它的 KL 最低，梯度范数和熵保持平稳。

SC + MIS + top-p 训练时的平均失配是同类采样器的 4.7 倍，且保持稳定。在每一次崩溃的运行中，失配都会增长 3 倍或更多。

Each group's share of the training samples and of the tokens with an importance ratio outside [0.2, 5], for SC + MIS + top-p on the mixed samplers. — 在混合采样器上运行 SC + MIS + top-p 时，每组在训练样本中以及重要性比率超出 [0.2, 5] 的 token 中所占的份额。

对于每组采样器，校正的工作量随失配而变化。一组在极端比率 token 中的份额更多地跟随其 KL，而非其在样本中的份额：H100 FP8 提供 10% 的样本，却贡献了这些 token 的 29%，而 B200 和 B300 NVFP4 提供 47% 的样本，却只贡献 16%。即便对于 H100 FP8，也只有约 0.3% 的 token 落在 [0.2, 5] 之外。校正按每个 token 自身的比率对其加权，因此同一校正能让混合中每一组的训练都保持稳定。

局限性

MIS、SC 和 top-p 都一次只处理一个 token，因此各自继承了逐 token 近似的局限：前缀和优势仍然来自采样器。序列级校正，例如 GSPO 和 DeepSeek-V3.2 的校正，以及对陈旧度的限制，可以在该误差过大时加以防范。我们的运行没有使用其中任何一种，仍保持稳定，上下文长度最高达 256K token，平均陈旧度约为 4。

带掩码重放的 top-p 采样也改变了训练内容：更新优化的是截断到 top-p 候选集上的策略，而该集合之外的 logits 在该位置得不到梯度。我们的评估从完整词表采样，而 SC + MIS + top-p 的评估表现与 SC + MIS 大致相当，后者在训练中也从完整词表采样。

校正也不改变样本本身：它改变的是更新如何使用每个 token，而不是采样器产生哪些轨迹。因此，选择某种精度或加速器仍然是一种权衡，但权衡的是吞吐量与样本效率，而非吞吐量与稳定性。

这些校正并非唯一的选择，但它们共享一个原则：失配可以很大，只要更新保持接近同策略更新，且失配保持有界。更好的方法很可能还会出现。

下一步

校正将采样器与训练器解耦，这改变了 RL 可以运行在什么之上。采样不再需要大型同构集群：任何能够服务该模型的加速器都可以贡献样本，因此采样集群的规模可以与训练器分开设定，并在运行中途随着运行需要更多样本（例如轨迹变长时）而扩展。推理优化也能直接带入 RL：更快的 kernel、投机解码、更低精度，或将 prefill 和 decode 拆分到不同加速器上，都能加速采样，而无需训练器与之匹配。训练器也可以迁移到最好的加速器及其自身的低精度，例如 FP8 或 NVFP4，而无需采样器跟随。随着加速器在厂商、代际和数字格式上不断多样化，这种解耦对于让所有人都能进行前沿训练变得至关重要。

我们正在攻克 RL 算法与大规模系统交汇处的难题。如果这些听起来像是你感兴趣的问题，加入我们。

致谢

我们要感谢 Miles 团队、Luke Huang、Martin Marek、Dhruv Gautam、Xiuyu Li、Byron Hsu、Ziang Li 和 Zilin Zhu 就本文草稿进行的讨论。