PaperlayerView exact revision on arXiv

Source-linked AI summary

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

arXiv:2610.07767v1cs.LGcs.CL

TL;DR

RL rollout incurs substantial computation and memory overhead, motivating low-precision rollout. TRACE uses rollout-side FP4 quantization outcomes to guide training-side rounding and reduce train–rollout discrepancy. Across multiple MoE models and RL tasks, it enables joint FP4 weight/activation and KV-cache rollout with performance comparable to BF16 rollout and up to 5.4× rollout speedup.

  • Problem

    Repeated generation of long trajectories makes RL rollout costly in computation and memory, motivating low-precision rollout.

    Source: §1 · p. 1

  • Method

    TRACE uses rollout-side FP4 quantization outcomes to guide training-side rounding decisions and caches selected deeper-layer mantissa and scale information to reduce guidance overhead.

    Source: §1 · p. 2, §3 · p. 3

  • Results

    Across multiple MoE models and RL tasks, TRACE preserves performance comparable to BF16 rollout while enabling joint FP4 weight/activation and KV-cache rollout, with up to 5.4× rollout speedup over BF16 rollout.

    Source: §1 · p. 2, §5 · p. 10

  • Takeaways & Limitations

    TRACE achieves stronger final FP4 task performance than pipelines that train with BF16 rollout and then apply post-hoc FP4 quantization.

    Source: §1 · p. 2

    The local guarantee does not ensure monotonic reduction of full-network discrepancy or policy-level divergence after downstream nonlinear transformations.

    Source: §3.1 · p. 4

Abstract

from arXiv · show

Reinforcement learning (RL) for post-training large language models (LLMs) incurs substantial computation and memory overhead during rollout generation, which motivates low-precision rollout for efficient RL training. However, existing FP4 RL methods suffer from a key limitation: they primarily optimize quantization accuracy on the training and rollout paths independently rather than directly reducing the discrepancy between the two quantized execution paths. In this work, we propose TRACE (Train-Rollout Quantization Alignment via Compact GuidancE), an FP4 quantization framework for RL training of Mixture-of-Experts (MoE) language models that addresses the limitation of existing FP4 RL methods. TRACE incorporates rollout-guided quantization-aware training that uses rollout-side quantization outcomes to guide training-side FP4 rounding decisions, directly reducing train-rollout discrepancy. Moreover, TRACE adopts an efficient quantization-information caching scheme that selectively retains mantissa and scale information from deeper layers to reduce the storage and communication overhead introduced by rollout guidance. We evaluate TRACE on four large-scale MoE language models across reasoning, coding, and long-horizon RL tasks. Our results demonstrate that TRACE enables joint FP4 weight/activation and FP4 KV-cache rollout with RL performance comparable to BF16 rollout, while achieving up to 5.4xrollout speedup and strong final FP4 performance compared with post-hoc FP4 quantization of BF16-trained policies.

1 Introduction

TRACE addresses the train–rollout discrepancy that per-path FP4 quantization accuracy does not resolve, using rollout outcomes to guide training-side rounding. Across four MoE models and varied RL tasks, it reports BF16-comparable performance with faster rollout and stronger final FP4 performance than post-hoc quantization.

1 Introduction

  • RL rollout generation incurs substantial computation and memory overhead, motivating low-precision rollout, while aggressive FP4 can introduce numerical discrepancies between training and rollout.The passage notes that MoE routing can further interact with numerical differences.Source: §1 · p. 1
  • Per-path quantization accuracy does not necessarily reduce discrepancy between the two quantized execution paths, leaving policy mismatch and degraded RL performance possible.The passage also notes that this mismatch can lead to training collapse.Source: §1 · p. 1
  • TRACE uses rollout-side quantization outcomes to guide training-side FP4 rounding and selectively caches mantissa and scale information from deeper layers.The selective cache reduces storage and communication overhead from rollout guidance.Source: §1 · p. 2
  • TRACE controls train–rollout discrepancy as training proceeds, allowing its initial performance gap relative to BF16 rollout to diminish and eventually close or reverse.Conventional FP4 QAT instead shows increasing policy mismatch and sharp test-score degradation on Terminal-Bench.Source: §1 · p. 2
  • The evaluation compares TRACE with FP4 rollout and post-hoc quantization baselines across four MoE models and reasoning, coding, and long-horizon RL tasks.The listed models are Qwen3.5-35B-A3B, Qwen3.5-122B-A10B, Qwen3.8-Flash-Next, and Qwen3.8-2.4T-A95B.Source: §1 · p. 2
  • Up to 5.4× rollout speedup over BF16 rollout accompanies BF16-comparable performance and stronger final FP4 task performance than post-hoc FP4 quantization of BF16-trained policies.TRACE maintains training dynamics and final task performance comparable to BF16 rollout.Source: §1 · p. 2

