TRIAGE: Direction-Aware Mismatch Stabilization of Native NVFP4 Reinforcement Learning

2026-10-04

AI Summary

This study addresses the training instability arising from learner-sampler misalignment in low-precision policy optimization. We propose a direction-aware stabilization method that employs a piecewise diagnostic mechanism to reveal the interaction between such mismatches and gradient directions, alongside selective rebalancing and bounded repair algorithms to correct severe deviations. This approach enables native NVFP4 (W4A4) quantized training, achieving stable optimization while preserving forward inference efficiency. Experimental results demonstrate that the proposed method attains full-precision performance on mathematical reasoning benchmarks, yielding throughput improvements of up to 2.3× over the BF16 baseline.

Abstract

Low-precision execution can substantially accelerate reinforcement learning (RL) for large language models, but discrepancies between learner and sampler execution can destabilize policy optimization. In this paper, we characterize the interaction between mismatch and the policy-gradient direction, distinguishing locally amplifying from contracting update contributions that mismatch magnitude alone cannot identify. In native NVFP4 runs, we observe an early imbalance between the two amplifying regions, favoring negative-advantage, negative-gap updates. Their tail tokens become concentrated in a small fraction of response segments before mismatch spreads globally. Motivated by these findings, we introduce TRIAGE, a direction-aware stabilization method that uses segment-level diagnosis to selectively rebalance policy-gradient updates and applies bounded repair to residual severe mismatch. TRIAGE modifies the optimization objective while retaining native NVFP4 weight-and activation 4-bit (W4A4) forward execution on both the sampler and learner. Experiments on Qwen3-4B and Qwen3-30B-A3B show stable optimization throughout the evaluated training horizon and achieve full precision level performance across five mathematical reasoning benchmarks, while native NVFP4 with TRIAGE provides up to 2.3x higher rollout throughput than BF16.

Problem

• reinforcement learning
• low-precision execution
• policy optimization instability
• NVFP4
• mismatch stabilization

Innovation

• Direction-Aware Stabilization
• Native NVFP4
• Policy-Gradient Mismatch
• Segment-Level Diagnosis
• Reinforcement Learning