阿里:QUADS稳定MoE强化学习
📖标题:QUADS: Stabilizing NVFP4 Reinforcement Learning for MoE via QUantization-error Alignment across Dual Sides
🌐来源:arXiv, 2607.15810v1
🛎️文章简介
🔸研究问题:如何在使用NVFP4(W4A4)低精度加速推理时,解决混合专家模型(MoE)强化学习中因训练与推理精度不匹配导致的训练崩溃问题?
🔸主要贡献:论文提出了QUADS框架,通过双侧量化误差对齐策略,在保持NVFP4高吞吐量的同时,实现了与BF16全精度相当的强化学习训练稳定性与准确率。
📝重点思路
🔸问题诊断与误差分析:研究发现直接应用NVFP4进行MoE RL rollout会在约150步后崩溃。通过消融实验证明,激活值(Activation)的量化误差而非权重误差是导致对数概率差距扩大和训练不稳定的主导因素,因为激活值在线重计算且受粗糙的E2M1网格放大影响。
🔸训练侧非对称量化感知训练:提出非对称QAT方案,仅在训练端对权重进行FP4伪量化以对齐推理端的权重量化路径,同时保持激活值为BF16。这避免了在训练端对激活值进行粗粒度伪量化从而加剧引擎漂移的问题。
🔸推理侧残差激活补偿:针对推理端剩余的激活量化误差,设计残差补偿机制。通过在线选择高残差通道,对其进行二次FP4量化校正并加回结果中。该过程通过融合内核实现,在保留原生W4A4 GEMM高效性的同时显著缩小对数概率差距。
🔎分析总结
🔸训练稳定性显著提升: naive NVFP4 RL在多个基准测试中迅速崩溃,平均pass@1仅为51.37%;而QUADS成功避免了崩溃,训练曲线平稳,平均pass@1达到72.86%,与BF16基线(73.15%)和FP8基线(72.88%)相当。
🔸有效缩小精度差距:实验显示,QUADS将训练与推理间的最大对数概率差距从naive方法的约1.3降低至0.86,接近BF16引擎间固有的漂移下限(0.74),证明了双侧对齐的有效性。
🔸吞吐量优势得以保留:尽管引入了残差补偿,QUADS的rollout吞吐量仍比FP8高出约16%,仅比纯NVFP4慢5-9%,证明了该方法在效率与精度之间的良好平衡。
💡个人观点
论文定位了FP4强化学习不稳定的根源是激活量化误差,采用了“训练侧重权重、推理侧重激活”的非对称双侧对齐策略。