Summary (Overview)

• TRACE (Train-Rollout Quantization Alignment via Compact GuidancE) is an FP4 quantization framework for reinforcement learning (RL) training of Mixture-of-Experts (MoE) language models that directly reduces the discrepancy between training and rollout execution paths.
• The core innovation is rollout-guided quantization-aware training, which uses rollout-side FP4 quantization outcomes to guide training-side rounding decisions, rather than independently optimizing quantization accuracy on each path.
• TRACE employs mantissa-only train-rollout communication, caching only mantissa and scale information from selected deeper layers, reducing storage/communication overhead by ~6× while preserving effectiveness.
• Evaluated on four MoE models (Qwen3.5-35B-A3B, Qwen3.5-122B-A10B, Qwen3.8-Flash-Next, Qwen3.8-2.4T-A95B) across reasoning, coding, and long-horizon RL tasks, TRACE achieves BF16-level performance with up to 5.4× rollout speedup.
• TRACE enables the policy to progressively adapt to FP4 execution, recovering initial quantization-induced degradation and matching or exceeding BF16 rollout performance after ~120 training steps.

Introduction and Theoretical Foundation

Background and Motivation

Reinforcement learning has become a critical post-training paradigm for LLMs, improving reasoning, coding, and complex capabilities. However, RL training incurs substantial computation and memory overhead due to repeated generation of long trajectories during rollout. Low-precision rollout (e.g., FP4 quantization with W4A4) offers a promising solution by quantizing the policy model during generation.

The Core Problem: Train-Rollout Discrepancy

Existing FP4 RL methods (e.g., QUADS, Rollout-ResQ, QaRL) suffer from a fundamental issue: they optimize quantization accuracy on training and rollout paths independently, minimizing error with respect to high-precision representations. However, minimizing per-path quantization error does not directly minimize the discrepancy between the two quantized execution paths.

The train-rollout quantized discrepancy is formalized as:

$$\mathcal{D}_{\mathrm{act}} = \left\| Q_{\mathrm{FP4}}^{\mathrm{train}}(X_{\mathrm{train}}) - Q_{\mathrm{FP4}}^{\mathrm{rollout}}(X_{\mathrm{rollout}}) \right\|_{F} \tag{1}$$

Key insight from Figure 3 examples: In one case, training activation = 2.40, rollout activation = 2.49. Vanilla FP4 maps both to 2 (zero discrepancy), but QUADS retains 2.40 on training side and reconstructs 2.50 on rollout side, increasing discrepancy to 0.10 despite reducing per-path error.

Discrepancy amplification: Small BF16 differences can be amplified by FP4 rounding. If training activation = 60.24 and rollout activation = 59.76 (difference 0.48), after normalization they may fall on opposite sides of a rounding boundary (2.51 → 3, 2.49 → 2), amplifying discrepancy from 0.48 to 24.

Methodology

1. Rollout-Guided Quantization-Aware Training

TRACE uses rollout-side quantization outcomes to guide training-side FP4 rounding decisions. For each captured training-side activation, let $\{q_-, q_+\}$ be its two neighboring normalized FP4 E2M1 codewords under the rollout-side scale, and $q_{\mathrm{rollout}}$ be the exact FP4 codeword from the rollout-side activation. Instead of round-to-nearest (RTN), TRACE selects:

$$q_{\text{TRACE}} = \arg\min_{q \in \{q_-, q_+\}} |q - q_{\text{rollout}}| \tag{2}$$

This guarantees:

$$\left| q_{\text{TRACE}} - q_{\text{rollout}} \right| \leq \left| q_{\text{RTN}} - q_{\text{rollout}} \right| \tag{3}$$

TRACE does not increase local quantized train-rollout discrepancy at each captured site relative to RTN.

2. Mantissa-Only Train-Rollout Communication

Motivation: Full rollout-side quantization information creates massive overhead. For Qwen3.5-35B-A3B with 40 MoE layers and max response length of 256K tokens, one generated token produces ~50 KB of guidance data. With 4,096 trajectories per RL step, this amounts to ~51 TB per step — far exceeding practical storage/communication capacity.

Key observation: For over 99% of mismatched quantized values across layers, training and rollout results differ by only one adjacent FP4 codebook entry. This means a small amount of mantissa code information suffices to identify the desired neighboring rollout-side codeword.

Design: TRACE communicates only mantissa and scale information from the latter half of model layers (deeper layers show predominantly more rounding corrections). This reduces cache size by ~6× (from ~45 KB to ~7.5 KB per token for Qwen3.5-35B-A3B).

Empirical Validation / Results

Overall Performance (Table 1: Qwen3.5-35B-A3B, Reasoning RL)

Larger-Scale Models (Table 2)

Efficiency Results

• Decoding throughput: Up to 5.4× higher than BF16 rollout at 128K output length
• End-to-end RL step time: Only 7.4% overhead (664 → 713 seconds) over vanilla joint FP4 rollout
• Rollout time breakdown: model forward = 86%, weight-side guidance = 4%, KV-side guidance = 2%, KV dequantization = 8%

Ablation: Modular Sensitivity (Table 5)

Comparison with Score Centering (Table 6)

Comparison with Post-Hoc FP4 Quantization (Table 7)

Generalization to MXFP4 Formats (Table 4)

• W4A8 + MXFP4 KV: Average 69.0 (QAT) → 75.1 (TRACE), ↑6.1
• W4A4 + MXFP4 KV: Average 67.2 (QAT) → 73.5 (TRACE), ↑6.3

Theoretical and Practical Implications

Theoretical Significance

1. Paradigm shift in QAT for RL: TRACE demonstrates that directly optimizing cross-path discrepancy (train vs. rollout) is more effective than independently minimizing per-path quantization error — a principle that extends beyond FP4 to other low-precision formats.

2. Quantization-induced adaptation: The observation that policies can adapt to low-precision execution environments when train-rollout discrepancy is controlled suggests that quantization-aware RL training can yield better final FP4 policies than post-hoc quantization of BF16-trained models.

3. Routing stability for MoE: By reducing numerical discrepancies that can alter router scores and expert activation, TRACE addresses a critical vulnerability specific to MoE architectures.

Practical Implications

1. Efficiency gains: Up to 5.4× rollout speedup with only 7.4% training overhead makes FP4 RL training practically feasible for large-scale MoE models.

2. Storage feasibility: The mantissa-only caching scheme reduces guidance data from ~51 TB/step to a manageable ~6× smaller footprint, making rollout guidance practical in production settings.

3. General applicability: TRACE works across model scales (35B to 2.4T), task types (reasoning, coding, long-horizon), and FP4 formats (NVFP4, MXFP4), suggesting broad industrial applicability.

Conclusion

TRACE addresses the fundamental limitation of existing FP4 RL methods by directly reducing train-rollout discrepancy through rollout-guided quantization-aware training. Key takeaways:

1. Direct discrepancy minimization outperforms independent per-path quantization accuracy optimization.
2. Compact guidance (mantissa-only, deeper-layer-focused caching) preserves effectiveness while keeping overhead minimal.
3. Policy adaptation to low-precision execution is achievable when discrepancy is controlled, yielding final FP4 policies that outperform BF16-trained + post-hoc quantized pipelines.

Future directions suggested by this work include:

• Extending rollout-guided alignment to other quantization formats (e.g., INT4, INT8)
• Investigating layer-adaptive guidance allocation strategies
• Exploring the interaction between rollout guidance and other RL stability techniques (e.g., expert routing replay, importance sampling) in more detail