TRACE makes FP4 reinforcement learning stable for Mixture-of-Experts LLMs by letting the low-precision rollout path dictate how the training path rounds its activations, closing the policy mismatch that normally wrecks aggressive 4-bit RL and delivering up to 5.4× rollout speedup versus BF16.
When you post-train an LLM with RL, most of the wall-clock time is spent in rollout: the model generates long trajectories, which get scored and fed back as gradient signal. People have been shrinking that cost by running the rollout in low precision (FP8, FP4) while keeping training in BF16. The catch: the training forward pass and the rollout forward pass now compute slightly different numbers for the same token, so the “policy” that acted and the “policy” being updated disagree. In Group Relative Policy Optimization (GRPO)-style RL, that mismatch shows up as a log-probability gap, and if it grows, training collapses.
This is especially bad for Mixture of Experts models. A tiny numerical nudge can flip which expert the router picks, and now you aren’t just comparing slightly different activations, you’re comparing entirely different subnetworks. Prior FP4 RL work (QUADS, QaRL, Rollout-ResQ) tried to fix this by making each path’s quantization more accurate with respect to the original BF16 values. The TRACE authors argue that’s the wrong target: two paths can each have tiny per-path quantization error and still land on different FP4 codewords, which is what actually matters for policy agreement.
The core observation is a rounding-boundary problem. The BF16 training activation and the BF16 rollout activation might differ by only 0.48, but if that difference straddles the midpoint between two FP4 codewords, one gets rounded up and the other down, and the FP4 gap balloons (the paper shows an example going from 0.48 to 24). Minimizing each path’s rounding error independently doesn’t help, because the two paths can round in opposite directions.
TRACE’s fix is rollout-guided rounding. During rollout generation, the system records which FP4 codeword each quantization site actually produced. Later, during the Quantization-Aware Training (QAT) forward pass on the training side, instead of standard round-to-nearest, TRACE picks whichever of the two neighboring FP4 codewords is closer to what rollout picked. By construction, the training-side quantized value cannot be farther from the rollout-side value than round-to-nearest would be. It doesn’t prove the full-network policy divergence shrinks, but it eliminates the “amplification by inconsistent rounding” source of discrepancy.
The practical obstacle is bandwidth. Logging full FP4 codes for every activation and every KV entry across a long rollout is huge: the paper estimates about 51 TB per RL step for one of the Qwen models, which blows past storage I/O budgets. The fix is based on the empirical observation that >99% of mismatched sites differ by only one adjacent FP4 codebook entry, so you don’t need the full code to disambiguate. TRACE stores only one mantissa bit plus scale metadata, and only for the deeper half of the layers (where the rounding corrections concentrate), compressing the cache by roughly 6×.
# During rollout: record compact code per site
for site in capture_sites_in_later_layers:
q_roll = fp4_quantize(x_rollout, scale_roll)
cache[traj, tok, site] = (low_bits(q_roll, b=1), scale_roll)
# During QAT training forward on paired activation x_train:
c, s = cache[traj, tok, site]
z = clip(x_train / s, -6, 6)
q_roll_hat = reconstruct_from_mantissa(c, z) # compact rollout ref
q_minus, q_plus = fp4_bracket(z) # two candidate codewords
q_star = argmin(abs(q - q_roll_hat) for q in [q_minus, q_plus])
x_quantized = straight_through_estimator(x_train, s * q_star)
Across four Qwen3 MoE models (ranging from a 35B total / 3B active model up to a 2.4T total / 95B active model) and reasoning, coding, and long-horizon RL tasks, the headline is that TRACE under joint FP4 weights/activations plus FP4 KV cache tracks the BF16 rollout reference, while prior FP4 methods lose several points.
•
On Qwen3.5-35B-A3B reasoning (average over LiveCodeBench, AIME24, AIME25, HMMT25): TRACE 75.3 vs BF16 74.9, vs QUADS 68.8, vs vanilla QAT 59.6. Biggest single gain on HMMT25 (+11 points over QUADS).
•
On larger models: TRACE essentially matches BF16 on DeepSWE (33.0 vs 33.4), on Terminal-Bench (70.6 vs 68.8), and on GDPval (90.2 vs 90.3), where plain QAT drops 2–5 points.
•
Training dynamics: vanilla QAT’s train–rollout log-prob gap grows during training and eventually crashes test scores; TRACE keeps that gap close to the BF16 reference across all 400 steps.
•
A notable qualitative finding: TRACE starts below the BF16 reference (quantization still hurts initially) but catches up around step 120 on HMMT25 and sometimes exceeds BF16 later. The authors interpret this as the policy adapting to the low-precision execution environment once the discrepancy is controlled.
•
Efficiency: rollout decoding up to 5.4× faster than BF16 at 128K output length; end-to-end RL step overhead vs vanilla FP4 rollout is only 7.4%.
•
Ablations: keeping only 1 mantissa bit for the latter 20 of 40 layers is close to full-information guidance (75.3 vs 75.8 average); dropping to 5 layers loses about 2 points. The method also carries over to MXFP4 W4A4 and W4A8 settings.
•
Versus post-hoc FP4 quantization of a BF16-trained policy (vanilla NVFP4, 4over6, H-Scale): all three sit around 70–71 average, TRACE at 75.3. Training the policy under FP4 beats quantizing it after the fact.
The comparisons are against the baselines the authors ran under the same stack (verl + Megatron + SGLang, GRPO with routing replay). The log-prob diagnostic is the clearest causal story: TRACE is the only method that keeps that gap small, and that correlates tightly with downstream scores not collapsing.
•
If you’re running RL on a MoE model and the rollout is your bottleneck, the paper is strong evidence that joint FP4 weights/activations/KV is viable, not just FP8 or weight-only FP4. Worth testing on your own stack if you can modify both the rollout engine and the trainer, since TRACE requires coordinated logging between them.
•
The rollout-guided-rounding idea is more general than FP4. Any time you have two execution paths of the same model at different precisions and need their outputs to agree, having the lower-precision path emit a hint for the higher-precision path’s rounding is a cheap intervention worth considering. The paper only evaluates it for RL rollout vs QAT training.
•
If you can’t modify training internals and only run post-hoc quantization through a serving stack, this result does not apply. The gain comes from co-training with FP4 in the loop; the post-hoc baselines in the paper are 4–5 points worse on the same model.
•
The compact-caching trick (1 mantissa bit, deeper layers only) is the practical enabler. If you reproduce this and skip the compression, expect decoding throughput to tank, which the paper shows explicitly for the full-information variant.
•
The paper does not link a public code release in the supplied text, so reimplementation requires engineering against VeRL/Megatron/SGLang yourself.
•
All evaluated models are Qwen MoE variants, and the rollout engine is SGLang on GB200 hardware. Generalization to dense models, other MoE routers, or non-Nvidia FP4 formats isn’t established, though the MXFP4 ablation is a positive signal.
•
The local rounding guarantee is per-site. The paper is explicit that it does not prove the full-network discrepancy or the policy KL shrinks monotonically; the empirical log-prob trace is the actual evidence.
•
TRACE doesn’t fix discrepancies caused by policy staleness (training policy drifting ahead of the rollout policy in async RL). It only removes the extra gap introduced by inconsistent FP4 rounding on top of whatever BF16-level drift exists.
•
Required infrastructure is nontrivial: you need the rollout engine to emit per-site quantization records, a cache tier that holds them, and a training engine that reads them back and overrides its own rounding. This is not a drop-in change to an existing RL pipeline.
•
Only routed experts and KV cache are quantized to FP4; attention and non-expert modules stay BF16. The claim is about this specific hybrid, not fully-FP4 execution.