MoE模型中的训练与推理数值一致性：数值漂移的根源

在混合专家（MoE）模型的训练与推理过程中，保持数值一致性至关重要。然而，由于浮点数加法不满足结合律，看似数学等价的优化操作可能会产生不同的数值结果，这种现象被称为数值漂移。本文基于Kimi K2.5和Qwen3.5-MoE的实际经验，详细分析了导致数值漂移的几种常见陷阱。

首先，全归约（all-reduce）拓扑差异是重要原因之一。在训练中，NCCL通常使用环状拓扑进行归约，不同数据块经历不同的累加顺序；而推理引擎如FlashInfer的自定义Lamport内核则采用统一的累加顺序。尽管数学上等价，但浮点运算的非结合性导致结果出现微小差异。这种差异在单次操作中很小，但经过多层变换（如61层Transformer）后会显著累积。

其次，通信与计算的融合也会引起数值漂移。将全归约与RMSNorm融合为一个内核可以提升性能，但融合内核的线程布局与独立RMSNorm内核不同，导致平方和归约的加法顺序发生变化，从而影响最终的归一化缩放因子。这一差异同样源于浮点数非结合性，且难以察觉。

对于MoE模型，多操作融合带来的问题更为严重。推理引擎可能将MoE最终汇聚、全归约和下一层的RMSNorm融合为一个内核，每个操作都采用整体内核决定的线程布局，而非各自最优布局。由于MoE路由对隐藏状态的微小变化高度敏感，数值漂移可能改变专家选择，产生级联效应。

为了量化这些影响，本文使用k3指标（KL散度的一种变体）来衡量数值偏离。实验表明，即使单个融合操作带来的误差很小，但在61层网络中的累积效应不可忽视。作者建议提供细粒度的控制选项，让用户根据需求选择是否启用特定融合优化。

总之，MoE模型的数值一致性需要谨慎维护。开发者应认识到“数学等价”不等于“数值一致”，并通过合适的度量工具和灵活的控制策略来平衡性能与精度。