FP8で大規模なモデルをトレーニングするのはどれくらい良いですか? Microsoft:BF16よりも64%高速、42%少ないメモリ

出典: Heart of the Machine

画像ソース: Unbounded AIによって生成

低精度のトレーニングは、大規模なモデルのトレーニングでモデルのサイズをスケーリングし、トレーニング コストを節約するための重要なテクノロジの 1 つです。 現在の 16 ビットおよび 32 ビット浮動小数点混合精度トレーニングと比較して、FP8 8 ビット浮動小数点混合精度トレーニングを使用すると、2 倍の速度を実現し、ビデオ メモリの 50% - 75% と通信コストの 50% - 75% を節約でき、NVIDIA の最新世代の Kawang H100 には優れた FP8 ハードウェア サポートが付属しています。 しかし、FP8トレーニングのサポートは、業界の大規模モデルトレーニングフレームワークではまだ非常に限られています。 最近、Microsoftは、LLMをトレーニングするためのFP8混合精度フレームワークであるFP8-LMを提案し、FP8を大規模モデルトレーニングのコンピューティング、ストレージ、および通信に可能な限り適用し、H100を使用してGPT-175BをBF16よりも64%高速にトレーニングし、メモリ使用量を42%節約しました。 さらに、オープンソースです。

大規模言語モデル(LLM)は、これまでにない言語理解と生成機能を備えていますが、これらの高度な機能を解き放つには、巨大なモデルサイズと計算負荷の高いトレーニングが必要です。 この文脈では、特にOpenAIが提案するスーパーインテリジェンスモデルの規模へのスケーリングを見ると、低精度のトレーニングは最も効果的で重要な手法の1つであり、メモリフットプリントが小さく、トレーニング速度が速く、通信オーバーヘッドが低いなどの利点があります。 Megatron-LM、MetaSeq、Colossal-AIなど、現在のほとんどのトレーニングフレームワークは、LLMのトレーニングにデフォルトでFP32フル精度またはFP16/BF16混合精度を使用しています。

しかし、それでも限界にとらわれず、NVIDIA H100 GPU のリリースにより、FP8 は低精度の特性評価のための次世代のデータ型になりつつあります。 理論的には、FP8は、現在のFP16/BF16浮動小数点混合精度トレーニングと比較して2倍の速度向上を実現し、メモリコストを50%から75%、通信コストを50%から75%節約できます。

それにもかかわらず、FP8トレーニングのサポートは現在制限されています。 NVIDIA の Transformer Engine (TE) は、GEMM の計算にのみ FP8 を使用し、エンドツーエンドのアクセラレーション、メモリ、および通信コストの節約を限定的に提供します。

しかし、現在、Microsoft のオープンソースの FP8-LM FP8 混合精度フレームワークは、この問題を劇的に解決します: FP8-LM フレームワークは、フォワード パスとバックワード パッシングのトレーニング全体で FP8 形式を使用するように高度に最適化されており、システムの計算、メモリ、および通信のオーバーヘッドを大幅に削減します。

*住所：

• オープンソースフレームワーク:

実験結果によると、H100 GPUプラットフォームでGPT-175Bモデルをトレーニングする場合、FP8-LM混合精度トレーニングフレームワークは実際のメモリフットプリントを42%削減するだけでなく、広く採用されているBF16フレームワーク(Megatron-LMなど)よりも64%高速に動作し、Nvidia Transformer Engineよりも17%高速に動作します。 さらに、事前学習や複数のダウンストリームタスクでは、FP8-LM学習フレームワークを使用して、現在の標準BF16混合精度フレームワークと同様の結果を持つモデルを取得できます。

コンピューティングリソースを考えると、FP8-LMフレームワークを使用すると、トレーニング可能なモデルのサイズを最大2.5倍まで簡単に増やすことができます。 GPT-5をFP8で学習させた場合、同量のH100しか使用してもモデルサイズがGPT-4の2.5倍になるという議論がTwitterで議論されている開発者もいます!

ハギングフェイスのR&Dエンジニアは、「それはクールだ。FP8のマストレーニング技術を使えば、計算による不正行為を実現できる!」と皮肉を言った。

FP8-LMの主な貢献:

