TRACE: Rollout-Guided Quantization-Aware Training for FP4 Reinforcement Learning of MoE Language Models
Xin Wang, Hao Yu, Zhengyang Zhuge, Bochao Mao, Zheng Li, Junda Feng, Yuyan Luo, Yi Zhang, Yizhong Cao, Mi Zhang, Dayiheng Liu, Jianwei Zhang
TRACE aligns FP4 training and rollout paths by using rollout-side quantization outcomes to guide training-side rounding.
How can we perform reinforcement learning on large MoE language models using FP4 precision for rollouts without the performance degradation caused by train-rollout quantization discrepancies?
Reinforcement learning for large language models requires generating long trajectories, but aggressive FP4 quantization creates numerical discrepancies between the training and rollout paths that destabilize policy updates. TRACE (Train-Rollout Quantization Alignment via Compact GuidancE) addresses this by recording rollout-side FP4 quantization outcomes and using them to guide the rounding decisions during quantization-aware training. This alignment enables joint FP4 weight and KV-cache rollout that matches BF16 performance while achieving up to 5.4× speedup in rollout generation.
Paper Primer
The core mechanism of TRACE is rollout-guided rounding: instead of standard round-to-nearest, the training engine selects the FP4 codeword that minimizes the distance to the specific codeword produced during the rollout generation. To manage the massive data volume of these records, TRACE caches only the mantissa and scale information from the latter half of the model layers, which is sufficient to identify the correct rounding direction for over 99% of mismatched values.
TRACE enables joint FP4 weight and KV-cache rollout with performance comparable to BF16 rollout.
Evaluations across reasoning, coding, and long-horizon RL tasks on four MoE models (e.g., Qwen3.5-35B-A3B, Qwen3.8-2.4T-A95B) show consistent score improvements over existing FP4 RL methods. Achieves up to 5.4× rollout speedup compared to BF16 rollout while maintaining stable training dynamics.
TRACE significantly reduces train-rollout discrepancy compared to existing quantization-aware training (QAT) methods.
Training dynamics analysis shows that while vanilla QAT and other baselines exhibit increasing log-probability divergence and training instability, TRACE maintains alignment close to the BF16 reference throughout the training process. Improves average benchmark scores by up to 11.0 points on specific reasoning tasks compared to baseline QAT.
Why does minimizing quantization error on each path independently fail to stabilize RL training?
Minimizing per-path error does not minimize the cross-path discrepancy; small differences in high-precision activations can be amplified when rounded to different FP4 codewords, leading to policy mismatch and training collapse.
What is the primary trade-off of using rollout-guided quantization?
The primary cost is the storage and communication overhead of the rollout-side quantization records, which TRACE mitigates by using a compact mantissa-only caching scheme that selectively retains information from deeper layers.
By directly aligning the quantized execution paths rather than just improving individual path accuracy, TRACE allows practitioners to deploy aggressive FP4 rollout for MoE models without sacrificing the stability or performance of high-precision training.
The FP4 Rollout Bottleneck
We expose the train‑rollout discrepancy that hampers FP4 RL and introduce TRACE to align them.
Reinforcement learning (RL) for large language models (LLMs) demands heavy computation and memory because each training step repeatedly generates long rollout trajectories. Low‑precision FP4 rollout promises speed and memory savings, yet the coarse quantization introduces a substantial train‑rollout discrepancy that destabilizes policy learning. Existing FP4 RL methods improve quantization accuracy on training and rollout paths separately, but they do not directly minimize this discrepancy.
The policy’s behavior diverges when the same model is run in high‑precision training mode versus low‑precision FP4 rollout mode, leading to mismatched actions and unstable RL updates.
How does Train‑Rollout Discrepancy differ from the usual train‑inference mismatch?
Train‑inference mismatch typically refers to differences between training‑time and inference‑time precision (e.g., BF16 vs. FP16). Train‑Rollout Discrepancy specifically concerns the divergence introduced when the rollout generation itself is quantized to FP4, affecting not only forward activations but also the KV‑cache and expert routing, which directly feeds back into the RL update loop.
**Figure 1.** Training dynamics of Qwen3.8-Flash-Next under FP4 rollout. Conventional FP4 QAT exhibits increasing train-rollout discrepancy and degraded test score on Terminal-Bench 2.1, while TRACE maintains stable alignment and enables the policy to progressively adapt to FP4 rollout.
The fundamental conflict is that FP4 offers efficiency but threatens RL policy stability via train‑rollout discrepancy.
The TRACE Framework
TRACE aligns quantized rollout and training paths by guiding rounding decisions with cached rollout information.
Low‑precision FP4 rollout introduces a numerical gap between the rollout execution path and the training‑time forward pass, and naïve quantization‑aware training (QAT) does not close this gap, causing policy divergence or instability.
TRACE records the exact FP4 rounding outcomes seen during rollout and then uses that record to steer the rounding decisions made while training, so the two paths stay in lockstep.
How does TRACE differ from the standard round‑to‑nearest (RTN) used in ordinary QAT?
Standard QAT rounds each activation to the nearest FP4 codeword independently of rollout outcomes. TRACE instead looks up the exact rollout‑side codeword and selects the training‑side codeword (among the two nearest candidates) that minimizes the distance to that rollout codeword, guaranteeing no larger rounding error than RTN.
Normalize both values: 60.24 → 2.51, 59.76 → 2.49.
Round to nearest FP4 codewords: 2.51 → 3, 2.49 → 2.
Compute FP4 discrepancy: $|3 - 2| = 1$ codeword, which corresponds to a numeric gap of roughly 24 after de‑scaling.
BF16 discrepancy before quantization was $|2.51 - 2.49| = 0.02$, i.e. a 0.48 raw difference.
This shows how a sub‑unit BF16 difference can be amplified by a factor of ~50 after FP4 rounding, motivating the need for aligned rounding decisions.
For each activation, TRACE picks the FP4 codeword that is closest to the rollout‑side codeword, ensuring the training rounding never deviates more than the naïve round‑to‑nearest would.
Compute distances: $|2.5 - 2.8| = 0.3$, $|3.0 - 2.8| = 0.2$.
Select the smaller distance: $q_{\text{TRACE}} = 3.0$.
Standard RTN would have chosen $q_{\text{RTN}} = 3.0$ as well, but if $q_{\text{rollout}}$ were $2.6$, RTN would pick $3.0$ while TRACE would pick $2.5$, reducing the error.
TRACE adapts the rounding direction to the rollout outcome, preventing the amplification that occurs when the two sides fall on opposite sides of a rounding boundary.
**Figure 2.** Overview of TRACE for FP4 RL (FP4 W / A + FP4 KV) Training of MoE Models .
**Figure 3.** Two Examples of Train–Rollout Discrepancy under Different FP4 RL Methods.
**Algorithm 1** TRACE: Rollout-Guided FP4 QAT **Require:** Policy parameters $W$, prompts $\mathcal{D}$, reward function $r$, captured sites $\mathcal{S}$, bits $b = 1$ **Ensure:** Updated policy parameters $W$ 1: **for** each RL iteration **do** 2: $\quad$ Generate trajectories $\mathcal{T}$ and log-probabilities $\ell_{\text{roll}}$ from $\mathcal{D}$ using $Q_{\text{FP4}}(W)$ 3: $\quad$ During rollout, record $\mathcal{R} = \{(\text{Pack}(\text{low}_b(q_{\text{roll}})), s_{\text{roll}})\}$ at $\mathcal{S}$ 4: $\quad$ Compute advantages $A$ from rewards $r(\mathcal{T})$ 5: $\quad$ **for** each policy-update minibatch $\mathcal{B} \subseteq \mathcal{T}$ **do** 6: $\quad\quad$ **for** each captured activation $x$ encountered during the QAT forward **do** 7: $\quad\quad\quad (c, s) \leftarrow \text{UNPACK}(\mathcal{R}[\text{trajectory, token, site}])$ 8: $\quad\quad\quad z \leftarrow \text{clip}(x/s, -6, 6)$ $\quad \triangleright$ Normalize with the rollout scale 9: $\quad\quad\quad \hat{q}_{\text{roll}} \leftarrow \text{RECONSTRUCT}(c, b, z)$ 10: $\quad\quad\quad (q_-, q_+) \leftarrow \text{BRACKETFP4}(z)$ 11: $\quad\quad\quad q^* \leftarrow \text{arg min}_{q \in \{q_-, q_+\}} |q - \hat{q}_{\text{roll}}|$ 12: $\quad\quad\quad$ Use $\text{STE}(x, sq^*)$ in the remaining forward computation 13: $\quad\quad$ **end for** 14: $\quad\quad$ Compute $\mathcal{L}_{\text{RL}}$ from the QAT log-probabilities, $\ell_{\text{roll}}$, and $A$ on $\mathcal{B}$ 15: $\quad\quad W \leftarrow \text{OPTIMIZERSTEP}(W, \nabla_W \mathcal{L}_{\text{RL}})$ 16: $\quad$ **end for** 17: **end for** 18: **return** $W$
Optimizing Train-Rollout Communication
Mantissa-Only Communication cuts rollout data transfer by sending only essential mantissa and scale info.
Preserving the full rollout‑side quantization information inflates data movement to impractical levels. For a 35 B‑parameter model ($D_1\!=\!2048$, $D_2\!=\!256$, 40 MoE layers) each generated token yields ~45 KB of activation guidance and ~5.6 KB of KV guidance, totaling ~50 KB. With 4,096 trajectories per RL step this reaches ~51 TB, turning storage I/O into the dominant bottleneck.
**Figure 5.** Data Communication Pipeline of Rollout-Side Quantization Information.
Instead of sending the full FP4 codeword for each rollout token, we send only its mantissa bits and the shared scale, which is enough to recover the neighboring training‑side codeword.
Transmit only the mantissa $m=2$ and the shared exponent $e=0$ from rollout to the training engine.
Training reconstructs the rollout value as $v_{\text{roll}} = (1 + m/4) \times 2^{e} = 1.5$.
Compare $v_{\text{roll}}$ with the training activation; because the training activation is slightly larger, the nearest higher codeword $(m=3, e=0)$ is selected.
Only one additional bit (the next mantissa) is needed to indicate the upward rounding direction.
Sending just the mantissa and exponent lets the training side recover the exact rollout value or its immediate neighbor, eliminating the need to transfer the full 4.5‑bit codeword.
**Figure 6.** Difference in FP4 Codebook Entries Between Training and Rollout Activations.
By communicating only mantissa and scale from the deeper half of the model, Mantissa-Only Communication slashes storage and bandwidth demands by orders of magnitude, yet still provides the precise signal needed for TRACE to align training and rollout quantization.
Experimental Setup and Baselines
We detail models, datasets, baselines, and implementation specifics for evaluating TRACE.
We evaluate TRACE on four large‑scale MoE language models across three RL task families, comparing against a wide range of NVFP4 and MXFP4 baselines.
NVFP4 is a 4‑bit floating‑point format; the baselines explore how naïve rollout, QAT, and post‑training tricks perform under this low‑precision regime.
Implementation uses the same VeRL, Megatron, and SGLang stack across all methods, with GRPO for policy optimization and R3 for replaying rollout‑side routing decisions.
We clip token‑level importance ratios at 5 and bound policy staleness with `max_version_diff` = 3; rollout‑guided quantization aligns rollout activations to training activations only for routed experts and KV caches.
Sampling budgets are 256×16, 128×16, 64×16, and 64×8 trajectories for the four models respectively; response length caps are 32 K tokens (reasoning) and context length caps 256 K tokens (coding and long‑horizon).
Performance Evaluation
TRACE recovers FP4 rollout performance, matching high‑precision BF16 across benchmarks.
The paper’s core issue is that quantized rollout (FP4) diverges from the training policy, causing instability. TRACE aligns the two, letting low‑precision inference keep the same quality as BF16.
TRACE consistently outperforms QAT, QaRL, and QUADS on all reasoning benchmarks.
Table 1 shows TRACE’s average score of 75.3 versus 68.8 for the best baseline.
**Figure 8.** Training dynamics of Qwen3.5-35B-A3B on reasoning RL tasks, including reward, response length, and test scores on HMMT, AIME24, AIME25, and LiveCodeBench.
TRACE consistently outperforms baselines across reasoning benchmarks.
Efficiency and Throughput
We quantify TRACE’s decoding throughput and RL training overhead.
We assess TRACE’s efficiency from two angles: decoding throughput across output lengths and end‑to‑end RL training step time.
TRACE adds only 7.4% overhead to end‑to‑end RL training step time.
Figure 9(b) shows step time rising from 664 ms to 713 ms.
**Figure 9.** Efficiency analysis of TRACE on Qwen3.5-35B-A3B. Decoding throughput is measured using only the rollout engine (SGLang) on 4 GB200 GPUs. End-to-end RL training step time is measured on the reasoning RL tasks using 72 GB200 GPUs.
The figure is a pie chart illustrating the distribution of different components: Forward (86%), Weight TRACE (4%), KV TRACE (2%), and KV Dequant (8%).
Table 3 provides an ablation of isolated NVFP4 weight and KV components, confirming that TRACE’s full configuration retains efficiency while delivering the best reasoning performance.
Ablation Study
We quantify how each TRACE component affects performance under FP4 quantization.
We first isolate the contribution of FP4 weight and KV‑cache quantization, then probe TRACE’s robustness across microscaling formats and varying mantissa‑bit coverage.
**Table 5.**
TRACE improves the average score by 6.1 points over the QAT baseline under the W4A8 + MXFP4 KV configuration.
Table 4 reports an average of 75.1 for TRACE versus 69.0 for the QAT baseline.
TRACE improves the average score by 6.3 points over the QAT baseline under the W4A4 + MXFP4 KV configuration.
Table 4 reports an average of 73.5 for TRACE versus 67.2 for the QAT baseline.
Next we examine how much rollout‑side mantissa information TRACE needs to retain for training.
**Table 5.** Modular sensitivity study of TRACE with FP4 MoE and FP4 KV on Qwen3.5-35B-A3B under reasoning RL tasks. TRACE (R-1bit-L20) is the default config. The best performance is marked in bold. The performance gain compared to the QUADS baseline is marked in green inside brackets.
TRACE with a single mantissa bit retained across all 40 layers attains an average score of 75.8.
Reported in the modular sensitivity study (Table 5).
Limiting mantissa‑bit retention to the last 10 layers reduces the average score to 74.1.
Table 5 entry for the R‑1bit‑L10 configuration.
Limiting mantissa‑bit retention to the last 5 layers reduces the average score to 73.4.
Table 5 entry for the R‑1bit‑L5 configuration.
Comparison with Score Centering
TRACE beats Score Centering by 1.2 pp on average under NVFP4 rollout.
TRACE outperforms Score Centering (SC) by 1.2 percentage points on average across the four benchmarks.
Table 6 shows TRACE achieving 75.3 average versus 74.1 for SC at step 380.
Score Centering reduces the train‑rollout log‑probability gap but still leaves a discrepancy of about 2, while TRACE keeps the gap near the BF16 baseline throughout training. Consequently TRACE maintains policy alignment more consistently.
**Figure 11.** Training dynamics of Qwen3.5-35B-A3B on reasoning RL tasks, including the minimum/maximum log-probability difference, reward, entropy loss, response length, and test scores on HMMT, AIME24, AIME25, and LiveCodeBench.
**Table 6.** Comparison of TRACE against baselines.
Post-Hoc Quantization Comparison
TRACE reaches 75.3 average score, beating all post‑hoc PTQ baselines.
TRACE outperforms all post‑hoc PTQ baselines on average FP4 policy score.
Table 7 shows TRACE achieves 75.3 versus 71.4 for the best PTQ method.
In contrast, the post‑hoc PTQ pipelines achieve average scores in the low‑70s, while TRACE maintains a substantially higher average by training with FP4 rollout throughout.
**Table 7.** Comparison of different PTQ algorithms and TRACE on various benchmarks.
Implementation Details
Additional implementation details, training diagnostics, and throughput results.
Algorithm 1 details the TRACE rollout‑guided quantization pipeline for MoE models. It operates on packed low‑bit rollout codes and per‑site scales, reconstructing a compact reference that steers the local FP4 rounding decision. The mantissa‑only regime is activated by setting $b = 1$ in the later layers.
TRACE: Rollout‑Guided FP4 QAT (simplified pseudocode)
On Qwen3.8‑Flash‑Next, vanilla QAT exhibits larger log‑probability gaps to the BF16 rollout reference, with the maximum difference widening over a 400‑step window. TRACE stays tightly aligned to the BF16 baseline, showing no comparable drift. Entropy declines sharply for QAT after step 100, while TRACE’s entropy drops more gradually and its responses remain longer, leading to higher Terminal‑Bench 2.1 scores throughout the window.
For Qwen3.8‑2.4T‑A95B, QAT’s maximum log‑probability difference hovers around 2.9, whereas TRACE and the BF16 reference stay near 0.95. QAT’s mean reward falls from 0.75 to 0.59 and its training entropy drops from 0.40 to 0.26, accompanied by a shortening of responses from 125 K to 93 K tokens. TRACE, in contrast, nudges reward upward to 0.78, keeps entropy around 0.38–0.40, and preserves response lengths in the 106 K–134 K range.
Throughput experiments compare four precision configurations on Qwen3.8‑Flash‑Next. Joint FP4 quantization outperforms FP8/BF16 by 1.09× at 8 K output, 1.41× at 16 K, 1.77× at 32 K, and 2.43× at 64 K; a single‑run 128 K test shows a 3.83× speedup (2013 tokens/s vs 526 tokens/s). The higher‑precision FP4/FP4 setting also raises the peak decoding batch size from 7 requests (FP8/BF16) to 64 requests, demonstrating the concurrency benefit of reduced KV memory.