RL on Heterogeneous Hardware - what if?

21 Oct, 2026

After looking through some works like The Art of Scaling Reinforcement Learning Compute for LLMs and RL at 1T which discuss in detail how does one go about post-training at scale, I was wondering about what it would take to post-train at scale on heterogeneous systems. A lot of the existing work on training-inference mismatch should carry over to such systems, but not all of it, so I went down the rabbit hole of the underlying system details that would affect the training reward and speed and efficiency of learning.

While I have no empirical results of my own so far, I have tried to dig more into how would the policy and the framework need to come together to close the gap between a trainer and samplers that run on different hardwares. As has been noted in works like Seer, RollPacker, and Jet-RL, Rollouts take anywhere between 65-90% of the RL budget (in wall-clock). And unlike Trainers, that are predominantly compute-bound, Rollouts can be a mix of both, with long-reasoning traces stressing the memory bandwidth like no other. A subtle yet evident example of this is in the AReal-Hex design where they prefer the H20 over H800s for rollouts since the former offers 4TB/s compared to 3.35TB/s of memory bandwidth but with only 15% of the H800's BF16 compute. The challenge in this and similar systems is the numerical mismatch this introduces. If you haven't already, I would recommend reading When Speed Kills Stability: Demystifying RL Collapse from the Training-Inference Mismatch to understand how catastrophic the numerical mismatch can be on the RL runs.

To preface, the trainer inference mismatch are not just a consequence of different systems. They exist even when run on the same GPUs. Something as little as a different decode and prefill kernel can and do contribute to the difference in output probabilities. Heterogeneous accelerators only make it worse. While there is sufficient work on exploring this heterogeneity across different Nvidia GPUs, none exists that cover the impact of bringing systems like the TPUs or d-matrix Corsair (heavily inference optimized) into the mix. So, most of the notes about heterogeneous hardware are more extrapolation than actual experiments.

Same weights, two policies

Take the clipped objective at the core of PPO and GRPO -

𝒥 (θ) = 𝔼_{o_{i} ~ π_{θ_{old}}} [\frac{1}{G} \sum_{i = 1}^{G} \frac{1}{| o_{i} |} \sum_{t = 1}^{| o_{i} |} ℓ_{i, t}]

with the per-token term

ℓ_{i, t} = min (r_{i, t} {\hat{A}}_{i, t}, clip (r_{i, t}, 1 - ε, 1 + ε) {\hat{A}}_{i, t})

and the ratio

r_{i, t} = \frac{π_{θ} (o_{i, t} ∣ q, o_{i, < t})}{π_{θ_{old}} (o_{i, t} ∣ q, o_{i, < t})}

$π_{θ_{old}}$ here appears in two places - one in the expectation of J, indicating the responses are sampled from $π_{θ_{old}}$ and the other in the ratio where it gives the probability of each sampled token. In an ideal world, they are the same policy. However, in a real RL system, they are computed by two different programs. The samples come from the inference engine - say, vLLM or SGLang, with its own kernels, batch sizes and parallelism and the denominator is usually recomputed by the trainer - Megatron or FSDP - on the same tokens with the same weights. And the two don't produce the same probability for the same token, which is why the trainer recomputes it. The GSPO paper states it directly: because of the precision differences between the engines, "we typically use the training engine to recompute the likelihoods of sampled responses under the old policy".

So the expectation is really over $π_{engine}$ , the policy that sampled the tokens, and the ratio is computed against $π_{train}$ - the policy that computes the gradient. As has been highlighted in the Thinking Machines blog, these numerical differences turn on-policy RL into off-policy RL. Overall, if the policy ignores this off-policyness due to the numerics mismatch, $π_{train}$ would be assumed to be the behavior policy $π_{engine}$ - which would immediately be incorrect due to the mismatch.

So, a standard correction is introduced in the form of truncated importance sampling and the objective function looks like -

