Toward RL on mixed-precision, mixed-accelerator samplers

Modal researchers demonstrated stable, cheaper reinforcement learning by training Qwen3.6-35B-A3B fully asynchronously with mixed-precision, mixed-accelerator samplers, using off-policy correction and autoscaling serverless GPUs. The samplers span six accelerator types — B200 and B300 serving NVFP4 weights, H100 and H200 serving FP8 and BF16 weights, and A100 and RTX PRO 6000 serving BF16 weights — each running its own kernels with an FP8 KV cache, while a Megatron-LM trainer updates the policy in BF16 on B200 GPUs. The work targets the 80–85% of GPUs consumed by sampling in large RL runs, training on 2.7K Xiaomi MiMo-V2.6 repository tasks with up to 256K tokens of context and evaluating on SWE-bench Pro V2 and its HARD-51 subset.

In large-model reinforcement learning RL , samplers usually run on the same accelerators and often the same precision https://microsoft.ai/pdf/mai-thinking-1.pdf page=52 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 https://reflection.ai/blog/introducing-beam :~:text=inference%2Dto%2Dtraining%20GPU%20ratios%20from%203.9%3A1%20to%205.4%3A1 – 85 https://microsoft.ai/pdf/mai-thinking-1.pdf page=51 % 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 https://modal.com/blog/truly-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 https://huggingface.co/Qwen/Qwen3.6-35B-A3B on agentic coding with up to 256K tokens of context using the 2.7K repository tasks Xiaomi released https://huggingface.co/datasets/XiaomiMiMo/MiMo-V2.6-RL-oss with MiMo-V2.6 https://arxiv.org/pdf/2610.11959 , after removing the repository history and other traces through which some tasks leak the reference fix https://www.vals.ai/blogs/mimo-reward-hacking . 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 https://modal.com/blog/scaling-to-1-million-concurrent-sandboxes-in-seconds . 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 https://labs.scale.com/leaderboard/swe bench pro public v2?tab=full , which has no overlap with the training set, and its hard subset HARD-51 https://labs.scale.com/leaderboard/swe bench pro public v2?tab=hard separately. Every checkpoint is evaluated in BF16. We use Miles https://github.com/radixark/miles as the RL framework, with SGLang https://github.com/sgl-project/sglang samplers and a Megatron-LM https://github.com/NVIDIA/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 https://arxiv.org/pdf/2402.03300 without any correction for the mismatch, and add three pieces: - Masked importance sampling https://arxiv.org/pdf/2510.18855v2 page=7 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 https://arxiv.org/pdf/2609.20807 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 https://arxiv.org/pdf/2512.02556v1 page=8 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: | Source | Examples | Effect | |---|---|---| | Kernels | Separate inference and training kernels; no batch invariance https://thinkingmachines.ai/blog/defeating-nondeterminism-in-llm-inference/ | Different outputs for the same input https://arxiv.org/pdf/2512.01374 , even at the same precision | | Precision | BF16; FP8 https://www.lmsys.org/blog/2025-11-25-fp8-rl ; MXFP4; NVFP4 https://humansand.ai/blog/nvfp4-rl | Formats differ in precision mantissa bits and dynamic range exponent bits, block scaling , so the two sides incur different rounding errors https://arxiv.org/pdf/2510.26788 ; mismatch can grow as more of the computation is quantized https://www.lmsys.org/blog/2026-07-29-mxfp8-nvfp4-rl/ | | Accelerator | Architecture-specific kernels and number formats: no FP8 tensor cores on A100 nor FP4 tensor cores on Hopper | Different mismatch on different accelerators with the same code | | Policy staleness | Fully asynchronous RL, where samplers run several training steps behind https://www.appliedcompute.com/research/staleness-in-fully-async-rl ; partial rollouts https://arxiv.org/pdf/2607.24653v2 page=13 ; in-flight weight updates https://arxiv.org/pdf/2509.19128 | Mismatch that can grow with staleness | | Routing | Top-k expert selection; the sparse-attention indexer’s top-k selection | Router inconsistency https://arxiv.org/pdf/2510.11370 : the two sides select different experts, or different key-value entries https://arxiv.org/pdf/2602.15763v2 page=12 , for the same token | | Top-p / top-k sampling | Tokens are sampled only from the top-p or top-k candidates | Mismatched action spaces https://arxiv.org/pdf/2512.02556v1 page=8 that destabilize training | | Multi-turn context | Detokenize–retokenize round trips; chat templates that re-render earlier turns or strip their reasoning | The trainer scores tokens in a context the model never saw https://www.lmsys.org/blog/2026-05-13-no-token-left-behind | 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 https://thinkingmachines.ai/blog/defeating-nondeterminism-in-llm-inference/ , 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 https://arxiv.org/pdf/2609.08368v1 page=24 and nearly identical for GLM-5 https://github.com/THUDM/slime/pull/2262 , and DeepSeek-V4 built batch-invariant, deterministic kernels https://arxiv.org/pdf/2606.19348v1 page=18 to make its training and inference bitwise reproducible. Matching every source, though, has a cost: | Source | How it is matched | Cost | |---|---|---| | Kernels | Batch-invariant, deterministic kernels on both sides | Slower sampling https://www.lmsys.org/blog/2025-09-22-sglang-deterministic/ , or kernels written for each model | | Precision | The same format on both sides | The trainer is tied to the samplers’ format, and no single low-precision compute format runs on every accelerator: A100 has no FP8 tensor cores, Hopper no FP4 | | Accelerator | Sampling and training on one accelerator generation | Rules out every other accelerator | | Policy staleness | Synchronous RL, one update per batch | Gives up the throughput of asynchronous RL | | Routing | Rollout routing replay https://arxiv.org/pdf/2510.11370 ; indexer replay | For MoE, a biased target policy https://arxiv.org/pdf/2512.01374 , plus memory and communication overhead; for the indexer, far more overhead https://arxiv.org/pdf/2602.15763v2 page=12 | | Top-p / top-k sampling | Top-p / top-k mask replay | Matches the candidate set, not the probabilities; mask storage and transfer | | Multi-turn context | Token-in-token-out https://www.lmsys.org/blog/2026-05-13-no-token-left-behind | Longer contexts if earlier reasoning is kept; more training compute if it is stripped | 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 https://arxiv.org/pdf/2610.11959v1 page=21 , and runs without staleness still correct the training–inference mismatch https://arxiv.org/pdf/2602.15763v2 page=11 . 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 https://arxiv.org/pdf/1802.01561 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 <sub θ</sub 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 <sub θ</sub : kernels, precision, accelerators and staleness all show up only as a different q . When tokens come from q instead of p <sub θ</sub , the policy-gradient update at each position picks up two errors https://arxiv.org/pdf/2609.20807 . 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 <sub θ</sub and grow otherwise. Importance sampling fixes both errors by weighting each token’s term by its ratio p <sub θ</sub / 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 https://arxiv.org/pdf/2512.01374 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 <sub θ</sub 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 https://arxiv.org/pdf/2507.18071 ’s and DeepSeek-V3.2 https://arxiv.org/pdf/2512.02556v1 page=8 ’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 https://modal.com/blog/flash-attention-4-faster , speculative decoding https://modal.com/blog/spec-is-all-u-need , 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 https://arxiv.org/pdf/2412.19437v2 page=14 or NVFP4 https://arxiv.org/pdf/2509.25149 , 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 https://modal.com/careers . 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.