2 Related Work

Prior work addresses MoE routing sensitivity, rollout cost, and train–rollout inconsistency through routing replay, optimization changes, and low-precision methods. TRACE is distinguished by directly using rollout-side FP4 outcomes to guide training-side rounding.

2 Related Work

  • Small numerical differences can change router scores and activate different experts in MoE RL; prior responses include routing replay in R3 and PR2 and sequence-level importance ratios and clipping in GSPO.The passage also identifies memory, communication, and rollout costs from MoE model capacity as motivations for efficient training and low-precision rollout.Source: §2 · p. 2
  • Prior low-precision rollout methods include routed-expert quantization, 8-bit approaches such as QuRL and Jet-RL, and 4-bit approaches including QeRL, QUADS, Rollout-ResQ, and QaRL.The passage describes these approaches as using adaptive clipping and update-aware quantization, unified FP8 computation, asymmetric QAT with activation compensation, sparse residuals, or matched low-precision kernels, respectively.Source: §2 · pp. 2–3
  • Unlike optimization-level correction or matched execution alone, TRACE directly uses rollout-side FP4 outcomes to guide training-side rounding and target rounding-outcome mismatch.The passage notes that small activation differences can cross FP4 rounding boundaries and produce different codewords.Source: §2 · p. 3

3 Train-Rollout Quantization Alignment via Compact GuidancE (TRACE)

TRACE aligns training and rollout FP4 computations by using rollout quantization outcomes to guide training-side rounding, then reduces guidance overhead by caching selected quantization information. Its design addresses the gap between independently minimizing each path’s quantization error and reducing discrepancy between the paths.

3 Train-Rollout Quantization Alignment via Compact GuidancE (TRACE)

  • TRACE records rollout-side FP4 activation and KV quantization outcomes to guide training-side rounding, while caching selected deeper-layer mantissas and scales to reduce overhead.This combines rollout-guided quantization-aware training with a compact information cache.Source: §3 · p. 3

3.1 Rollout-Guided Quantization-Aware Training

  • QUADS examples show that lower per-path quantization error need not reduce train–rollout discrepancy: the examples retain discrepancies of 0.10 and 0.99.In one example, vanilla FP4 maps both activations to 2, while QUADS reconstructs the rollout activation as 2.50 and leaves the training activation at 2.40.Source: §3.1 · pp. 3–4
  • A 0.48 BF16 activation gap can become 24 after FP4 rounding across a codeword boundary; TRACE instead selects the neighboring codeword closest to the rollout result.The example’s normalized values, 2.51 and 2.49, round to 3 and 2, respectively.Source: §3.1 · p. 4
  • With full rollout-codeword information, TRACE does not increase local discrepancy versus RTN at captured sites, but this does not guarantee lower full-network discrepancy or policy divergence.The local guarantee applies under the same candidate set and does not imply monotonic improvement after downstream nonlinear transformations.Source: §3.1 · p. 4

3.2 Mantissa-Only Train-Rollout Communication

  • For Qwen3.5-35B-A3B, 4,096 trajectories can generate up to approximately 51 TB of rollout quantization information per RL step.The passage estimates that writing or reading this volume at a representative 5 GB/s would take roughly three hours serially.Source: §3.2 · pp. 4–5
  • Because over 99% of mismatched values differ by one adjacent FP4 codebook entry, TRACE communicates mantissa and scale information only from the latter half of model layers.The compact information is used to reconstruct a rollout-side reference and guide training-side quantization, reducing storage and communication overhead.Source: §3.2 · p. 5