𝒥 (θ) = 𝔼_{o_{i} ~ π_{engine}} [\frac{1}{G} \sum_{i = 1}^{G} \frac{1}{| o_{i} |} \sum_{t = 1}^{| o_{i} |} w_{i, t} ℓ_{i, t}]

with

w_{i, t} = min (\frac{π_{train} (o_{i, t} ∣ θ_{old})}{π_{engine} (o_{i, t} ∣ θ_{old})}, C)

Where does the gap come from? (Most of it is widely published.)

Floating-point addition isn't associative. Matmul, attention and RMSNorm are all reductions, and kernels split those reductions differently depending on the batch size (split-K, split-KV), parallelism configs and more and these differences cause changes in the outputs.
Traditionally different kernels. Training and inference frameworks use vastly different kernels based on the characteristics of the workload and in some cases also different underlying hardware.
BF16. This FP16 paper demonstrated how the sequence-level KL varies between the engines from 7.64 in BF16 to 0.32 in FP16 - reducing by almost 24x and the gap gets worse as response gets longer.
The LM head. MiniMax traced their gap to high-magnitude activations in the LM head and moved it to FP32 (MiniMax-M1).
MoE routing. MoE router are a huge source of discrepancy as even the slightest numerical differences could cause the tokens to choose a different expert. For instance, Qwen3-30B-A3B shows about 10% of router decisions differed between SGLang and Megatron, and 94% of tokens had a different expert in at least one layer (R3).
Low-precision rollouts, primarily for performance. FP8 or INT4 on the sampler with a BF16 trainer, because rollouts are where the time goes.

How big does the KL gap tend to get?

A good instance of this is in the MiniMax-M1 report where they initially observed the correlation between the trainer and inference to be around 0.987 which didn't grow the reward at all. On tracing it to high-magnitude activations and changing the LM head to FP32, the correlation improved to 0.997 and the model started training again. Similarly, for MoE models, a well-known technique is to replay the rollout’s expert choices. MoE models start with a bigger gap because tiny numerical differences flip which experts are selected for a token - the R3 paper measured an engine-trainer KL of 1.5e-3 for Qwen3-30B-A3B, against 6.4e-4 for the dense Qwen3-8B. Replaying the engine’s expert choices in the trainer brought it down to 7.5e-4, close to the dense model.

While I am a believer of not needing the trainer-inference to be 0KL and that the RL policy should account for the mismatch, a couple of things to remember is that any correction techniques like Truncated Importance Sampling (TIS) or Masked Importance Sampling (MIS) add bias to the gradient estimates by capping/dropping the importance weights. And if the estimator is unbiased, in the trainer/inference async (or mismatch), the variance can be unbounded. The only way to get both - no bias and no extra variance - is to eliminate the mismatch altogether. This is what NVIDIA and Inferact in this vLLM blog: the training forward calls vLLM's inference kernels and allows both sides to produce bitwise identical logprobs. Their DeepSeek-v4-Flash baseline also showed the KL exploded from 2e-3 to 29.9 with just routing replay and without correction techniques like TIS, MIS, or Score Centering, while their consistent run stayed at exactly 0 at every step. The catch is that it needs the same kernels on both sides, on the same hardware, at TP=1 and with one update per step.

So the gap is already there on a single accelerator. The question with mixed accelerators only builds on top of this about which fixes allow an RL-able platform.

The step going forward

Two kinds of things cross from the sampler to the trainer - numbers (the logprobs) and decisions (which token got sampled, which experts were picked, which tokens survived top-p). Replaying the decisions (R3) and making the numbers bitwise identical (NVIDIA and Inferact) are covered above. Beyond those, the fixes published over the last year or so fall into two groups - making the trainer round like the sampler, and correcting for what's left in the loss.

fix

Fig. 1: The sampler and the trainer hold the same weights but run different programs, so π_θold means two things - the policy that sampled the tokens, and the policy the trainer recomputes. Two kinds of things cross the boundary, numbers and decisions, and the weights come back. Dashed boxes are fixes that not every recipe uses - pick one to see which part of this picture it acts on.

