🤖 machine learning

Full-Stack FP4: Stable LLM Pretraining with Quantized Projections, Optimizers, and Attention

本論文は、NVFP4量子化を線形投影だけでなく、安定したオプティマイザ状態、Root/Muon計算、および混合精度アテンションへと拡張するモジュール式フレームワークであるFull-Stack FP4を紹介し、単一のRTX 5090において、大幅なメモリ削減と高速化を実現しつつ、BF16に近い事前学習性能を達成しています。

原著者： Siyu Ding, Mingchuan Ma, Jiabo Tong, Xingrun Xing, Ziming Wang, Guoqi Li

公開日 2026-08-11

📖 1 分で読めます☕ さくっと読める

CC BY 4.0

原著者： Siyu Ding, Mingchuan Ma, Jiabo Tong, Xingrun Xing, Ziming Wang, Guoqi Li

原論文は CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) でライセンスされています。 ✨ これは以下の論文のAI生成解説です。著者が執筆または承認したものではありません。技術的な正確性については原論文を参照してください。免責事項の全文を読む

あなたは、巨大なデジタル脳に話し方、書き方、そして考え方を教えようとしているところだと想像してください。この「脳」は大規模言語モデル（LLM）であり、学習するためには山のような情報を処理しなければなりません。しかし、ここには落とし穴があります。学べば学ぶほど、この脳はコンピュータのメモリと電気をより多く欲しがるようになるのです。それは、巨大な高精細の映画を、小さくて古いノートパソコンで再生しようとしているようなものです。画面はぼやけ、ファンは悲鳴を上げます。これを解決するために、科学者たちはコンピュータが計算に使う数字を小さくしようと試みてきました。非常に精密で重い数字（詳細な写真のようなもの）を使う代わりに、より小さく軽い数字（ピクセル化されたスケッチのようなもの）へと切り替えてきたのです。これは「低精度学習（low-precision training）」と呼ばれます。最近では、「FP4」（4ビット浮動小数点）と呼ばれる新しい種類の極小の数字が登場し、これらの巨大な脳をより小さく、より速くすることを約束しています。しかし、問題があります。脳の主要な部分はこれらの極小の数字を扱えるのですが、「補助」となる部分――例えば、学習のために保持するメモリや、文章を理解するための複雑な計算――は、極小の数字を強制されると、クラッシュしたり混乱したりしてしまうのです。それは、レーシングカーのエンジンを持っているのに、自転車のチェーンで走ろうとしているようなものです。エンジンは準備万端ですが、他の機械の部分が追いつけません。

この論文は、これらのデジタル脳のための熟練したメカニックとして機能する、Full-Stack FP4という巧妙な新しいツールキットを紹介しています。研究者たちは、脳のすべての部分に対して同じ極小数字のルールを適用することはできないということに気づきました。主要な接続部分のような部分は、極小の数字でも十分に扱えます。しかし、昨日学んだことを記憶するための「メモリ」のような他の部分は非常に敏感であり、特別なケアを必要とします。チームは、脳の各部分に対して異なる「レシピ」を使用するモジュール式のシステムを作成しました。主要な接続については、LoRA-SVDというトリックを使用します。これは、最も重要な詳細については高解像度のスケッチを維持しつつ、それ以外の部分にはピクセル化されたブロックを使用するような手法です。メモリについては、数字が紛失しないように、極小の箱に押し込む前に数字を「滑らかにする」方法を考案しました。そして、脳が適切な言葉に集中するのを助ける計算（アテンション）については、最も敏感な部分を高精細のまま維持し、残りの部分はピクセル化された状態のままにすることに決めました。

彼らがこのフルシステムを、30億パラメータのモデル（中規模のデジタル脳）と640億語の学習データを用いてテストしたところ、結果は素晴らしいものでした。「Full-Stack FP4」を備えた脳は、従来の重量級バージョンとほぼ全く同じように学習しました。学習スコアの差はわずか0.838%であり、ほとんど気づかないレベルです。実際、見たことがない質問に対して脳がどれだけうまく答えられるかをテストしたところ、新しい極小数字バージョンは、はるかに少ないメモリを使用しているにもかかわらず、正解の単語を予測する能力がむしろわずかに向上していました。単一の強力なグラフィックスカード（RTX 5090）において、この新しい手法は、特定のタスクにおいて脳の学習プロセスを2.5倍から2.8倍高速化し、標準的な高精度手法よりもメモリを38%から42%削減しました。

