技術サマリー: FORGE (Fused On-Register Gradient Elimination)

1. 問題点: 勾配のフロア（Gradient Floor）

現在のメモリ効率の高い大規模言語モデル（LLM）の学習手法は、主にGPU上の4つのメモリプール（オプティマイザ状態、アクティベーション、モデル重み、および実体化された重み勾配）のうち、3つを対象としています。しかし、これらは4番目のプールである**実体化された重み勾配（materialized weight gradient）**を未だに手付かずのままにしています。

標準的な逆モード自動微分では、バックプロパゲーション（逆伝播）においてすべてのレイヤーのフル勾配テンソル ($\nabla W$) を計算し、それを高帯域幅メモリ（HBM）に書き込み、その後オプティマイザが重みを更新するためにそれを読み戻します。この二段階のスケジュールは、「勾配のフロア」を生み出します。つまり、バックプロパゲーションとオプティマイザ・ステップの境界において、すべてのレイヤーの勾配がメモリ上に共存しなければならないという現象です。8B（80億）パラメータのモデルの場合、この一時的な勾配プールだけで約16 GiBを消費し、これが学習構成を単一のGPUに収めることができるかどうかの決定要因となります。既存の手法（LOMOやFlashOptimなど）は、バックプロパゲーションとオプティマイザ・ステップをレイヤーまたはパラメータ全体の粒度で融合させていますが、これらにおいても、更新が行われる前にパラメータのフル勾配テンソルが依然としてメモリ上に実体化されています。

2. 手法: FORGE

FORGE (Fused On-Register Gradient Elimination) は、重み勾配の計算に直接オプティマイザ・ステップを融合させることで、タイル粒度（具体的には $128 \times 128$ タイル）において、実体化された勾配テンソルを完全に排除します。

コアメカニズム

FORGEは、フル勾配行列を計算して保存する代わりに、重み行列を一度に1つのタイルずつ処理します：

1. タイル単位の勾配累積: 特定の重みタイルに対して、アップストリームの勾配 ($\Delta Y$) と入力アクティベーション ($X$) の部分積を、FP32 GPUレジスタに直接累積します。
2. インレジスタ最適化: タイルの勾配がレジスタ内で完全に累積されると、そのタイルに対してオプティマイザの更新（例：AdamW）がレジスタ内ですぐに適用されます。
3. 即時破棄: 勾配タイルは生成された瞬間に消費されます。それはテンソルとしてHBMに書き込まれることはありません。
4. ライトバック: 更新された重みタイルと、（オプションの量子化された）オプティマイザ状態タイル（モーメント）のみがグローバルメモリに書き戻されます。

精度と忠実度

• 厳密性: 要素ごとに分離可能なオプティマイザ（各重みの更新がその重みの座標のみに依存するもの。例：AdamW, SGD, Lion）については、タイル単位の更新は、標準的なフル勾配ステップと数学的に同一です。唯一の差異は、トークンに対する浮動小数点加算の順序ですが、これはあらゆるGEMM実装に固有のものであり、オプティマイザのロジックの近似を意味するものではありません。
• 低精度における忠実度: BF16またはINT8の学習レジームにおいて、FORGEは以下の2つの方法でフル精度の忠実度を維持します：
  • 勾配は最適化の前にBF16に切り詰められることなく、FP32レジスタ内で計算・消費されます。
  • 「デッドゾーン」（BF16において小さな更新がゼロに丸められてしまう現象）を防ぐため、FORGEは重みのライトバックに**確率的丸め（stochastic rounding）**を採用しており、これによりサブULP（最小単位）の更新の偏りのない蓄積を保証します。
  • INT8オプティマイザ状態については、量子化誤差を最小限に抑えるために、標準的な8ビットオプティマイザ（ブロックサイズ2048）よりも小さなブロックサイズ（64）を使用します。

分散学習との互換性

FORGEは、テンソル並列（TP）およびシーケンス並列（SP）の下で証明可能な厳密性を持ちます。なぜなら、シャードのローカル勾配はバックプロパゲーション時に完結しているからです。しかし、現在の形式ではデータ並列（DP）やコンテキスト並列とは互換性がありません。これらは更新を適用する前に部分的な勾配のAll-reduceを必要とするため、勾配テンソルの実体化が必要になるからです。

3. 主な貢献