Make the trainer rounding similar to the sampler

Low-precision rollouts are the standard way to make RL cheaper. Slime, a popular training framework for RL scaling, recommends the production path to be BF16 training with FP8 rollouts, and similarly VERL reports more than 35% faster rollouts from FP8 on a 30B MoE. However, as we have seen aplenty, this speed up doesn't come for free. Jet-RL found that BF16 training with FP8 rollouts "suffers from severe training instability and catastrophic accuracy collapse" under long-horizon rollouts, and even on runs that did converge, Llama3.1-8B's average score dropped by about 10 points. Qwen's QUADS saw NVFP4 rollouts with BF16 training collapse after roughly 150 steps.

One fix proposed in a series of works includes making the trainer's forward pass use the sampler's quantization, usually with fake quantization and a straight-through estimator for the backward -

Kimi K3 runs QAT through all of post-training: "During RL, rollout and training share the same quantization scheme — eliminating the train–inference mismatch."
MiMo-V2.6 serves MXFP4 experts in rollout and applies quantize-dequantize to the trainer's experts after every update, "so the two engines see identical expert weights".
Composer 2 runs the trainer's MoE forward in NVFP4 because "it is necessary that the trainer match the inference for numerical stability", and keeps the MoE backward in MXFP8.
Jet-RL uses the same FP8 precision flow for training and rollout, and gets a 16% end-to-end speedup over BF16.

A quirky result in this regard is that from Huawei's HiFloat4 paper (in simulation). It shows that raising the trainer back to high precision while keeping the rollout in FP4 gave worse accuracy than doing everything in FP4 - essentially, matching the sampler matters more than higher precision on the trainer.

However, something which I have observed personally and as confirmed in Nvidia's report: FP8-RL The MoE router needed to be maintained in at least BF16 instead of FP8. NeMo-RL always keeps the LM head out of FP8.

Fake quantization has a limit. It matches the format, not the kernels - the trainer's matmuls still accumulate in their own order. And even the weights may not come out identical. Miles' INT4 QAT docs note that the two paths are separate implementations, so matching the format "does not by itself establish bitwise-identical train and rollout weights".

Correct for it in the loss

We have seen techniques like Truncated Importance Sampling (TIS (Feng Yao, Liyuan Liu et al., Aug 2025)) and Masked Importance Sampling (MIS) for correcting the trainer-inference mismatch in the policy. MIS was introduced in When Speed Kills Stability which argues for sequence level importance sampling and drops the whole sequences that fall outside a trust region. Sequence-level masks have their own length problem but out of scope for today. Some more works that correct for this mismatch in policy -

IcePop (Ring-1T) zeroes out tokens whose ratio falls outside [0.5, 5] - about 1 to 2 tokens per thousand. GLM-5 uses it with [0.5, 2].
Off-policy sequence masking (DeepSeek-V3.2) drops negative-advantage sequences whose average log-ratio against the engine's probabilities is above a threshold.
NeMo-RL has seq_logprob_error_threshold, which drops any sequence whose mean $\exp (| Δ \log p |)$ is above the threshold - set to 2 in the Nemotron 3 Nano recipe.
Score centering (Together AI, Sep 2026) argues the instability comes from drift - "a persistent bias between training and inference engines that accumulates with every training step" - and centers the score function to remove it. With INT8 weights and activations and an INT4 KV cache, it reached 30% against 12% for TIS, for under 1% extra wall-clock.

A lot of this work has been discussed more from the perspective of similar architecture families deployed across training and inference. However, there are works like AReaL-Hex that run the rollouts on H20s and training on H800s - memory-bound decode on the bandwidth-heavy card, compute-bound training on the FLOPs-heavy one - and reports up to 1.5× the throughput at the same budget, or 1.46× lower cost at the same throughput”. INTELLECT-2 ran inference on permissionless workers, consumer GPUs included, and verified it with TOPLOC. Popular frameworks like Slime also support heterogeneous trainer-inference pairs: their docs mention there is no requirement for the training GPUs and rollout GPUs to be the same model, or even from the same vendor.

