MAC END

Implementing Mixed-Precision RL with Diverse Accelerators for Enhanced Efficiency

Gate: Machine Learning
Departs: Oct 13, 2026
Crew: Nan Jiang
Passengers: 749

Research explores the use of mixed-precision, mixed-accelerator samplers in reinforcement learning, aiming for improved performance and cost-efficiency.

Implementing Mixed-Precision RL with Diverse Accelerators for Enhanced Efficiency

In large-model reinforcement learning (RL), samplers typically operate on the same accelerators and usually maintain a consistent precision as the trainer. This alignment helps stabilize training due to numerical proximity. However, outside of RL environments, systems are more flexible, often optimizing for speed and cost by utilizing low-precision formats such as FP8, MXFP4, and NVFP4. These formats can be paired with specialized kernels and accelerators specifically designed for long-context prefill tasks.

Sampling currently accounts for approximately 80%–85% of the GPU usage in extensive RL operations, indicating substantial optimization opportunities. The focus of this research is to determine if we can implement mixed-precision, mixed-accelerator samplers that employ different kernels while ensuring training remains stable and effective.

Samplers across multiple types of accelerators work together to support a single BF16 trainer on B200.

Identifying Challenges

Two significant challenges are evident in this approach:

The first challenge is related to learning. The integration of various optimizations can lead to mismatches between training and inference, resulting in samples that deviate from the optimal policy. This off-policy correction is a viable solution, as it allows a trainer to learn from data generated by different policies, like outdated samples in asynchronous training, without making assumptions about the discrepancies in the two policies.

The second challenge arises from the systems architecture. Traditional RL infrastructures often rely on homogeneous accelerator clusters, while mixed samplers necessitate a system of autoscaling accelerators capable of operating with their own precisions and kernels.

Through the combination of off-policy correction and our autoscaling, serverless GPUs, we demonstrate the feasibility of achieving stable, high-performance RL at a reduced cost using mixed-precision, mixed-accelerator samplers. Our experiments involve fully asynchronous training of an open-weight model in agentic coding tasks.

Experimental Setup

Our training focuses on the Qwen3.6-35B-A3B, engaging with agentic coding across a maximum of 256K tokens of context. We utilize the 2.7K repository tasks provided by Xiaomi released with the MiMo-V2.6, while ensuring the removal of the repository history to prevent data leakage. In every episode, the model constructs and submits a patch through a bash tool in a Modal Sandbox. The model earns a reward of 1 for submitting a passing patch and 0 for any failures. We evaluate the performance using SWE-bench Pro V2, which does not overlap with our training dataset, and a difficult subset termed HARD-51. Each checkpoint's performance is assessed in BF16 format.

For our RL framework, we employ Miles along with SGLang samplers and a Megatron-LM trainer that updates the policy in BF16 on B200 GPUs. This training setup operates fully asynchronously without any restrictions on staleness, utilizing 128 prompts with 8 trajectories per update. Each sampler operates with its own accelerator, precision, and kernels, while the trainer refrains from using their MoE routing decisions.

The samplers function within two configurations:

The mixed samplers encompass six types of accelerators: B200 and B300 utilizing NVFP4 weights, H100 and H200 using FP8 and BF16 weights, and A100 and RTX PRO 6000 serving BF16 weights. Each sampler employs its accelerator's specific kernels and maintains its Key-Value (KV) cache in FP8. A routing mechanism allocates each agent session to one of these groups based on their serving capacities, ensuring a balanced share of requests throughout training.
In contrast, the homogeneous samplers are confined to B200 with BF16 weights and a BF16 KV cache. They align with the trainer's accelerator and precision to ensure that only differences in kernels and staleness distinguish them.

Each group of samplers and its share of the requests served. — Each group of samplers and the share of requests they serve.

Hosted on Modal, each sampler group scales dynamically as required, ensuring they can serve the trainer's most recent weights in their respective precision formats.

Ensuring Stability in Mixed Sampler RL

Our starting approach utilizes the GRPO framework without corrections for inconsistencies. We implement three pivotal strategies:

