Modal 验证混合精度、混合加速器采样器上的稳定 RL 训练

In large-model reinforcement learning (RL), samplers usually run on the same accelerators and often the same precision as the trainer, keeping the two numerically close for stable training. Outside of RL, however, serving systems are not constrained in this manner: they’re free to optimize generation for speed and cost using low-precision formats such as FP8, MXFP4 and NVFP4, co-designed kernels, and even accelerators built just for long-context prefill.

Since sampling already takes 80–85% of the GPUs in large RL runs, this leaves a lot on the table! We ask whether we can use mixed-precision, mixed-accelerator samplers, each running its own kernels, while keeping training stable and performant.

Two problems stand in the way.

The first problem is learning. Each of these optimizations adds to training–inference mismatch, making samples more off-policy, which often destabilizes training. Off-policy correction is the natural solution; it lets a trainer learn from data generated by a different policy, such as stale samples in asynchronous training. Since it makes no assumption about why the two policies differ, it should apply equally to different precisions or accelerators.

The second problem is the system. RL infrastructure is usually built around clusters of homogeneous accelerators. Mixed samplers require the opposite: autoscaling accelerators that burst when needed and run using their own precisions and kernels.

Through off-policy correction and our autoscaling, serverless GPUs, we show that stable, performant, and cheaper RL is possible with mixed-precision, mixed-accelerator samplers by training an open-weight model fully asynchronously on agentic coding tasks.

Our setup

We train Qwen3.6-35B-A3B on agentic coding with up to 256K tokens of context using the 2.7K repository tasks Xiaomi released with MiMo-V2.6, after removing the repository history and other traces through which some tasks leak the reference fix. In each episode, the model creates a patch for an issue over multiple turns and submits it through a bash tool in a Modal Sandbox. It receives a reward of 1 if the submitted patch passes the task’s tests in a fresh sandbox, and 0 otherwise. We evaluate on SWE-bench Pro V2, which has no overlap with the training set, and its hard subset HARD-51 separately. Every checkpoint is evaluated in BF16.

We use Miles as the RL framework, with SGLang samplers and a Megatron-LM trainer that updates the policy in BF16 on B200 GPUs. Training is fully asynchronous with no limit on staleness, and each update uses 128 prompts with 8 trajectories each. The samplers keep their own accelerators, precisions and kernels, and the trainer does not replay their MoE routing decisions.

The samplers run in one of two setups:

1. The mixed samplers span six kinds of accelerators: B200 and B300 serve NVFP4 weights, H100 and H200 serve FP8 and BF16 weights, and A100 and RTX PRO 6000 serve BF16 weights. Each sampler runs its accelerator’s own kernels and keeps its KV cache in FP8. A router assigns each agent session to one of these eight groups, in proportion to the group’s serving capacity, so each group serves a steady share of the requests throughout training.
2. The homogeneous samplers run only on B200 with BF16 weights and a BF16 KV cache. They match the trainer’s accelerator and precision, so only kernels and staleness separate them.

On Modal, each group of samplers scales up or down as needed, even mid-run, and serves the trainer’s newest weights in its own precision.

Making mixed-sampler RL stable

We start from GRPO without any correction for the mismatch, and add three pieces:

• Masked importance sampling (MIS) corrects the weighting. It multiplies each token’s term by its importance ratio and masks tokens whose ratio falls outside a fixed range, so on the tokens it keeps, the expected update follows the training policy and every ratio stays bounded. Without the mask, a few tokens the samplers rarely produce could carry very large ratios and dominate the update.
• Score centering (SC) removes drift. At each position, it subtracts the expected gradient of the log-probability under the sampling policy, so in expectation only differences in advantage between tokens move the policy. Importance ratios also remove drift, but masking reintroduces it.
• Top-p sampling with mask replay (top-p) keeps rare tokens out and gives both policies the same support. The samplers sample each token only from the top-p candidate set, so rare tokens, which can carry extreme ratios, never enter the data. The trainer then renormalizes its own probabilities over the same set. Runs without it sample from the full vocabulary.

Each sampled token adds its advantage times the gradient of its log-probability to the update. Put together, MIS reweights that term by the importance ratio, SC centers the gradient, and top-p restricts both policies to the same candidates.

Without correction, training collapses on both setups: on the mixed samplers, GRPO collapses soon after training starts. On the homogeneous samplers, training lasts much longer before ultimately collapsing as well. Each piece of the correction on its own delays the collapse, but every such run ends at or below the base model. Only SC + MIS, with or without top-p, trains stably to the end: pass@4 rises over roughly the first 100–140 steps, then levels off at about 65–75% on HARD-51. The full set shows the same pattern from a higher start, from 89% to about 94%.

Every collapse shares three signs: the training reward falls, the sampler–trainer KL grows 3× or more, and most episodes end on repeated format errors. Trajectory length does not show a consistent pattern: some collapsed runs end their episodes early, while others start repeating themselves and run longer. The stable runs show none of these signs. Matching the trainer’s accelerator and precision only delays the collapse: GRPO on the homogeneous samplers starts from a KL about 4× lower, but follows the same path from about step 100.

