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モデルの数値的一貫性には慎重な管理が必要です。開発者は「数学的に等価」が「数値的に一致」を意味しないことを認識し、適切な測定ツールと柔軟な制御戦略を用いて性能と精度のバランスを取るべきです。