Masked importance sampling (MIS) adjusts weights based on importance ratios. This involves multiplying each token's term by its importance ratio and masking tokens whose ratio lies outside a predetermined range, ensuring that the remaining tokens maintain expected updates in line with the training policy. If unmasked, rare tokens can significantly skew results.
Score centering (SC) mitigates drift. At every token position, it deducts the expected gradient of the log-probability under the sampling policy, maintaining that only differences in advantage influence the policy adjustments. While MIS addresses weighting discrepancies, it may inadvertently introduce drift.
Top-p sampling with mask replay (top-p) prevents the inclusion of rare tokens and standardizes support across both policies. The samplers draw from a finite top-p candidate set, ensuring rare tokens don’t feature in the sampled data. The trainer then renormalizes probabilities over this same candidate selection. Runs lacking this feature typically sample from the entire vocabulary.

Every token generated influences its associated update by adding its advantage multiplied by the gradient of its log-probability. By integrating MIS, SC, and top-p together, we achieve a system that reweights based on importance, centers the gradient, and uniformly restricts candidate sets across policies.

SWE-bench Pro V2 pass@4 over training on HARD-51 (left) and the full set (right). — Performance of SWE-bench Pro V2 pass@4 during training on HARD-51 and the full dataset.

Experiments indicate that without these corrections, the training collapses across both setups—GRPO quickly fails in the mixed sampler configuration, while the homogeneous approach sustains longer before eventual failure. Notably, combinations of SC and MIS, whether or not coupled with top-p, manage to maintain stability, resulting in pass@4 scores on HARD-51 that peak at approximately 65% to 75%, whereas the entire dataset rises from 89% to around 94%.

Training reward, sampler–trainer KL, trajectory length and episodes ending on repeated format errors, for every run. — Visualization of training reward, KL divergence between sampler and trainer, trajectory lengths, and the frequency of episodes terminating on repeated format errors.

Collapsed episodes typically exhibit distinctive behaviors: training rewards decline, sampler–trainer KL diverges significantly, and episodes frequently conclude with repetitive format errors. No consistent length patterns emerge for trajectories, though some runs finish early while others might continue long but with repetitive tasks. Notably, stable runs do not display these negative signs. The homogeneous samplers’ GRPO, despite starting at a lower KL divergence, remains subject to collapse, following similar trajectories as the mixed samplers.

Comparison of healthy and collapsed episodes. — Visualization showing the differences between successful and collapsed episodes.

In collapsed episodes, malformed commands are rarely produced. Instead, the model tends to stop calling its tool midway through tasks, leading to an episode end once there are three consecutive turns without a tool call. For instance, in the GRPO’s step-40 checkpoint using mixed samplers, around 62% of evaluation episodes end prematurely, averaging just 22 turns in contrast to the 57 turns for successful episodes, with 98% of the last replies being empty.

Cost Efficiency of Mixed Samplers

Training reward, sampler–trainer KL, gradient norm and entropy for SC + MIS + top-p on the mixed and homogeneous samplers. — Comparison of training reward, KL divergence, gradient norm, and entropy between mixed and homogeneous samplers.

Utilizing the combination of SC, MIS, and top-p, the mixed samplers achieve SWE-bench Pro V2 scores that are comparable to those obtained by homogeneous samplers, with pass@4 rates remaining closely aligned throughout training. Additionally, metrics like training rewards and gradient norms reflect similar trajectories. Notably, this occurs despite a 4.7× greater average KL divergence compared to homogeneous setups. Until approximately step 160, entropic behaviors remain consistent, after which they diverge, with entropy increasing in homogeneous setups while remaining stable in mixed samplers.

SWE-bench Pro V2 pass@4 against cumulative sampling cost for SC + MIS + top-p on HARD-51 (left) and the full set (right). — Evaluation of SWE-bench Pro V2 pass@4 performance relative to cumulative sampling cost.

Decoupling samplers from the trainer has allowed for significant cost reductions during the training phase. Throughout the majority of the training process, mixed samplers demonstrated an approximately 15% reduction in cost per update compared to their homogeneous counterparts. It's essential to recognize that while our experimental setup primarily targeted showcasing the potential of mixed-precision, mixed-accelerator samplers, ample opportunities for cost savings exist; for example, each sampler can be optimized individually, with the possibility of shifting the mix to prioritize lower-cost options.

Understanding Training-Inference Mismatch

Mismatched token probabilities during sampling can lead to discrepancies between the outcomes from the sampler and those from the trainer. These discrepancies stem from multiple potential sources within the RL system:

The mixed samplers inherently comprise several of these mismatch sources, resulting in an average mismatch that is 4.7 times greater than that experienced by homogeneous samplers.

Sampler–trainer KL for each group of samplers, averaged over training with SC + MIS + top-p. — KL divergence between the sampler and the trainer, averaged over training using SC + MIS + top-p.