A collapsed episode rarely writes a malformed command. Instead, partway through a task, the model stops calling its tool, and after three replies in a row without a tool call, the agent ends the episode. At GRPO’s step-40 checkpoint on the mixed samplers, 62% of evaluation episodes end this way, a median of 22 turns in, as opposed to 57 turns for passing episodes. The last reply is empty in 98% of them, while only 6% contain a malformed call.

On par with homogeneous samplers at lower cost

With SC + MIS + top-p, the mixed samplers reach SWE-bench Pro V2 scores on par with the homogeneous samplers’: the two runs’ pass@4 stays close throughout training. Their training reward and gradient norm also follow similar paths, despite a 4.7× larger sampler–trainer KL on average. Entropy matches until about step 160; afterwards, it rises on the homogeneous samplers and stays flat on the mixed ones. Importantly, nothing in how the correction works is specific to this model or these accelerators, so we expect this pattern to hold more broadly.

Decoupling the samplers from the trainer also makes these results cheaper to reach. For most of training, the samplers’ cost per update was about 15% lower on the mixed samplers than on the homogeneous samplers. Note that our setup was built to validate the performance of using mixed-precision, mixed-accelerator samplers rather than for minimal cost, so there are more savings to be had: every sampler can be tuned on its own, and the mix can shift toward the most cost-effective samplers.

The next three sections explain these results: where training–inference mismatch comes from, why we don’t use true on-policy RL, and how off-policy correction works.

Where training–inference mismatch comes from

A sampled token can get one probability from the sampler and another from the trainer. The causes span the whole RL system:

Mixed samplers combine most of these at once, and their mismatch is on average 4.7× that of the homogeneous samplers.

Why not true on-policy RL?

Training is on-policy when every token is sampled from the policy being trained. When the samplers and the trainer are separate systems, being on-policy requires two things: the samplers to compute the same bitwise probabilities as the trainer (i.e., true on-policy RL), and every sample to come from the trainer’s current weights.

Bitwise-identical sampling and training has been demonstrated. Open RL frameworks make the two identical for small models and nearly identical for GLM-5, and DeepSeek-V4 built batch-invariant, deterministic kernels to make its training and inference bitwise reproducible. Matching every source, though, has a cost:

Most of these sources change only token probabilities, which correction accounts for. Multi-turn context and top-p / top-k sampling break the shared context and shared support that correction relies on, so our runs match them: token-in-token-out, and mask replay whenever we sample with top-p. For the rest, matching is a choice.

In practice, large runs match some of these sources and attempt to correct for the rest. Runs that match the precision, routing and top-p mask still bound staleness and mask tokens with extreme importance ratios, and runs without staleness still correct the training–inference mismatch. At scale, bitwise matching has meant kernels are written to be specific to the model. But even then, an asynchronous run remains off-policy through staleness.

The common view is that the lower this mismatch, the more stable the training, which is why runs keep the samplers numerically close to the trainer. However, stability depends less on the magnitude of the mismatch than on whether it keeps growing. Without correction, the mismatch biases the update, which can widen the mismatch until training collapses. A smaller mismatch only delays this: GRPO on the homogeneous samplers started with a KL about 4× lower than on the mixed samplers but still collapsed later. With correction, training on the mixed samplers stays stable and the performance is on par with the homogeneous samplers despite the substantial mismatch.

The magnitude of the mismatch and its growth have different causes. The magnitude comes from the system: the accelerator, precision and kernels each sampler runs. On the other hand, the growth comes from the update, which is what correction fixes. So accelerators and precisions need no new kind of fix: off-policy correction, which large runs already apply, extends to them. The goal is to keep the mismatch bounded, not to keep it low.

So we don’t pursue true on-policy RL: the samplers and the trainer are built for different jobs, and each keeps improving at its own. Matching ties them together, and each gives up some of its optimizations; with correction, each keeps them.

How off-policy correction works

Learning from data another policy generated is an established problem in RL. IMPALA separated the actors that generate experience from the learner and corrected for policy lag rather than keeping them in lockstep. The sampling policy q generates the tokens, and the training policy p_θ learns from them. With mixed samplers, each token comes from one sampler at one weight version, where each sampler records the token’s probability, so q can differ from token to token. The correction never needs to know why q differs from p_θ: kernels, precision, accelerators and staleness all show up only as a different q.

When tokens come from q instead of p_θ, the policy-gradient update at each position picks up two errors.

Drift. Under the training policy, the gradient of the log-probability averages to zero at each position; under q, it averages to a step toward q. So part of the update moves the training policy regardless of which tokens earned the reward: toward the sampling policy when the expected advantage is positive, and away from it when it is negative. Normalizing advantages within a group does not remove this, since the expected advantage after a given prefix is generally not zero. The samplers then load the trainer’s new weights and add their error again, so the trainer keeps moving toward a target that moves with it.

The wrong weighting. The update weights each token by how often the samplers produce it, not by the training policy. A token can keep receiving a large gradient after the training policy has made it unlikely.