• 新しい FP8 混合精度トレーニング フレームワーク。 8ビットの重み、勾配、オプティマイザー、分散トレーニングを徐々に解き放つ加算的な方法で簡単に使用できます。 この 8 ビット フレームワークは、ハイパーパラメーターとトレーニング メソッドを変更することなく、既存の 16/32 ビット混合精度メソッドの対応する部分を単純かつ簡単に置き換えることができます。 さらに、Microsoftチームは、ユーザーが少量のコードで8ビットの低精度トレーニングを実装できるPyTorch実装をリリースしました。
• FP8でトレーニングされたGPTスタイルのモデルのファミリー。 新たに提案したFP8スキームを用いて、SFTやRLHFなどのGPTの事前学習と微調整を行ったところ、70億から1750億のパラメータを持つあらゆるサイズのモデルに対応できる可能性が示されました。 FP8は、テンソル、パイプライン、シーケンス並列化などの一般的な並列コンピューティングパラダイムをサポートしているため、ユーザーはFP8を使用して大規模なベースモデルをトレーニングできます。 また、Megatron-LMをベースにした最初のFP8 GPTトレーニングコードベースをオープンソースでリリースしました。

FP8-LMの実装

具体的には、FP8 を使用して混合精度と分散トレーニングを簡素化することを目標に、3 つのレベルの最適化を設計しました。 これら 3 つのレベルは、8 ビットの集合通信オプティマイザと分散並列学習を段階的に統合できます。 最適化レベルが高いほど、LLMトレーニングで使用されるFP8が多くなります。

さらに、大規模なトレーニング(数千のGPU上のGPT-175Bなど)では、テンソル、トレーニングパイプライン、トレーニングなど、FP8精度の低桁並列化を提供し、次世代の低精度並列トレーニングへの道を開きます。

テンソル並列化とは、モデルのレイヤーを複数のデバイスに分散させ、重み、勾配、活性化テンソルのシャードを異なる GPU に配置することです。

テンソル並列化で FP8 をサポートするために、Microsoft チームはシャードの重みと活性化テンソルを線形層計算用の FP8 形式に変換し、FP8 が前方計算と後方勾配集合通信に使用されるようにしました。

一方、シーケンスの並列化では、入力シーケンスを複数のチャンクにスライスし、サブシーケンスを異なるデバイスにフィードしてアクティベーションメモリを節約します。

図2に示すように、シーケンスの並列化とテンソルの並列化は、利用可能なメモリを最大限に活用し、トレーニング効率を向上させるために、Transformerモデルのさまざまな部分で実行されています。

ZeRO(Zero Redundancy Optimizer)の場合、FP8パーティショニングに伴うスケーリングファクタの扱いが難しいため、FP8を直接適用することはできません。 したがって、各テンソルのスケーリング係数は、FP8が分割される方法に沿って分布する必要があります。

この問題を解決するために、研究者らは、ZeROアプローチのように複数のサブテンソルに分割するのではなく、各テンソルを全体として複数のデバイスに分散させる新しいFP8割り当てスキームを実装しました。 このメソッドは、アルゴリズム 1 に示すように、FP8 テンソルの割り当てを貪欲な方法で処理します。

具体的には、このメソッドは最初にサイズに基づいてモデル状態のテンソルを並べ替え、次に各 GPU の残りのメモリ量に基づいてテンソルを異なる GPU に割り当てます。 この割り当ては、より多くのメモリが残っている GPU が新しく割り当てられたテンソルを受け取る可能性が高いという原則に従います。 このようにして、テンソルのスケーリング係数をテンソルに沿ってスムーズに割り当てることができ、通信と計算の複雑さを軽減できます。 図3は、ZeROテンソルの分割方法の違いを示しています。

FP8でLLMをトレーニングするのは簡単ではありません。 データのアンダーフローやオーバーフローなど、多くの困難な問題が関係しています。 また、ダイナミックレンジの狭さやFP8データフォーマット特有の精度低下による量子化誤差もあります。 これらの課題は、トレーニングプロセスにおける数値の不安定性と不可逆的な乖離につながる可能性があります。 これらの問題に対処するために、Microsoft は、重要な情報の損失を防ぐために、精密な分離と自動スケーリングの 2 つのテクノロジを提案しました。

精密デカップリング

精度のデカップリングでは、重み、勾配、オプティマイザの状態などのパラメーターに対するデータ精度の影響をデカップリングし、精度の影響を受けないコンポーネントに精度の低下を割り当てます。

精度のデカップリングについては、勾配統計では低い精度で済むが、ソブリン重みでは高い精度が求められるという指針を見つけたと研究チームは述べている。

