返回资讯中心
外部精选
软件工程
#开发者实时流

Toward RL on mixed-precision, mixed-accelerator samplers

In large-model reinforcement learning (RL), samplers usually run on the same

Hacker News 最新torutofu2026/10/16 · 21 天前30 分钟阅读

以下正文同步自 Hacker News 最新,版权归原站所有,已转换为易读排版。

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.

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.

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.

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.

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

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.

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.

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.

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.

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

正文由 探How 从来源站点 RSS 同步,内容未经改写;遇到排版缺失或需要图片、视频时请以原文为准。

打开原文
去社区讨论 →