A GPU trainer with TPU samplers

Naturally, this brings to my next topic - what if we could put the samplers on TPUs and leave the trainers on GPUs (I am sure someone is doing this already?). Cooper et al. conducted experiments running the same models on A100, L40S, and H100, in BF16, and observed that for a median of about 85% of the problems, at least one pair of GPUs gave a different greedy output. So, obviously this means using a heterogeneous Trainer-Inference pair is not going to be a cakewalk. But at least, it's a good thought experiment about where the state of the frameworks today (as reported) and what would need to change to support this.

Potential Setup (deliberately convoluted)

Model: Qwen3-30B-A3B - 48 layers, 128 experts, 8 per token. R3 uses this and also seems like vLLM's TPU backend lists it as passing on both v6e and TPU7x.
Trainer: Megatron in BF16 on GB200s.
Pool A: vLLM on GB200 with FP8 rollouts
Pool B: vLLM on TPU7x (Ironwood) with FP8 rollouts, which Ironwood supports natively.

Every token carries the logprob from the pool that sampled it and the correction becomes per pool?

w_{i, t} = min (\frac{π_{train} (o_{i, t} ∣ θ_{old})}{μ_{k (i)} (o_{i, t} ∣ θ_{old})}, C_{k (i)})

where $k (i)$ is the pool that sampled response $i$ and $μ_{k}$ is that pool's sampling distribution.

Chips are trajectories, colored by the pool that sampled them.

Fig. 2: The setup above - two sampler pools on different hardware, one trainer. Each pool has its own sampling distribution μ_k, so every token carries the logprob from the machine that sampled it, along with the pool, the policy version and its expert IDs, and the correction can be set per pool.

The routes

We have already seen how vital it is for the router choice to match across the trainer and inference engines. On TPUs, as noted in the Tunix repository, a stock jnp.einsum picks a different MXU tiling dependent on the number of rows, M = B.T. So the router logits for the same token depend on primarily which pool (GB200 Rollouts or TPU rollouts) sampled it and the number of rows delivered. If I was to extrapolate R3's 10% different routing decisions from two engines on the same kind of GPUs, I wouldn't be surprised it's much higher across all the pools now.

vLLM has additional flags like MOE_APPROX_TOPK which would again impact the expert choice. It's also important to be mindful of the volume of data that crosses between the pools to the trainer - a quick estimate for one byte per expert ID and up to 256 experts = 48 x 8 = 384 bytes per token (layers x topk). This comes out to about 3.2GB of routes per step for 1024 rollouts of just 8k tokens. This is no joke.

Another thing that can come to byte here is the dtype and how the outputs are combined across all active experts. For instance, Nvidia (NeMo-RL) and Inferact pinned this exact thing to FP32 and topk-slot order to reach zero on just GB200s. Important to keep in mind.

trainer

Fig. 3: One token at one MoE layer. The 8th and 9th experts are nearly tied, so the TPU7x sampler picks a different expert than the trainer would. Replay makes the trainer run the sampler's experts, but the gate weights and the way outputs are summed stay its own. The scores are illustrative.

The sampling mask and logprobs

To compute its probabilities similar to the sampler, the trainer needs the set of tokens that survived the truncation on the sampler. On GPUs, vLLM and SGlang both have flags to enable this - --return-sampling-mask. Seems like the TPU vLLM backend has processed_logprobs - so $μ_{k}$ reflects the truncation but I didn't see an option to return the kept set. Obviously, you could set top-p 1.0 with top-k disabled (-1 in vLLM). Also find this interesting bug where vLLM on TPU was silently ignoring the temperature - More here.

