Skip to content

Back to the catalog

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

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

TRACE (Train-Rollout Quantization Alignment via Compact GuidancE) is an FP4 quantization framework for reinforcement learning post-training of Mixture-of-Experts language models. It uses rollout-guided quantization-aware training to align training-side FP4 rounding with rollout-side quantization outcomes, reducing the train-rollout discrepancy that limits prior low-precision RL methods, and adds a caching scheme that selectively retains mantissa/scale information from deeper layers to cut communication overhead. Evaluated on four large MoE models across reasoning, coding and long-horizon RL tasks, it matches BF16-rollout RL performance while achieving up to 5.4x rollout speedup with joint FP4 weight/activation and FP4 KV-cache quantization.

TL;DR A new FP4 quantization-aware training method closes the train-rollout gap for RL post-training of MoE LLMs, hitting BF16-level performance with up to 5.4x rollout speedup.
Open original ↗
Document

What it claims

TRACE (Alibaba Token Hub / Ohio State) targets a specific failure in FP4 RL rollout for MoE models. Training-side and rollout-side activations that differ only slightly can still round to different FP4 codewords, and that difference is then amplified. Prior FP4 RL methods (QUADS, Rollout-ResQ) minimise each path's error against its own BF16 value, which does not minimise the gap between the two quantized paths. TRACE instead makes training-side rounding follow what rollout actually produced.

Tested models:

  • Qwen3.5-35B-A3B
  • Qwen3.5-122B-A10B
  • Qwen3.8-Flash-Next
  • Qwen3.8-2.4T-A95B

Baselines:

  • RL-time: vanilla QAT, QaRL, QUADS, Score Centering.
  • Post-hoc FP4 quantization of a BF16-trained policy: vanilla NVFP4, 4over6, H-Scale.

How it works

Why rounding amplifies discrepancy. In the paper's example, train and rollout activations of 60.24 and 59.76 differ by 0.48. Under the same global and block scale they normalise to 2.51 and 2.49. These fall on opposite sides of a rounding boundary, round to 3 and 2, and the dequantized gap becomes 24.

QUADS counterexample. With train = 2.40 and rollout = 2.49, vanilla FP4 maps both to 2, so the discrepancy is 0. QUADS keeps 2.40 on the training side and reconstructs rollout as 2.50, so the discrepancy rises to 0.10.

Rollout-guided rounding. During rollout, TRACE records the FP4 outcomes for routed-expert activations and for the KV cache. In QAT (FP4 fake-quant forward, BF16 backward), each captured training value has two neighbouring E2M1 codewords under the rollout-side scale. TRACE picks whichever of the two is closest to the rollout codeword, instead of round-to-nearest. This guarantees the local discrepancy is never worse than RTN at each captured site.

Cost of storing the guidance. Full records are very large. For Qwen3.5-35B-A3B with 256K-token responses, the paper estimates about 50 KB per token (45 KB activation, 5.6 KB KV). At 4,096 trajectories that is about 51 TB per RL step. The authors' arithmetic:

  • GPU→CPU at 300 GB/s: about 3 minutes.
  • Storage I/O at 5 GB/s: about 3 hours, which makes storage the bottleneck.

Compression. TRACE keeps only mantissa and scale bits, and only for selected deeper layers. The two record types reach storage by different routes:

flowchart LR
  A[FP4 activation per forward] --> B[Reused GPU buffer]
  B -->|async offload| C[CPU cache]
  K[Persistent FP4 KV cache] -->|on request end| C
  C --> D[Disk]
  D -->|reload per microbatch| E[QAT rounding guidance]

Evidence

The fetched text shows dynamics, not results tables:

  • Figure 1 (Qwen3.8-Flash-Next, NVFP4 W/A + NVFP4 KV):
  • With conventional QAT, the train–rollout log-prob difference grows and the Terminal-Bench 2.1 score degrades.
  • TRACE stays aligned. It starts below BF16 rollout but closes the gap, first matching BF16 at step 240. The authors say the gap can eventually be "closed or even reversed."
  • Figure 4b (Qwen3.5-35B-A3B): vanilla NVFP4 shows large relative excess train–rollout MSE across layers, and TRACE substantially reduces it.
  • Headline claim: up to 5.4× rollout speedup over BF16. The authors also claim better final FP4 performance than post-hoc FP4 quantization of BF16-trained policies. The fetched text gives no per-benchmark numbers for either.

Limits

  • The guarantee is local only. The authors state it does not imply reduced full-network or policy-level divergence after nonlinearities.
  • Guidance is approximate in practice. The guarantee assumes the exact rollout codeword, but the deployed scheme uses partial (mantissa, deeper-layer) information.
  • The fetched page is truncated. It cuts off mid-Section 3.2: the layer-selection rule, the "over 99% of mismatched values differ by…" result, the results tables, the ablations and the overhead measurements are all missing.
  • Scope of quantization. It covers routed experts and the KV cache; attention appears to stay BF16 per Figure 2.
  • Infrastructure required. The method needs storage and I/O between the rollout and training engines.
  • Release status. No code or release is mentioned in the fetched text.

Technologies named here

  • TRACE
  • Mixture of Experts

Techniques named here

  • Rollout-guided quantization-aware training
  • Quantization-information caching scheme

Related in training