より具体的には、1次の勾配モーメントはより高い量子化誤差を許容でき、低精度のFP8を装備できますが、2次のモーメントはより高い精度を必要とします。 これは、Adam を使用する場合、モデルの更新時に勾配の方向が振幅よりも重要になるためです。 テンソルスケーリング機能を備えたFP8は、高精度テンソルとして1次モーメントの分布を効果的に保持できますが、精度がいくらか低下します。 通常、勾配値は小さいため、2 次の勾配モーメントの勾配の 2 乗を計算すると、データのオーバーフローの問題が発生する可能性があります。 したがって、数値精度を維持するためには、より高い 16 ビット精度を割り当てる必要があります。

その一方で、主権の重みを維持するために高精度を使用することも重要であることもわかりました。 根本的な理由は、トレーニング中にウェイトの更新が非常に大きくなったり、非常に小さくなったりすることがあり、ソブリンウェイトの場合、精度が高いほどウェイトが更新されたときに情報が失われるのを防ぎ、より安定した正確なトレーニングが可能になるためです。

この実装では、ソブリン ヘビー級には、FP32 の完全精度を使用するか、テンソル スケーリングで FP16 を使用するかの 2 つの実行可能なオプションがあります。 テンソルスケーリングを使用したFP16の利点は、精度を犠牲にすることなくメモリを節約できることです。 したがって、新しいフレームワークの既定の選択は、テンソル スケーリングで FP16 を使用して、オプティマイザーにソブリン重みを格納することです。 トレーニングでは、FP8 混合精度オプティマイザーの場合、各パラメーターに 6 バイトのメモリが必要です。

この新しい低桁オプティマイザは、以前のソリューションと比較してメモリフットプリントを2.6倍に削減します。 これは、LLMトレーニング用の最初のFP8オプティマイザーです。 実験によると、FP8オプティマイザーは、1億2,500万から1,750億のパラメータの範囲の幅広いモデルサイズでモデルの精度を維持することが示されています。

自動スケーリング

自動スケーリングは、勾配値をFP8データ形式の表現範囲に保存することですが、これにはテンソルスケーリング係数の動的な調整が必要であり、all-reduce通信中のデータのアンダーフローとオーバーフローを減らすことができます。

具体的には、トレーニング中の状況に応じて変更できる自動スケーリング係数μを導入しました。

実験結果

新たに提案されたFP8低精度フレームワークを検証するために、研究者らは、事前学習や教師あり微調整(SFT)など、GPTスタイルのモデルの学習にFP8を使用する実験を行いました。 この実験は、Azure クラウド コンピューティング用の最新の NDv5 H100 スーパーコンピューティング プラットフォームで実施されました。

実験結果は、新しいFP8法が効果的であることを示しています:以前の広く使用されているBF16混合精度トレーニング法と比較して、新しい方法には、実メモリ使用量の27%〜42%の削減(たとえば、GPT-7Bモデルでは27%、GPT-175Bモデルでは42%の減少)などの明らかな利点があります。 重み付けされた勾配通信のオーバーヘッドは63%から65%減少しました。

学習率や重み減衰などのハイパーパラメーターを変更しなくても、FP8 でトレーニングされたモデルのパフォーマンスは、事前トレーニング済みのタスクか下流のタスクかに関係なく、BF16 の精度でトレーニングされたモデルのパフォーマンスに匹敵します。 GPT-175Bモデルのトレーニング中に、新しく提案されたFP8混合精度フレームワークは、TEメソッドと比較してH100 GPUプラットフォームでトレーニング時間を17%短縮し、メモリフットプリントを21%削減できることは注目に値します。 さらに、図1に示すように、モデルのサイズが拡大し続けるにつれて、低精度のFP8を使用することでさらにコストを削減できます。

微調整では、FP8混合精度を使用して命令の微調整と、ヒューマンフィードバック(RLHF)を使用した強化学習を行い、事前学習済みのLLMを端末タスクやユーザーの好みに合わせました。

その結果、AlpacaベンチマークとMT-Benchベンチマークでは、FP8の混合精度微調整を使用したモデルのパフォーマンスが、半精度BF16の微調整を使用したモデルのパフォーマンスに匹敵し、FP8を使用したトレーニングも27%高速であることがわかりました。 さらに、FP8の混合精度は、トレーニング中に複数のモデルを読み込む必要があるプロセスであるRLHFに大きな可能性を示しています。 FP8をトレーニングに使用することで、人気のあるRLHFフレームワークであるAlpacaFarmは、モデルの重みを46%削減し、オプティマイザ状態のメモリ消費量を62%削減することができました。 これは、新たに提案されたFP8低精度トレーニングフレームワークの汎用性と適応性をさらに実証しています。

また、アブレーション実験を行い、各成分の有効性を検証しました。

FP8の低精度学習は、今後、大型モデル開発の新たなインフラとなることが予想されます。