Thinking Machines が Nvidia GB300 ハードウェアでのエンドツーエンド 4 ビットトレーニングを詳細に説明します

Thinking Machines Lab は、エンドツーエンドの 4 ビットトレーニングに関する技術的な詳細を公開し、NVFP4 が Nvidia GB300 ハードウェアでの LLM トレーニングをどのように高速化するかを説明しました。

元 OpenAI 最高技術責任者のミラムラティ氏が設立した調査会社 Thinking Machines Lab は、大規模な言語モデルをほぼ完全に 4 ビット精度でトレーニングする方法についての詳細な技術的説明を公開しました。同社によれば、この技術により Nvidia の最新データセンターハードウェアの速度とメモリが大幅に節約されます。

「エンドツーエンドの 4 ビットトレーニング」と題され、11 月 13 日に同社の Web サイトに公開されたこのブログ投稿では、低精度トレーニングの背後にある数値的トレードオフについて説明し、チームが Inkling 2 Small モデルのトレーニング中にどのようにそれらを管理したかについて説明しています。最新の AI 開発を追いかけている読者のために、この投稿では、フロンティア規模のトレーニングの実行に実際にどれくらいの費用がかかるかを決定するエンジニアリング上の決定について、めったに公開されることはありません。

ビット数が少ないほどトレーニングが高速になる理由

大規模な言語モデルをトレーニングするには、膨大な量の数値を計算して保存する必要があります。各値の使用ビット数を減らすと、メモリ使用量が削減され、GPU の演算が高速化されます。犠牲となるのは精度です。ビットが少なくなると、表現できる数値が少なくなるため、値をより積極的に丸める必要があります。

同社はこの問題を簡単な例で説明している。形式が数値 3 と 4 しか表現できない場合、3.2 を格納するには 3 に丸める必要があり、量子化誤差 0.2 が発生します。単一の丸め誤差は些細なものですが、トレーニング中にこれらの誤差はレイヤーや最適化ステップ全体に蓄積され、学習プロセス自体に干渉する可能性があります。

「これが低精度トレーニングの中心的な課題です。学習を妨げるエラーを引き起こすことなく、高速な演算から恩恵を受けるにはどうすればよいでしょうか?」と同社は投稿に書いた。

3 つのフォーマット、1 つのトレーニング実行

Thinking Machines は、単一の精度レベルにこだわるのではなく、トレーニングのさまざまな段階で 3 つの数値形式 (BF16、MXFP8、および NVFP4) を使用します。各形式の精度は固定されていますが、その速度はそれを実行しているハードウェアによって異なります。

同社は、NVFP4 での行列乗算を高速化するために特別に設計された Nvidia の Grace Blackwell Ultra プラットフォーム GB300 でトレーニングを行っています。投稿によると、NVFP4 は BF16 よりも 3.6 倍少ないビット数で値を保存し、GB300 システムでは最大 6 倍高速に動作します。同社は、生のスループットに関して、GB300 のパフォーマンスを NVFP4 で約 15 PFLOPS、MXFP8 で 5 PFLOPS、BF16 で 2.5 PFLOPS としています。

問題は、NVFP4 の積極的な圧縮により、より広範なフォーマットよりもはるかに多くの量子化エラーが発生することです。このエラーは、トレーニングの一部の段階では許容されますが、他の段階では許容できないため、同社は各段階を個別に扱います。

異なるステージ、異なるルール

このブログ投稿では、トレーニングを 4 つの段階 (事前トレーニングの前方パスと後方パス、強化学習サンプラー、RL トレーナー) に分割し、それぞれにエラーの最小化、バイアスの最小化、速度の向上、メモリ使用量の削減の間で異なる優先順位のバランスを割り当てています。

速度は、数兆のトークンを実行する事前トレーニングと、大量のロールアウトを生成する RL サンプリングといった大量のステップで最も重要です。ただし、後方パスは特殊なケースです。この投稿では、バックワードパスではバッチ内の多数のトークンの平均が計算されるため、ランダムに分散された丸め誤差の影響が弱められますが、値が一貫して一方向に丸められるときに発生する体系的なドリフトである量子化バイアスの影響を非常に受けやすいと説明されています。ランダムノイズとは異なり、そのバイアスはレイヤー、サンプル、トレーニングステップ全体で複雑になる可能性があります。

同社は、Apache 2.0ライセンスに基づいて今月初めにリリースしたオープンウェイトの専門家混合モデルであるInkling 2 Smallをトレーニングする際に、これらのトレードオフを管理したと述べた。この投稿では、実行のすべての部分が 4 ビットで実行されたとは主張していません。むしろ、各コンポーネントの精度がデフォルトで継承されるのではなく、意図的に選択されるエンドツーエンドの戦略について説明しています。

低精度トレーニングが重要な理由

ここでの経済学は単純です。コンピューティングはフロンティアモデル開発における主なコストであり、ハードウェアベンダーは新しいチップ世代ごとに低精度のアクセラレーションを追加するために競い合ってきました。 Nvidia は、Blackwell Ultra プラットフォームでの NVFP4 サポートを目玉機能として位置付けており、AMD、Google などの競合アクセラレータは、独自の低精度フォーマットを宣伝しています。

研究機関がより多くのトレーニングパイプラインを 4 ビット演算に安全にプッシュできれば、同じハードウェア予算でより多くの実験、より多くのトークン、またはより大きなモデルを購入できます。これは、最先端のアクセラレータへのアクセスが不足し、高価である市場においては、大きな利点となります。また、これは、固定ハードウェアからより多くの機能を絞り出す、量子化を意識した推論や投機的デコードなどの効率化技術に対する業界の広範な推進とも重なっています。

この投稿は、研究コミュニティがどこに向かっているのかを示すものでもある。丸め誤差により学習が損なわれると考えられていたため、低精度のトレーニングは推論を超えるものとしてはリスクが高すぎると長い間考えられていました。 Thinking Machines が公開した分析のように、段階ごとに注意深く分析することで、その仮定がどのようにテストされ、可能であれば覆されるかがわかります。

ブログ投稿からオープンウェイトまで

この開示が注目に値するのは、それが純粋に理論的なものではないということです。同社はこの技術を実際に出荷したモデルに直接結びつけました。 Inkling 2 Small は 11 月初旬にダウンロード可能なウェイトとともにリリースされ、新しい投稿ではトレーニングの背後にある数値戦略について説明しています。モデルをダウンロードした実践者は、事実上、同社が説明する低精度の方法論の実用的な成果物を入手することになります。

Thinking Machines は、最先端規模のラボとしては異例のオープンさで、モデルのリリースと並行して技術的なブログ記事を公開しています。この透明性により、外部の研究者は珍しいものを得ることができます。それは、現代のトレーニング実行が混合精度演算の厄介な数値現実をどのように処理するかについて文書化された再現可能な説明です。

より広範な業界にとって、4 ビットトレーニングは研究への好奇心から生産技術へと移行したというメッセージです。同様の戦略を採用するラボが増えるにつれて、理論上のハードウェアのピークパフォーマンスと実際のトレーニングスループットとの間のギャップは縮小し、トレーニング可能なモデルのコストは低下し続けるでしょう。