技術概要：MXFP4 による大規模言語モデルの事前学習

問題提起

大規模言語モデル（LLM）の急速なスケールアップにより、ハードウェア効率の高いトレーニングアルゴリズムへの需要が高まっている。FP8 は実用的な低精度代替手段として浮上しているが、4 ビット形式（MXFP4 など）を用いたトレーニングは依然として重大な課題である。主な難点は、アクティベーションと勾配のアウトライヤーが量子化誤差を増幅し、最適化を不安定化させる可能性があることだ。初期の研究では、粗く量子化されたアクティベーションのアウトライヤーが量子化範囲を支配し、浮動小数点の品質を著しく低下させることが示唆されている。混合精度や段階的スケジュールを通じてノイズを軽減する努力がなされているにもかかわらず、根本的な問いが残っている：なぜ、フォワードアクティベーションとアクティベーション勾配が安定しているように見える場合でも、フルパイプラインの FP4 トレーニングはしばしば発散するのか？ さらに、トレーニングパイプラインのどの特定のコンポーネント（フォワード伝播、アクティベーション勾配、または重み勾配）が収束の劣化に最も責任があり、どのような安定化メカニズムが有効であるかは不明である。

手法

著者は、MLPerf LLM 事前学習設定に従い、C4 データセットを用いたLlama 3.1–8Bの事前学習について制御された研究を実施した。本研究は、ソフトウェアエミュレーションに依存した先行研究とは異なり、AMD Instinct MI355X GPU上でネイティブな FP4 テンソルコアサポートを活用して実行された。

核心的な実験デザインは、MXFP4 有効化の段階的な「ラダー」であった：

1. ベースライン： フル FP8 トレーニング。
2. 段階的量子化： 3 つのコンポーネントに対して MXFP4 を段階的に有効化し、他を FP8 に固定する：
   • フォワード伝播（Fprop）のみ。
   • Fprop + アクティベーション勾配（Dgrad）。
   • Fprop + Dgrad + 重み勾配（Wgrad）。
3. 安定化評価： フル MXFP4 パイプライン（Fprop + Dgrad + Wgrad）の下で、安定性を回復させるための各種介入をテストした：
   • 確率的丸め： randomness を導入するため。
   • ランダム化ハダマード回転： アウトライヤーを非相関化するため。
   • 決定論的ハダマード回転： 具体的に $16 \times 16$（H16）および $32 \times 32$（H32）行列を使用。

指標： 収束は、検証パープレキシティが3.3に達することとして定義された。効率性の主要指標はトークンオーバーヘッド（FP8 ベースラインと比較して目標到達に必要なトークンの相対的な増加）であった。二次指標には、トレーニングステップのスループットとエンドツーエンドの高速化が含まれる。

主要な貢献と知見

1. 不安定化源の特定

本研究は、重み勾配（Wgrad）の量子化を収束劣化の主要な駆動力として特定した。

• Fprop のみに対して MXFP4 を有効化した場合、トークンオーバーヘッドは**8–9%**であった。
• Fprop + Dgradを有効化した場合、同様に**10–11%**のオーバーヘッドが生じた。
• **フルパイプライン（Fprop + Dgrad + Wgrad）を安定化なしで有効化した場合、トークンオーバーヘッドは26–27%**に跳ね上がるか、完全に発散（「収束せず」）した。
• 結論： 不安定性は Fprop または Dgrad の量子化に固有のものではなく、Wgrad の量子化によって支配されている。

2. 安定化戦略の評価

著者は、確率性や構造化された回転を追加することでフルパイプラインを安定化できるかをテストした：

• 確率的丸め： フルパイプラインを安定化できず、実行は発散した。
• ランダム化ハダマード回転： フルパイプラインに適用しても同様に収束しなかった。
• 決定論的ハダマード回転： 安定した最適化を成功裡に回復させた。決定論的ハダマード変換（具体的には H16）をフル MXFP4 パイプラインに適用することで、トークンオーバーヘッドを**8–9%**まで戻し、FP8 ベースラインに密接に追従させた。
• 洞察： 確率的介入は、アウトライヤーの影響を軽減するのではなく、むしろ実効的な量子化誤差を増幅するようだ。不安定性は、敏感な勾配経路に沿った構造化されたマイクロスケーリング誤差によって駆動されており、決定論的回転によって修正可能である。

3. 効率性とスループット

決定論的ハダマード（H16）による安定性問題の解決により、著者は FP4 の実用的な効率性向上を実証した：

• スループット： MI355X GPU 上での FP8 ベースラインと比較して、MXFP4 + H16 構成はトレーニングステップスループットが20% 高いことを達成した。
• エンドツーエンドの高速化： トークンオーバーヘッドが最小化（8–9%）されたため、スループットの優位性は9–10% のエンドツーエンド高速化に転換された。
• 主要な洞察： FP4 の高速化は「安定性ゲート付き」である。Wgrad の不安定性を制御しなければ、収束に必要なトークン数の増加が、低精度による生計算利点を相殺してしまう。

意義と主張

本論文は、ソフトウェアエミュレーションなしに MI355X クラスのハードウェア上でフル MXFP4 事前学習を実証した最初の事例であると主張している。その意義は、FP4 トレーニングダイナミクスに関する理解をシフトさせる点にある：

• 失敗のメカニズム： FP4 事前学習における不安定性は、不十分な確率性によるものではなく、特に Wgrad における敏感な勾配経路の構造化されたマイクロスケーリング誤差によるものである。
• 実用的なレシピ： 決定論的ハダマード回転は、フルパイプライン MXFP4 トレーニングにとって必要かつ効果的な安定化剤であり、ハードウェア効率の向上を実現する。
• 文脈的限界： 著者は控えめに、この安定化レシピは設定依存性があると指摘している。C4 データセット上の Llama 3.1–8B に対しては有効であるが、異なるモデル、データセット、または微調整手法では挙動が異なる可能性があり、さらなる証拠なしに普遍的であると仮定すべきではない。

要約すると、この研究は、Wgrad 量子化誤差を確率的ノイズ注入ではなく決定論的な構造的介入によって制御する限り、大規模事前学習において FP8 を MXFP4 に置き換えることが可能であることを確立している。