4 Experiment Analysis

Experiments compare TRACE with FP4 rollout and post-hoc quantization baselines across four MoE models and reasoning, coding, and long-horizon RL tasks. They assess performance, efficiency, ablations, and comparisons with Score Centering and post-hoc FP4 quantization.

4.1 Overall Comparison

  • TRACE raises the reasoning average score from 68.8 to 75.3 and matches BF16 rollout on Qwen3.5-35B-A3B.It outperforms QAT, QaRL, and QUADS across four benchmarks; on HMMT25, the score increases from 59.0 to 70.0.Source: Table 1 · p. 6, §4.1 · p. 6
  • TRACE matches BF16 rollout across three larger MoE models: 33.0 vs. 33.4 on DeepSWE, 70.6 vs. 68.8 on Terminal-Bench, and 90.2 vs. 90.3 on GDPval.On Terminal-Bench, TRACE improves the score from 66.4 to 70.6 compared with QAT.Source: Table 2 · p. 6, §4.1 · pp. 6–7

4.2 Efficiency Analysis

  • TRACE achieves up to 5.4× higher decoding throughput than BF16 rollout at 128K output length, while increasing RL step time from 664 to 713 (7.4% over vanilla joint FP4 rollout).Its rollout-time breakdown is 86% model forward pass, 4% weight-side guidance, 2% KV-side guidance, and 8% KV dequantization.Source: Figure 9 · p. 8, §4.2 · p. 8, Figure 10 · p. 8

4.3 Ablation Study

  • Retaining one mantissa bit across all 40 layers achieves an average score of 75.8; coverage limited to the latter 20 layers causes small degradation, while 10 and 5 layers score 74.1 and 73.4.TRACE preserves strong performance with compact mantissa information from a subset of deeper layers.Source: Table 5 · p. 9, §4.3 · p. 9
  • TRACE remains BF16-comparable in isolated FP4 settings and raises average scores from 69.0 to 75.1 (W4A8) and 67.2 to 73.5 (W4A4) under MXFP4.The joint FP4 configuration is more challenging than either isolated setting.Source: §4.3 · pp. 8–9, Table 3 · p. 8, Table 4 · p. 9

4.4 Comparison with Score Centering Algorithm

  • Against Score Centering, TRACE keeps train–rollout log-probability differences near BF16 and scores 75.3 versus 74.1 average at step 380.Score Centering prevents vanilla QAT's severe collapse but leaves a larger discrepancy throughout training.Source: Table 6 · p. 9, §4.4 · p. 9, Figure 11 · p. 10

4.5 Comparison with Post-Hoc FP4 Quantization

  • TRACE achieves a final average score of 75.3, compared with 70.4, 71.0, and 71.4 for vanilla NVFP4, 4over6, and H-Scale post-hoc quantization.The post-hoc pipelines first train with BF16 rollout; TRACE trains throughout with joint FP4 weight/activation and KV-cache rollout.Source: Table 7 · p. 10, §4.5 · p. 10

5 Conclusion

TRACE combines rollout-guided QAT with compact quantization-information caching to reduce train–rollout discrepancy and its overhead. Across multiple MoE models and RL tasks, it enables joint FP4 weight/activation and KV-cache rollout with BF16-comparable performance and improved rollout efficiency.

5 Conclusion

  • TRACE combines rollout-guided QAT and efficient quantization-information caching to reduce train–rollout discrepancy and its associated overhead.Experiments across multiple MoE models and RL tasks show joint FP4 weight/activation and KV-cache rollout with performance comparable to BF16 rollout and substantially improved rollout efficiency.Source: §5 · p. 10

A.1 Pseudocode of TRACE

TRACE records compact rollout-side quantization information during trajectory generation and uses it during QAT to guide FP4 rounding decisions. The algorithm then computes the RL loss and updates policy parameters.