On the Inefficiency of True On-Policy RL

True on-policy RL exists when each token is sampled directly from the policy undergoing training. When samplers and the trainer operate as distinct systems, adhering to true on-policy requires two critical conditions: samplers must compute identical probabilities to the trainer (essentially achieving true on-policy RL), and every sample must originate from the trainer’s current weight configuration.

Instances of bitwise-accurate sampling and training exist. Although demonstrated effectively for smaller models, open RL frameworks have managed to achieve near-identical outputs even for models like GLM-5. Moreover, available frameworks have incorporated batch-invariant and deterministic kernels to ensure both training and inference yield consistent results. However, pursuing absolute alignment for every source incurs certain costs:

Source	Examples	Effect
Kernels	Differences in inference and training kernels; lack of batch invariance	Different outputs for identical inputs, even at the same precision
Precision	BF16; FP8; MXFP4; NVFP4	Variations in precision (mantissa bits) and dynamic range (exponent bits, block scaling), leading to divergent rounding errors; mismatch can escalate as computations increasingly utilize quantization
Accelerator	Architecture-specific kernels and number formats: absence of FP8 tensor cores on A100 or FP4 tensor cores on Hopper	Variation in mismatch across different accelerators utilizing identical code
Policy Staleness	Asynchronous RL frameworks wherein samplers lag behind training steps; partial rollouts; in-flight weight updates	Escalating mismatch correlated with policy staleness
Routing	Top-k expert selection; inconsistencies with the sparse-attention indexer’s top-k selection	Router inconsistency: two systems might select different experts, or different key-value entries, for a single token
Top-p / top-k Sampling	Tokens are selected exclusively from top-p or top-k candidates	Discrepancies in action spaces that can destabilize training
Multi-turn Context	Detokenizing and retokenizing processes; chat templates that reformat previous turns or obscure reasoning	Training scores may evaluate tokens in contexts that the model has not seen

Many of these sources affect solely token probabilities, which can be corrected. However, challenges arising from multi-turn contexts and top-p / top-k sampling disrupt cohesion and support, thereby complicating the correction process. To address this, our experiments keep token-by-token evaluation aligned with a token-in-token-out strategy, while implementing mask replay on top-p outputs for synchronization.

In practical settings, large-scale runs often attempt to align some mismatch sources while using corrections for others. For example, setups that maintain alignment across precision, routing, and top-p masking still impose limits on staleness, while asynchronous runs continue to rely on off-policy corrections. A common belief persists that a minimized mismatch enhances training stability, leading many to keep samplers numerically close to trainers. However, what truly matters is whether this mismatch remains stable rather than merely its magnitude. Without proper correction mechanisms, a bias in this mismatch can gradually widen until training fails. Interestingly, models that achieve a smaller mismatch also follow suit with eventual collapse—GRPO utilizing homogeneous samplers initially exhibited lower KL divergence but succumbed later in the training process. Implementing correction allows training on mixed samplers to maintain stability and achieve performance levels akin to those of homogeneous samplers, despite the more notable mismatch.

Training reward and sampler–trainer KL for GRPO and SC + MIS + top-p, across both mixed and homogeneous samplers. — Contrasting training rewards and KL values for GRPO alongside SC + MIS + top-p, comparing mixed and homogeneous configurations.

It's pivotal to differentiate the sources of mismatch magnitude from those contributing to its growth. The system's design governs the magnitude, including factors such as the accelerator's type, precision, and kernels in use. Meanwhile, update procedures dictate how efficiently this mismatch can be addressed, and existing off-policy corrections can still apply to these components. The objective isn’t to minimize mismatch per se, but rather to constrain its expansion.

This approach is why we don't pursue strict true on-policy RL. Given that the samplers and trainer are structured for distinct purposes, enforcing excessive similarity limits individual optimization efforts. By utilizing correction mechanisms, both components can improve independently.

Mechanics of Off-Policy Correction

Learning from data generated by an alternate policy is a well-established concern in RL. Existing models like IMPALA have successfully decoupled the actors generating experience from the learner, allowing for corrections related to policy lag rather than executing synchronous operations. The sampling policy q generates the tokens, while the training policy p_θ is trained using these tokens. In configurations with mixed samplers, each token is produced by designated samplers at various weight settings, enabling the recording of token probabilities for each q. Thus, the correction process does not need to ascertain why q diverges from p_θ: discrepancies arise naturally from factors like differing kernels, precision, accelerator types, and staleness that affect q.

