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模型的數值一致性需要謹慎維護。開發者應認識到“數學等價”不等於“數值一致”，並通過合適的度量工具和靈活的控制策略來平衡性能與精度。