As for the logprobs and the LM Head, this is another tricky territory. For instance, JAX docs show that the FP32 computation is done in BF16. I am sorry what? This is a bit perplexing but this should certainly lead to a mismatch on the logprobs and if MiniMax are to be believed, then RL rewards are sensitive to the LM head datatypes. This doesn't bode well with GB200 and TPU rollouts, both in the mix. Important for the correction factors in the policy account for this variation.

Correction and Staleness

I am not sure if any of the existing strategies, like ICEPop or others, alone would account for the mismatch across these different pools of hardware. I am wondering if the scheduler would need to actively discard rollouts from one pool or the another. For instance, if the TPU pool turns out to have worse numerics on the long rollouts assigned to it, these are discarded and regenerated on GPUs?

GB200 per-token ratio1.0005

TPU7x per-token ratio1.0010

cap C5.0

GB200 hits the cap at–tokens

TPU7x hits the cap at–tokens

masked, GB200––

masked, TPU7x––

Fig. 4: For corrections that multiply per-token ratios over the sequence, like MIS, a small per-token gap ρ grows as ρ^T and crosses the cap C at T = ln C / ln ρ. The pool with the worse numerics loses its longest rollouts first. The values are illustrative.

Oh, score centering would also need to adapt now - each pool would have its own and the trainer is going to see a weighted sum of them. The scheduler will play an important role here as the primary goal of cutting tail latency always remains a P0 while ensure training reward doesn't collapse.

This heterogeneous pool is certainly going to contribute to staleness differently as well (at least I would expect so). This might also lead to a staleness cap that is the lower of the two. The conventional mechanism of carrying the policy version with every token sampled and its replayed routes should continue to remain best practices.

The recipe, knob by knob

Knob	Single-hardware recipes today	GB200 + TPU7x (this thought experiment)
Behavior logprob	The engine's logprob ships with every token, the trainer recomputes its own, and TIS or MIS corrects the gap	The logprob from the pool that sampled the token, corrected against that pool's $μ_{k}$
MoE routing	Replay the engine's expert choices in the trainer (R3)	Replay from both pools - 384 bytes per token, about 3.2 GB per step for 1,024 rollouts of 8K tokens. `MOE_APPROX_TOPK` off. The expert-output combine still differs
Sampling mask	Return the kept set (`--return-sampling-mask` in vLLM) and apply it in training	GB200: the same. TPU7x: no kept set found, so top-p 1.0 with top-k disabled
LM head	FP32 LM head (MiniMax-M1)	JAX's default precision on TPU runs FP32 matmuls in BF16 - set the precision explicitly, and expect a residual gap
Trainer forward	Fake-quantize to the rollout format (Kimi K3, MiMo-V2.6, Composer 2, Jet-RL)	Both pools are FP8, but fake quantization matches the format, not the kernels
Loss correction	One TIS cap, IcePop band or sequence threshold	Probably per pool - or the scheduler discards the worse pool's long rollouts and regenerates them on GPU
Score centering	One persistent train-inference bias	One bias per pool, weighted by the scheduler's pool mix
Staleness	The policy version carried with every token	Per-pool staleness, with a cap set to the lower of the two

I think this is a good point to stop hypothesizing. There is a good blog on Async RL for dense models on TPUs from the Marin community here. So, hoping we will soon see an MoE run, on a more recent TPU.

Conclusion

To sum it up, the perpetual and hurtful trainer-inference mismatch - as tricky it is on homogeneous accelerators - should only get worse with heterogeneous hardware. This is irrespective of whether they are running the same policies or not. Closing that gap on homogeneous hardware cost about 10% of the throughput, at least for some models and layouts. With heterogeneity, my suspicion is it will be a tradeoff between the speedup offered by a new hardware and the impact on rewards/learning speed and efficiency due to trainer-inference mismatch and the gymnastics needed to reduce it.

Can someone run the MoE RL across different accelerators already?