When observing tokens generated according to q rather than p_θ, the policy-gradient update at each position tends to introduce two significant errors:

Drift. Under the training policy, the average gradient of the log-probability at each position should average to zero. However, under q, this average might incline toward an update in the direction of q. Consequently, updates can asymmetrically shift the training policy, propelling it towards the sampling policy when expected advantages are positive and against it when they are negative. Normalizing advantages across samples does not fully eliminate this drift, as the expected advantage derived from a specific prefix will generally deviate from zero. When the samplers subsequently update the trainer's weights, they reiterate this error, causing continuous adjustments toward a target that is also mutable.

Incorrect weighting. The updating process attributes each token with a weight reflective of its sampling frequency as opposed to the training policy's probability, which may lead to a token receiving a disproportionately high gradient after its likelihood of occurrence has diminished.

These errors converge towards zero when q matches p_θ but can escalate otherwise.

To address these issues, importance sampling comes into play, adjusting each token's contribution by weighing it with the ratio p_θ / q, conditional upon samplers having access to similar tokens as the training policy. When conducted effectively, weighting over a full trajectory using the product of ratios yields precise results, but its variance can sharply increase with trajectory length; thus, larger setups employ first-order approximations, applying the ratio incrementally for each token. However, rare tokens sampled from q can acquire weights as substantial as 1 / q. Truncating or masking these weights allows for variance management, though it can also reintroduce drift, which MIS addresses by ensuring ratios remain within specified boundaries, setting them to zero beyond certain thresholds.

Truncating and masking the importance ratio. The dashed line is the unbounded weight. — Illustration of truncating and masking the importance ratio. Dashed line indicates the unbounded weight.

Score centering, conversely, tackles the drift problem by subtracting the anticipated gradient of the log-probability from q at each token position. This adjustment ensures that a constant advantage does not shift the policy during updates. While the wrong weighting persists, as it remains unaltered by SC, this necessitates the pairing of SC with MIS: centering effectively neutralizes drift irrespective of weight variances, including those introduced through masking. Together, SC and MIS eliminate drift while reweighting the retained tokens based on the training policy's guidelines. Subsequently, for tokens not produced frequently, q probabilities are approximated using the trainer’s data, rescaled back to correspond with the remaining probability mass of the samplers.

Visual comparison demonstrating the effect of drift without correction (left) versus the influence of score centering (right). Adapted from Martin Marek's animation.

The implementation of top-p sampling modifies the support shared between the policies, and reintroducing mask replay restores this consistency. By sampling solely from a top-p candidate set, tokens outside this group are assigned zero probability under q. As a result, reintegrating this mask within the trainer allows for a preservation of identical probabilities across the same subset. Since the rarest tokens, with potential for largest ratios, are not included in the sampled data and probabilities are accurately recorded for all candidates, score centering becomes more reliable.

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. — Trainer's probabilities comparison with and without mask replay. Tokens outside the shaded area are excluded from the samplers' top-p set; dashed outlines depict probabilities pre-renormalization.

Insights from Our Runs

Share of tokens with an importance ratio outside [0.2, 5], sampler–trainer KL, gradient norm, and entropy for each method on mixed samplers. — Visualization of token shares with an importance ratio outside [0.2, 5], KL divergence, gradient norm, and entropy across different methods in mixed samplers.

GRPO retains both errors, with its gradient norm escalating approximately 20 times shortly after the training phase commences.
GRPO + MIS rectifies the weighting issues, yet fails to mitigate the drift introduced by its masking method. Consequently, while the gradient norm remains manageable, the model’s entropy diminishes as KL divergence jumps sharply.
SC is effective at removing drift but does not correct the weighting error, leading to entropy declines and significant KL growth before the run's collapse.
GRPO + top-p effectively filters out rare tokens; despite GRPO experiencing rapid growth in its gradient norm, this method retains smaller variation but eventually experiences KL increases and subsequent failure.
SC + MIS effectively addresses both drift and incorrect weighting, enabling stable performance with only brief spikes in gradient norms, while KL increases moderately.
SC + MIS + top-p serves as the optimal method, yielding the lowest KL divergence alongside stable gradient norms and entropic behaviors.