1. 圧縮ではなく構造的排除: FORGEは、実体化された勾配を数学的な必然性ではなく、学習スケジュールのアーティファクト（副産物）として特定しました。これにより、$2P$（パラメータ数）の勾配メモリ項を排除し、勾配プールをモデルサイズに依存しない $O(T_R T_C)$ （タイルサイズ）へと削減します。
2. 厳密性の証明: 本論文は、要素ごとに分離可能なオプティマイザに対して、タイル融合スケジュールが標準的なFP32勾配上のAdamWと同一の更新を行うことを証明しています。この厳密性は、テンソル並列およびシーケンス並列の下でも保持されます。
3. 低精度におけるフル精度への忠実度: 勾配をFP32レジスタ内に保持することで、FORGEは勾配をBF16に切り詰めてから最適化することによって生じるバイアスを回避します。確率的丸めと組み合わせることで、BF16やINT8の状態を使用している場合でも、FP32マスターとの収束性を一致させます。
4. 汎用性: FORGEは線形層に対してアーキテクチャに依存せず、8つの完全融合された要素別オプティマイザ（AdamW, NAdam, RAdam, Lion, RMSprop, AdaGrad, SGD, SGD-momentum）をサポートしています。

4. 実験結果

NVIDIA H200およびB200 GPUを用い、0.6Bから14Bパラメータの範囲のモデル（Llama-3.1-8B, Qwen3ファミリー）で評価を行いました：

• メモリ削減: FORGEは、融合AdamWと比較してピークメモリを38–53%削減します。Llama-3.1-8Bの場合、ピークメモリは約75 GiBから35.4 GiBへと減少します（53%の削減）。
• スループット: FORGEは、個別のオプティマイザカーネルの起動と勾配のHBMラウンドトリップを排除することにより、ステップタイムにおいて1.52倍の高速化（99.1 ms vs 150.9 ms）を達成しました。
• スケーラビリティ:
  • 単一のH200上で、FORGEはQwen3-14B（標準的な融合AdamWではメモリ不足（OOM）が発生するモデル）を63.2 GiBで学習可能です。
  • 4-GPU A100-40GB構成において、FORGEはマイクロバッチサイズ4（24.6 GiBに適合）を可能にしますが、標準的な融合AdamWはマイクロバッチ1しか適合しません。これにより、FORGEは同じハードウェア上で標準的なベースラインよりも2.42倍高いスループットを達成できます。
• 収束性: FORAGEは、継続的な事前学習タスクにおいて、融合AdamWの収束を0.001 nats（平均偏差）以内で追跡しており、低精度の偏差が実行間のノイズ範囲内であることを確認しています。
• 構成可能性: FORGEは、FP8量子化およびアクティベーション・チェックポインティングとクリーンに構成可能です。標準的なオプティマイザが通常受けるデクオンタイゼーション（逆量子化）のコストを吸収します。

5. 意義と主張

本論文は、FORGEが漸進的な最適化ではなく、構造的な転換であることを主張しています。勾配のフロアを取り除くことで、オプティマイザのメモリオーバーヘッドを、より大きなバッチサイズ、より長いシーケンス、または既存のハードウェア上でのより大きなモデルのための利用可能なヘッドルームへと変換します。

• 能力の境界: 著者らは、勾配のフロアの排除は単なる定数倍の改善ではなく、能力の境界線であると述べています。例えば、単一のH200における14Bモデルの標準的な学習構成は不可能ですが、FORGEはこれを可能にします。
• 直交性: この手法は、他のメモリ節約技術（状態の量子化やアクティベーションの再計算など）と直交しており、それらと組み合わせてメモリフットプリントをさらに削減できます。
• 制限事項: 現在のところ、この手法は要素ごとに分離可能なオプティマイザに限定されており、修正なしではデータ並列をサポートしません。また、入力勾配を重み更新の前に読み取る必要があるため、単一のステップで複数の勾配消費者を持つタイド・エンベディング（tied embeddings）やMoEエキスパートは除外されます。

要約すると、FORGEは「勾配のフロア」が現行の学習ループにおける不要なアーティファクトであることを示しています。タイルレベルでバックプロパゲーションにオプティマイザを融合させることで、広範なクラスのオプティマイザに対して数学的な厳密性を維持しながら、大幅なメモリ削減と高速化を実現しています。