研究者たちは、これが30億パラメータまでのモデルには有効であるものの、さらに大きな脳や、より長い学習時間に対しても機能するかどうかはまだ証明されていないという点に注意深く言及しています。また、単にすべてのものに極小の数字を投げつければよいわけではないことも分かりました。もし、これらの特別なレシピなしに、脳の最も敏感な部分に極小の数字を強制しようとすれば、学習は崩壊してしまいます。しかし、現時点では、この「フルスタック」のアプローチは、システム全体に一律のサイズを強いるのではなく、脳の異なる部分に対して実際に必要な精度レベルを個別に扱うことで、よりスマートで、速く、安価なAIを構築できることを示しています。

技術要約: Full-Stack FP4: 量子化された射影、オプティマイザ、およびアテンションを用いた安定したLLM事前学習

1. 問題提起

NVIDIA Blackwell GPUにおけるネイティブなNVFP4（4ビット）Tensor Coreサポートの進展により、大規模言語モデル（LLM）の事前学習における実用的なターゲットとして浮上しています。既存の研究では、Transformerの線形射影（Linear Projections）における安定した4ビット事前学習が成功していますが、他の重要な学習モジュールには依然として大きなギャップが存在します。現在のほとんどのNVFP4レシピは、オプティマイザの状態、オプティマイザの計算、およびアテンション機構に対して、より高精度な形式（例：BF16またはFP8）を維持しています。

主な課題は、これらのモジュールが、単一の量子化ルールでは対処できない独自の数値構造を持っていることです。

線形射影 (Linear Projections): 学習が進むにつれて勾配信号が弱まり、量子化ノイズが一定であるため、低精度学習の有効性が低下します。
AdamWの状態: 二次モーメント（second-moment momentum）は非負であり、ヘビーテイル（重い裾）を持ち、持続的で、分母の値に敏感であるため、直接的な4ビット保存は歪みが生じやすくなります。
Root/Muon オプティマイザ: これらは繰り返しのNewton–Schulz行列積に依存しており、低精度の反復誤差が蓄積して収束を不安定にする可能性があります。
アテンション (Attention): この演算子は、ソフトマックスに敏感なパス（例：$PV$, $dOV^\top$ ）を含む一方で、他のパス（例： $QK^\top$ ）は量子化に対してより堅牢です。

その結果、これら（射影、オプティマイザ、アテンション）の全工程において、混合精度へのフォールバックに頼ることなく、安定したネイティブNVFP4事前学習を可能にする、統一されたモジュール式のフレームワークが欠如しています。

2. 手法: Full-Stack FP4

著者らは、各学習モジュールの特定の数値的問題に対処するために設計された、4つの独立した構成可能なレシピからなるモジュール式フレームワーク、Full-Stack FP4を提案しています。

2.1 LoRA-SVD 線形射影

学習後半における勾配信号の劣化を軽減するために、著者らはLoRA-SVDを提案しています。

メカニズム: この手法は、重みを $W = W_{res} + \beta L_2 L_1$ とパラメータ化します。主要な部分空間（ $L_2 L_1$ ）をBF16で保持しつつ、支配的な残差（ $W_{res}$ ）をNVFP4で計算します。
バックプロパゲーション (Backward Pass): 部分空間に制約を設けたバックプロパゲーションでは、Cholesky-QR制約を使用してBF16の低ランクブランチの更新が主要部分空間と一致するようにし、残差はNVFP4コンポーネントを更新します。
再整列 (Realignment): 部分空間は、ランダム化SVDを用いて定期的に（2,048ステップごと）再整列され、BF16の部分空間が常に最適であることを保証します。

2.2 量子化状態 AdamW (Q-AdamW)

持続的でヘビーテイルなモーメント状態のNVFP4保存を安定させるために：

一次モーメント ( $m_t$ ): 確率的丸め（Stochastic Rounding; SR）および適応的な4/6ビット量子化の前に、分布を滑らかにするためのHadamard混合を適用します。
二次モーメント ( $v_t$ ): NVFP4保存の前に、以下の3段階のプリプロセッシング・パイプラインを利用します。
1. 平方根 (Square Root): ヘビーテイルを短縮します。
2. タイル平均 (Tile Mean): 正の共通成分（BF16で保存）と、近似的にゼロ平均の残差を分離します。
3. Hadamard変換 (Hadamard Transform): 残差内の局所的なエネルギーを拡散させ、より対称的にします。
再構成 (Reconstruction): 変換と平均分離を逆転させることで状態を再構成し、その結果を二乗します。このパイプラインは、SVDベースの部分空間操作を必要としません。

2.3 NVFP4 Root オプティマイザ (Q-Root)

Rootオプティマイザで使用されるNewton–Schulz (NS) 反復に対して：