Through the use of SC + MIS + top-p, training continues effectively, even with a mismatch average 4.7 times higher than that of the homogeneous samplers. In instances of collapse, the mismatch shows an increase of threefold or more.

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. — Comparison of each sampler group’s share of training samples and of tokens with an importance ratio beyond [0.2, 5] using SC + MIS + top-p.

Within each sampler group, the correction's impact tends to align closely with the mismatch. The ratio of extreme-ratio tokens within a group correlates more with its KL divergence than with the sample allocation: for instance, H100 FP8 contributes 10% of samples but is responsible for 29% of tokens, whereas B200 and B300 NVFP4 account for 47% of samples yet only 16% of extreme-ratio tokens. Notably, even H100 FP8 has just around 0.3% of its tokens classified outside the [0.2, 5] range. The associated correction evaluates tokens based on their individual ratios, thereby ensuring stability during training among all sampler groups.

Limitations of the Current Approach

While MIS, SC, and top-p methods operate on a token-by-token basis, each inherits limitations associated with per-token approximations: the approach’s prefixes and advantages stem from the samplers. Sequence-level corrections—including strategies like GSPO and DeepSeek-V3.2—alongside limiting staleness can guard against excessively large errors. Our experiments did not employ these strategies yet maintained stability, navigating contexts spanning up to 256K tokens with an average staleness ratio near 4.

Additionally, top-p sampling combined with mask replay alters the nature of trained outputs: updates are optimized exclusively for the top-p candidate pool, and tokens positioned outside this group do not receive gradients. For evaluation purposes, the full vocabulary was utilized for sampling during testing phases, resulting in comparable performances between SC + MIS + top-p and SC + MIS methods, both yielding solid evaluations.

The established corrections change only the token usage during updates without altering the trajectories generated by samplers. This means that the selection of precision or accelerator type remains a balancing act, more about throughput versus sample efficiency rather than a simple throughput-stability comparison. Though our proposed corrections are not exhaustive, they share a fundamental principle: maintaining mismatch at significant levels is permissible as long as the update mechanisms remain closely aligned with on-policy methodologies, ensuring a bounded mismatch.

Future Directions

This correction mechanism enables distinct separation between samplers and trainers, fundamentally reshaping the capacities and opportunities for RL implementations. As a result, sampling systems no longer demand large clusters of homogeneous setups; rather, any suitable accelerator able to serve the model can engage in sample generation, permitting the size of the sampling fleet to evolve independently from the trainer's requirements and facilitating dynamic scaling during runs (e.g., increasing sampling need as trajectory lengths extend). Moreover, inference optimizations transfer seamlessly into RL settings: faster kernels, speculative decoding, lowered precision processing, and the bifurcation of prefill and decoding across various accelerators can enhance sampling efforts without requiring a corresponding adjustment from the trainer. The RL trainer can also adapt to the most efficient accelerators and its own preferred lower-precision formats, such as FP8 or NVFP4, without requiring the samplers to realign with its choices. Given the ongoing diversification of accelerators across vendors and generations, this decoupling becomes essential for leveraging advanced training methodologies.

We are actively addressing complex challenges at the intersection of RL algorithms and large-scale systems. Stakeholders interested in these advancements can join our team and contribute to this rapidly evolving field.

Acknowledgments

We extend our gratitude to the Miles team, Luke Huang, Martin Marek, Dhruv Gautam, Xiuyu Li, Byron Hsu, Ziang Li, and Zilin Zhu for their insightful discussions and contributions to the development of this post.

Source	How it is matched	Cost
Kernels	Use batch-invariant, deterministic kernels across both platforms	Resulting in slower sampling, or the need for model-specific kernel development
Precision	Maintain identical formats on both ends	The trainer must be constrained to the samplers’ format, with no single low-precision compute format universally applicable across all accelerators, such as the absence of FP8 tensor cores on A100 or FP4 on Hopper
Accelerator	Sampling and training conducted on the same generation of accelerator	Limits the use of all other accelerators
Policy Staleness	Synchronous RL, with only one update per batch	Results in throughput trade-offs against asynchronous RL
Routing	Utilization of routing replay for rollouts	Would necessitate biased target policies, alongside memory and communication overheads; for the indexer, this incurs substantial overheads
Top-p / top-k sampling	Employing top-p / top-k mask replay	While this aligns the candidate sets, it does not synchronize probabilities; entails storage and transfer overheads for mask data
Multi-turn context	Implementation of token-in-token-out method	Requires longer contexts if the earlier reasoning is preserved; entails higher computational requirements if omitted