The 4-Bitter Lesson: Balancing Stability and Performance in NVFP4 RL
3 days ago
- #reinforcement learning
- #quantization
- #low-precision training
- The core challenge is balancing throughput and stability when using low-precision quantization like NVFP4 in reinforcement learning (RL) for large language models, as sampling and training instabilities compound.
- A baseline recipe uses NVFP4 for MoE layers (97% of parameters) with per-token activation scaling to avoid issues like token-dependent quantization and future-to-past leakage.
- Gradient stability is improved via dequantized backward passes, aligning backward operands with forward quantization decisions to reduce mismatches and gradient norm spikes.
- The four-over-six (4/6) technique reduces quantization error by adaptively choosing between maximal values of ±4 or ±6 for both weights and activations, crucial for RL due to disproportionate impacts on pre-trained weights.
- Selective layer precisions keep sensitive layers (e.g., final layers and shared MoE experts) in higher precision (BF16) to optimize memory-performance trade-offs without custom per-model handling.
- Combining optimizations yields stable gradient norms and reward curves comparable to BF16, while enabling efficient online NVFP4 serving without calibration or separate pipelines.
- Future improvements may include using MXFP8 activations with NVFP4 weights to reduce overhead and co-designing optimizers for low-precision training challenges.
- The work emphasizes avoiding the one-step trap by focusing on interactions between trainers and samplers, and long-horizon, multi-agent dynamics rather than standalone tasks.
- Open-source implementations and collaborations with RadixArk and NVIDIA are key, with contributions to TransformerEngine, FlashInfer, SGLang, and others.