直接NVFP4実行: 混合精度ルートではなく、NS行列積を直接NVFP4で実行します。
安定化: 安定性は、形状依存の係数（特定の行列アスペクト比に対して最適化された $a, b, c$ ）とp99アウトライヤー・クリッピングによって達成されます。このアプローチは、補助的な混合精度パスを必要とせずに、連続するNSステップ間の誤差増幅を抑制します。

2.4 混合精度アテンション (Q-Attn)

精度と効率のバランスを取るために：

精度の分離: フレームワークは、堅牢なパス（ $QK^\top$ および $dS$）にNVFP4を割り当て、ソフトマックスに敏感なパス（$PV$, $P^\top dO$ , $dOV^\top$ ）にはBF16を保持します。
一貫性: 量子化された $Q, K, V$ のビューをフォワードパスとバックプロパゲーションパスの間で再利用することで、テンソルの整合性を確保し、冗長なメモリトラフィックを回避します。

3. 主な貢献

LoRA-SVD 射影: 標準的なNVFP4と比較して、線形のみの損失ギャップを1.40%から0.61%に低減する、マージ可能な再パラメータ化を実現しました。
量子化状態 AdamW: 変換された状態パイプラインを利用して二次モーメントの再構成誤差を約60%削減し、統合された事前学習実行における安定したNVFP4 AdamWモーメント状態の保存を初めて実証しました。
NVFP4 Root: 形状認識係数とクリッピングによって可能となった、直接的な低精度Newton–Schulzパスを提供し、反復誤差を約10%から約3%に低減しました。
学習可能な混合精度アテンション: ソフトマックスに敏感なパスをBF16で保護しつつ、適切なフォワードおよびバックワードパスを量子化する戦略により、一貫したNVFP4事前学習オペレータを実現しました。

4. 実験結果

著者らは、64Bトークンで学習された3BパラメータのTransformerを用いてフレームワークを評価しました。

学習損失 (Training Loss): BF16のベースラインが2.267に達したのに対し、Full-Stack FP4は2.286に達し、相対的なギャップは**0.838%**でした。
ダウンストリーム性能: ゼロショット評価において、Full-Stack FP4の平均パープレキシティは26.665（BF16は26.675）であり、平均精度はBF16ベースラインよりわずか0.10パーセントポイント低い結果となりました。
アブレーション研究:
- LoRA-SVD単体では、線形のみの損失ギャップを1.40%から0.61%に低減しました。
- これらのモジュール式レシピは、独立して有効化することも、組み合わせて使用することも可能であり、「Full-Stack」はその集約的なカバレッジを表しています。
ネイティブハードウェア効率 (RTX 5090):
- Rootの高速化: Rootオプティマイザのステップにおいて、Full-Stack FP4は最適化されたBF16に対して2.50倍〜2.83倍、NVIDIA Transformer Engine (TE) に対して1.83倍〜2.17倍の高速化を達成しました。
- メモリ節約: 本フレームワークは、BF16と比較してAdamWのピークメモリを**37.9%〜42.5%**削減しました。
- 時間節約: ステップ時間の短縮と、BF16の損失に一致させるために必要なトークン数のわずかな増加を考慮すると、一致した損失を得るための推定時間は**56.9%〜61.9%**削減されました。

5. 意義と限界

意義:
本論文は、射影、オプティマイザの状態、オプティマイザの計算、およびアテンションをカバーする「Full-Stack」なNVFP4事前学習が実現可能であることを示しています。これは、BF16と同等の性能（損失ギャップ1%未満）を達成しつつ、Blackwellクラスのハードウェアにおいて大幅なメモリおよびスループットのメリットを提供できることを確立しました。本研究は、実務者が段階的に低精度を採用できるようにするためのモジュール式のレシピセットを提供しています。

限界:

スケール: 実験は最大3Bパラメータおよび64Bトークンのモデルに限定されています。本論文は、7B以上のスケールや、より長い学習期間における挙動を確立していません。
効率の測定: ネイティブハードウェア効率は、完全な分散モデルではなく、単一のRTX 5090上の4つのデコーダブロックで測定されました。
推定値: 時間節約の計算は、固定されたバッチサイズに基づく推定値であり、解放されたメモリを使用してバッチサイズを増やした場合のスループット向上の可能性を考慮していません。
カーネル融合 (Kernel Fusion): 一部のカーネル（LoRA backward、attention）はまだ完全に融合されておらず、実装はまだすべてのFlashAttention-4の最適化を活用していません。
フォーマットの特定性: これらのレシピと測定値はNVFP4に特化したものであり、MXFP4への転用可能性を主張するものではありません。