The 4-Bitter Lesson: Balancing Stability and Performance in NVFP4 RL
3 days ago
- #reinforcement learning
- #quantization
- #low-precision training
- 核心挑战在于，在使用NVFP4等低精度量化进行大语言模型（LLM）的强化学习（RL）时，如何平衡吞吐量和稳定性，因为采样和训练不稳定性会相互叠加。
- 基准方案对MoE层（占参数的97%）使用NVFP4，并通过每token激活缩放来避免token依赖量化和未来到过去的泄漏等问题。
- 通过反量化反向传播改进梯度稳定性，使反向操作数与正向量化决策对齐，减少不匹配和梯度范数尖峰。
- 四比六（4/6）技术通过自适应选择权重和激活的最大值为±4或±6来降低量化误差，这对于RL至关重要，因为对预训练权重的影响不成比例。
- 选择性层精度保持敏感层（例如最终层和共享MoE专家）使用更高精度（BF16），以优化内存-性能权衡，而无需针对每个模型进行定制处理。
- 组合这些优化策略可得到与BF16相当的稳定梯度范数和奖励曲线，同时支持高效的在线NVFP4服务，无需校准或单独流水线。
- 未来改进可能包括将MXFP8激活与NVFP4权重结合使用以减少开销，并协同设计优化器以应对低精度训练挑战。
- 这项工作强调避免单步陷阱，关注训练器和采样器之间的交互，以及长期、多智能体的动态性，而非孤立任务。
- 开源实现以及与RadixArk和NVIDIA的合作是关键，对TransformerEngine、FlashInfer、SGLang等项目都有贡献。