Both errors vanish when q = p_θ and grow otherwise.

Importance sampling fixes both errors by weighting each token’s term by its ratio p_θ / q, provided the samplers can sample every token the training policy can. Weighting a whole trajectory by the product of its ratios is exact, but its variance can grow exponentially with length, so large runs apply the ratio one token at a time, a first-order approximation whose error grows with the mismatch and with trajectory length. Even per token, a token the samplers rarely produce can carry a weight up to 1 / q. Bounding the weight, by truncating or masking it, controls the variance but reintroduces drift. MIS masks: it keeps the ratio inside a fixed range and sets it to zero outside.

Score centering removes drift instead. At each position, it subtracts the expected gradient of the log-probability under q, so in expectation a constant advantage no longer moves the policy. It keeps the wrong weighting, which is why it pairs with MIS: centering the weighted term removes drift whatever the weight, including the drift that masking brings back. SC + MIS removes drift and weights the tokens it keeps by the training policy. Computing the expected gradient under q needs q for every token. The samplers record it for their most likely tokens; for every other token, q is approximated by the trainer’s probabilities, rescaled to the samplers’ remaining probability mass.

Top-p sampling breaks the shared support, and mask replay restores it. The samplers sample only from the top-p set, so tokens outside it have zero probability under q. Replaying the mask in the trainer renormalizes p_θ over the same set. The rarest tokens, which can carry the largest ratios, never enter the data, and the samplers record every candidate’s probability, so centering is exact rather than estimated.

What our runs show

• GRPO keeps both errors. Its gradient norm grows about 20× soon after training starts.
• GRPO + MIS fixes the weighting but not the drift its mask brings back. Its gradient norm stays within a few times its start, but partway through training its entropy collapses as its KL and its share of tokens with ratio outside [0.2, 5] jump more than 10×.
• SC removes drift but keeps the wrong weighting. Its entropy falls steadily, and its KL grows about 3× before the run collapses.
• GRPO + top-p keeps the rarest tokens out. While GRPO’s gradient norm blows up, its own stays 3–10× smaller, but its KL climbs steadily and it collapses later.
• SC + MIS removes drift and fixes the weighting on the tokens it keeps. It stays stable apart from one brief gradient spike, while its KL rises about 2×.
• SC + MIS + top-p also centers exactly. It has the lowest KL, and its gradient norm and entropy stay flat.

SC + MIS + top-p trains on an average mismatch 4.7× that of the homogeneous samplers and stays stable. In every run that collapses, the mismatch grows 3× or more.

Per group of samplers, the correction’s work follows the mismatch. A group’s share of the extreme-ratio tokens follows its KL more than its share of the samples: H100 FP8 supplies 10% of the samples but 29% of these tokens, while B200 and B300 NVFP4 supply 47% of the samples and 16%. Even for H100 FP8, only about 0.3% of its tokens fall outside [0.2, 5]. The correction weights each token by its own ratio, so the same correction keeps training stable with every group in the mix.

Limitations

MIS, SC and top-p all work one token at a time, so each inherits the per-token approximation’s limit: the prefixes and the advantages still come from the samplers. Sequence-level corrections, such as GSPO’s and DeepSeek-V3.2’s, and bounding staleness guard against that error when it grows too large. Our runs use none of them and stay stable, with contexts of up to 256K tokens and an average staleness of about 4.

Top-p sampling with mask replay also changes what is trained: the update optimizes the policy truncated to the top-p candidate set, and the logits outside it get no gradient at that position. Our evaluation samples from the full vocabulary, and SC + MIS + top-p evaluated about as well as SC + MIS, which samples from the full vocabulary in training too.

Correction also leaves the samples as they are: it changes how the update uses each token, not which trajectories the samplers produce. Choosing a precision or an accelerator therefore remains a trade, but between throughput and sample efficiency rather than between throughput and stability.

These corrections are not the only options, but they share one principle: the mismatch can be large, as long as the update stays close to the on-policy one and the mismatch stays bounded. Better ones are likely to follow.

What’s next

Correction decouples the samplers from the trainer, and that changes what RL can run on. Sampling no longer needs a large homogeneous cluster: any accelerator that can serve the model can contribute samples, so the sampling fleet can be sized separately from the trainer and grown mid-run as the run needs more samples (e.g., as trajectories get longer). Inference optimizations also carry straight into RL: faster kernels, speculative decoding, lower precision, or prefill and decode split across different accelerators can all speed up sampling without the trainer having to match them. And the trainer can move to the best accelerators and its own low precision, such as FP8 or NVFP4, without the samplers following. As accelerators keep diversifying across vendors, generations and number formats, this decoupling becomes essential to enable frontier training for all.

We’re working on hard problems where RL algorithms meet systems at scale. If those sound like the kinds of problems you’re interested in, join us.

Acknowledgement

We’d like to thank the Miles team, Luke Huang, Martin Marek, Dhruv Gautam, Xiuyu Li, Byron Hsu, Ziang Li and Zilin Zhu for discussions on a draft of this post.