FP8 は、独自の数値表現により、一定レベルの精度を維持しながら大規模モデルのトレーニング速度を向上させ、メモリ使用量を節約し、最終的にトレーニング コストを削減できます。

AI大規模モデル開発システムであるColossal-AIは、混合精度トレーニングをアップグレードし、主流のBF16(O2) + FP8(O1)の次世代混合精度トレーニング方式をサポートしました。

わずか 1 行のコードで、主流の LLM モデルの平均 30% の高速化を実現し、大規模モデルの開発コストを削減し、トレーニングの収束を確実にすることができます。

追加の手書き CUDA 演算子を導入する必要がないため、長い AOT コンパイル時間と複雑なコンパイル環境構成を回避できます。

オープンソース アドレス: https://github.com/hpcaitech/...

FP8 混合精度トレーニング

低精度コンピューティングは、GPU ハードウェア開発において常にトレンドとなっています。

初期の FP32 から現在一般的な FP16/BF16、そして Hopper シリーズ チップ (H100、H200、H800 など) がサポートする FP8 に至るまで、低精度の計算速度はますます高速化し、必要なメモリはますます少なくなっており、これは大型モデルの時代のハードウェア要件と非常に一致しています。

現在、FP8 混合精度トレーニングのトレーニング結果に影響を与える最大の要因はスケーリング方法であり、一般的な方法は 2 つあります。

• 遅延スケーリング
• リアルタイムスケーリング

遅延スケーリングは、スケーリングの更新を行列乗算（GEMM）と統合しながら、前の時間ウィンドウ内のスケーリング値を用いて現在のスケーリングを推定します。この計算手法は非常に効率的ですが、スケーリングを推定するため、収束性に大きな影響を与えます。

リアルタイムスケーリングでは、現在のテンソル値を直接使用してスケーリングを計算するため、計算効率は低下しますが、収束への影響は小さくなります。NVIDIAのレポートによると、これら2つのスケーリング手法の計算効率の差は10%以内です。

Colossal-AI は、他の遅延スケーリング実装に匹敵するパフォーマンスを実現しながら、トレーニングの収束への影響を最小限に抑えるリアルタイム スケーリング スキームを採用しています。

単一のH100カードを用いた行列乗算のテストでは、行列の次元が大きいほどFP8の高速化効果が顕著であることが示されました。さらに、図1に示すように、Colossal-AI実装のパフォーマンスはTransformer Engineとほぼ同じです。ただし、Transformer Engineは複雑なAOTコンパイル環境の設定と長いコンパイル時間を必要とします。

△図1. シングルカードGEMMパフォーマンステスト

実験結果をより現実的なものにするために、Colossal-AI は実際のトレーニングのために主流の LLM で直接テストされました。

テストはまず単一のH100カードで実施されました。その後のテストでは、Transformer Engine (TE) はデフォルトのレイテンシスケーリングスキームを使用しました。

収束テストも実施され、図4に示すように、FP8混合精度トレーニングの損失曲線は基本的にbf16と同じであることがわかります。

△図4. H100シングルカードLLaMA2-7Bの混合精度トレーニング損失曲線

Colossal-AIは、H800上でマルチGPU並列学習シナリオにおけるパフォーマンスもテストしました。LLaMA2-7Bを8GPU搭載のH800マシン1台で学習させたところ、Colossal-AI FP8はColossal-AI BF16と比較して35%、Torch FSDP BF16と比較して94%のスループット向上を示しました。

単一の 8 カード H800 マシンで LLaMA2-13B をトレーニングした場合、Colossal-AI FP8 は Colossal-AI BF16 と比較して 39% のスループット向上を示しました。

Cohere Command-R 35B を 2 台のマシン、16 枚のカードの H800 システムでトレーニングすると、図 7 に示すように、Colossal-AI FP8 は Colossal-AI BF16 と比較して 10% のスループット向上を示します。

NVIDIA のレポートとテスト経験に基づいて、FP8 混合精度トレーニング パフォーマンスの最適化に関する予備的な洞察が得られました。

• テンソル並列処理の使用を最小限に抑え、パイプライン並列処理に置き換えます。
• モデルの隠れたサイズが大きいほど、加速効果は大きくなります。
• 行列乗算の割合が高いモデルでは、高速化の効果が大きくなります。

上記の実験では Command-R 35B がテンソル並列処理を使用したため、高速化の効果はそれほど大きくありませんでした。

Colossal-AIはFP8を幅広くサポートしており、様々な並列手法がFP8混合精度学習と互換性があります。プラグインの初期化時にFP8を有効にするだけで使用できます。

 from colossalai.booster.plugin import GeminiPlugin, HybridParallelPlugin, LowLevelZeroPlugin ... plugin = LowLevelZeroPlugin(..., use_fp8=True) plugin = GeminiPlugin(..., use_fp8=True) plugin = HybridParallelPlugin(..., use_fp8=True)

さらに、追加のコードや AOT コンパイルは必要ありません。

オープンソース アドレス: https://github.com/hpcaitech/...