FP8 を使用して大規模モデルをトレーニングするとどれくらい良いのでしょうか? Microsoft: BF16 より 64% 高速、メモリは 42% 削減

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

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

それにもかかわらず、FP8 トレーニングのサポートは現在制限されています。 NVIDIA の Transformer Engine (TE) は GEMM 計算に FP8 のみを使用するため、エンドツーエンドの高速化、メモリおよび通信コストの節約という点で得られる利点は非常に限られています。

しかし、現在では、Microsoft のオープン ソース FP8-LM FP8 混合精度フレームワークによってこの問題は大幅に解決されています。FP8-LM フレームワークは高度に最適化されており、トレーニングの順方向パスと逆方向パス全体で FP8 形式を使用しているため、システムのコンピューティング、ビデオ メモリ、通信のオーバーヘッドが大幅に削減されます。

• 論文アドレス: https://arxiv.org/abs/2310.18313
• オープンソースフレームワーク: https://github.com/Azure/MS-AMP

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

特定のコンピューティング リソース セットが与えられている場合、FP8-LM フレームワークを使用すると、トレーニング可能なモデルのサイズを最大 2.5 倍まで簡単に増やすことができます。一部の開発者は Twitter で次のように議論しました: GPT-5 が FP8 トレーニングを使用する場合、同じ数の H100 しか使用されなくても、モデル サイズは GPT-4 の 2.5 倍になります。

Huggingface の研究開発エンジニアは冗談を言った。「すごいですね! FP8 の大規模トレーニング技術により、計算による欺瞞が実現できるのです!」

FP8-LMの主な貢献:

• 新しい FP8 混合精度トレーニング フレームワーク。 8 ビットの重み、勾配、オプティマイザー、分散トレーニングを徐々に追加的にロック解除していくので、非常に便利です。この 8 ビット フレームワークは、ハイパーパラメータやトレーニング方法を変更することなく、既存の 16/32 ビット混合精度メソッドの対応する部分を簡単かつ直接置き換えることができます。さらに、Microsoft チームは、ユーザーが少量のコードで 8 ビットの低精度トレーニングを実行できるようにする PyTorch 実装もリリースしました。
• FP8 を使用してトレーニングされた GPT スタイルのモデル ファミリ。彼らは、新たに提案された FP8 スキームを使用して GPT の事前トレーニングと微調整 (SFT と RLHF を含む) を実行し、その結果、新しい方法は 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 は、一次モーメントの分布を高精度テンソルに効果的に保存できますが、ある程度の精度の低下も引き起こします。通常、勾配値は小さいため、2 次勾配モーメントの勾配の 2 乗を計算すると、データ アンダーフローの問題が発生する可能性があります。したがって、数値の精度を維持するには、より高い 16 ビットの精度を割り当てる必要があります。

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

この実装では、マスター ウェイトには FP32 フル精度を使用するか、テンソル スケーリングで FP16 を使用するかの 2 つのオプションがあります。テンソル スケーリングを使用した FP16 の利点は、精度を犠牲にすることなくメモリを節約できることです。したがって、新しいフレームワークのデフォルトの選択は、テンソル スケーリングを備えた FP16 を使用して、マスター ウェイトをオプティマイザーに格納することです。トレーニング中、FP8 混合精度オプティマイザーでは、各パラメータに 6 バイトのメモリが必要です。

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

自動スケーリング

自動スケーリングは、FP8 データ形式の表現範囲内で勾配値を保存するもので、テンソル スケーリング係数の動的な調整を必要とし、これにより、全削減通信中に発生するデータ アンダーフローおよびオーバーフローの問題が軽減されます。

具体的には、研究者らはトレーニング中に変更できる自動スケーリング係数 μ を導入しました。

実験結果

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

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

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

微調整では、事前トレーニング済みの LLM を最終タスクとユーザーの好みにより適合させるために、命令の微調整に FP8 混合精度を使用し、人間によるフィードバックによる強化学習 (RLHF) を使用しました。

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

また、各コンポーネントの有効性を検証するためにアブレーション実験も実施しました。

FP8 低精度トレーニングは、将来的には大規模モデル開発のための新たな基盤となることが予想されます。

詳細については、元の論文を参照してください。