# TRACE: Rollout-Guided Quantization-Aware Training for FP4 Reinforcement Learning of MoE Language Models

> TRACE achieves BF16-level reinforcement learning performance on MoE language models by directly aligning FP4 quantization between training and rollout paths, with up to 5.4x rollout speedup.

- **Source:** [arXiv](https://arxiv.org/abs/2610.07767)
- **Published:** 2026-10-24
- **Permalink:** https://picx.dev/p/NDSVuL
- **Whiteboard:** https://picx.dev/p/NDSVuL/image

## Summary

## 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)

| Rollout Config | Algo. | LiveCodeBench | AIME24 | AIME25 | HMMT25 | Average |
|---|---|---|---|---|---|---|
| BF16 W/A + BF16 KV | Vanilla | 67.1 | 83.8 | 81.3 | 67.5 | 74.9 |
| NVFP4 W/A + NVFP4 KV | QAT | 55.1 | 70.7 | 63.4 | 49.2 | 59.6 |
| NVFP4 W/A + NVFP4 KV | QaRL | 54.1 | 73.3 | 57.9 | 49.2 | 58.6 |
| NVFP4 W/A + NVFP4 KV | QUADS | 60.4 | 80.5 | 75.4 | 59.0 | 68.8 |
| **NVFP4 W/A + NVFP4 KV** | **TRACE** | **66.4 (↑6.0)** | **86.3 (↑5.8)** | **78.5 (↑3.1)** | **70.0 (↑11.0)** | **75.3 (↑6.5)** |

### Larger-Scale Models (Table 2)

| Model | Task | BF16 | QAT | QUADS | TRACE |
|---|---|---|---|---|---|
| Qwen3.5-122B-A10B | DeepSWE | 33.4 | 28.8 | 29.1 | **33.0** (↑3.9) |
| Qwen3.8-Flash-Next | Terminal-Bench | 68.8 | 60.4 | 66.4 | **70.6** (↑4.2) |
| Qwen3.8-2.4T-A95B | GDPval | 90.3 | 88.7 | 85.9 | **90.2** (↑1.5) |

### 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)

| Config | Average Score |
|---|---|
| TRACE (R-4bit-L40) | 75.9 |
| TRACE (R-1bit-L40) | 75.8 |
| **TRACE (R-1bit-L20) — default** | **75.3** |
| TRACE (R-1bit-L10) | 74.1 |
| TRACE (R-1bit-L5) | 73.4 |

### Comparison with Score Centering (Table 6)

| Method | LiveCodeBench | AIME24 | AIME25 | HMMT25 | Average |
|---|---|---|---|---|---|
| BF16 | 67.1 | 83.8 | 81.3 | 67.5 | 74.9 |
| TIS+SC | 66.7 | 81.7 | 81.3 | 66.7 | 74.1 |
| **TRACE** | 66.4 | **86.3** (↑4.6) | 78.5 | **70.0** (↑3.3) | **75.3** (↑1.2) |

### Comparison with Post-Hoc FP4 Quantization (Table 7)

| Method | Average |
|---|---|
| BF16 | 74.9 |
| Vanilla NVFP4 (PTQ) | 70.4 |
| 4over6 (PTQ) | 71.0 |
| H-Scale (PTQ) | 71.4 |
| **TRACE (FP4 rollout throughout)** | **75.3** (↑3.9) |

### 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

---

_Markdown view of https://picx.dev/p/NDSVuL, served by PicX — AI-generated visual whiteboard summaries of research papers._