A.1 Pseudocode of TRACE

  • Each RL iteration generates trajectories and rollout log-probabilities using the FP4 policy, while recording packed low-bit rollout codes and scales at captured sites.Advantages are computed from trajectory rewards before policy-update minibatches are processed.Source: §A.1 · p. 14
  • For each captured training activation, TRACE retrieves its rollout record, normalizes by the rollout scale, reconstructs a local rollout-side reference, and selects the nearer adjacent FP4 rounding candidate.The algorithm uses the selected scale in the remaining forward computation through the straight-through estimator.Source: §A.1 · p. 14
  • After QAT computes the policy loss from training log-probabilities, rollout log-probabilities, and advantages, the optimizer updates the policy parameters.The cached rollout record is indexed by trajectory, token, and quantization site.Source: §A.1 · p. 14

A.2 Training Dynamics on Qwen3.8-Flash-Next

On Qwen3.8-Flash-Next coding RL, TRACE stays closer to the BF16 reference in train–rollout log-probability diagnostics than vanilla QAT over the recorded 400-step window. It also maintains longer responses and improves reward on average, while Terminal-Bench test scores increase.

A.2 Training Dynamics on Qwen3.8-Flash-Next

  • Across the 400-step window, vanilla QAT has larger log-probability differences than TRACE and BF16, with its maximum rising and minimum becoming more negative toward the run's end.TRACE remains closer to BF16 in both diagnostics, without comparable widening.Source: §A.2 · p. 14
  • After approximately step 100, vanilla QAT's entropy decreases markedly and its responses shorten, while TRACE's entropy decreases more gradually and its responses remain long.TRACE improves reward on average over the recorded trajectory.Source: §A.2 · p. 14
  • TRACE's Terminal-Bench 2.1 test scores increase over the recorded training window.The six panels also show training diagnostics and test scores over 400 steps.Source: §A.2 · p. 14

A.3 Training Dynamics on Qwen3.8-2.4T-A95B

On Qwen3.8-2.4T-A95B long-horizon RL, TRACE stays closer to the BF16 rollout reference than vanilla QAT in maximum log-probability difference. Its reward, entropy, and response-length trajectories also avoid the declines reported for vanilla QAT.

Training–Rollout Consistency

  • Near 0.95 versus vanilla QAT’s near 2.9 maximum log-probability difference, TRACE stays close to the BF16 reference across 50 steps.The comparison covers joint NVFP4 weight and KV-cache quantization.Source: §A.3 · pp. 14–15, Figure 13 · p. 15

Reward, Entropy, and Response Length

  • Mean reward rises from about 0.75 to about 0.78 with TRACE, while vanilla QAT falls from about 0.75 to about 0.59.These trajectories support the reported training-stability trend for long-horizon RL.Source: §A.3 · p. 15
  • TRACE keeps entropy near the BF16 reference at about 0.38–0.40 and responses in the 106K–134K range; vanilla QAT falls to about 0.26 entropy and 93K tokens.Vanilla QAT starts at about 0.40 entropy and 125K response length.Source: §A.3 · p. 15

A.4 Throughput Analysis on Qwen3.8-Flash-Next

On Qwen3.8-Flash-Next, joint FP4 improves output throughput over FP8/BF16 across the reported output lengths. The measurements use a single-GB200 setup with fixed prompt and request settings.

Dependence on Output Length

  • 1.09×, 1.41×, 1.77×, and 2.43× throughput gains for joint FP4 over FP8/BF16 at 8K, 16K, 32K, and 64K outputs, respectively.Figure 14 compares four routed-expert/KV precision combinations.Source: Figure 14 · p. 15, §A.4 · p. 16

Setup and Metric

  • The comparison uses one GB200, a shared checkpoint, BF16 PLE/ngram embeddings offloaded to CPU pinned memory, and 64 synthetic 8K-token prompts.GPU memory fraction is 0.85 with automatic KV allocation; each request generates 8K–128K output tokens.Source: §A.4 · pp. 15–16, Figure 14 · p. 15

Dependence on Output Length

  • At 128K output, joint FP4 reaches 2013 tokens/s versus 526 tokens/s for FP8/BF16 (3.83×), with peak decoding batches of 64 versus 7 requests.The reported gain includes concurrency enabled by KV capacity; the 128K result is a single-run supplement.Source: §A.4 · p. 16

Related papers

Loading 2610.07767v1…