NVIDIAがMoE訓練のベンチマークを公開、B200 8基の最適構成で基準の2.21倍

NVIDIAは9月24日、BioNeMoによる混合エキスパート（MoE）の訓練手順とベンチマーク結果を公開した。B200 8基でMixtral-8x7Bを実行した場合、最適な設定でGPU当たり毎秒9,050トークンに達し、記事内のHugging Faceベースラインの約2.21倍となった。これは特定の訓練構成について同社が公表したスループットの結果であり、生物学分野のタスクで目標品質に達するまでの時間を同じ比率で短縮できることは、まだ示されていない。[公式ベンチマーク](https://developer.nvidia.com/blog/efficient-moe-training-for-biological-foundation-models/)

比較表からは、性能向上の要因を区別する必要があることも分かる。ベースラインはBF16と8基のGPUによる完全シャード化データ並列を採用し、スループットは4,096だった。精度と並列化設定を維持したままTransformer Engineに切り替えると4,447となり、約8.6%向上した。9,050を達成した設定では、エキスパート並列とMXFP8へ同時に変更している。このため、2.21倍という数値は構成全体の違いを反映しており、カーネル融合だけの効果とは言えない。[比較条件](https://developer.nvidia.com/blog/efficient-moe-training-for-biological-foundation-models/)

MoEはトークンごとに一部のエキスパートだけを選択するが、小規模なカーネルの起動やGPU間のデータ交換によって訓練が遅くなる可能性がある。BioNeMoのモデル実装では、グループ化した線形演算でエキスパートを処理し、エキスパートの重みをエキスパート次元に沿って分割して複数のGPUに配置する。データ並列とエキスパート並列を組み合わせた2次元構成も可能だ。デフォルトのディスパッチャーはNCCLを介してトークンを交換し、より高度な処理経路にはDeepEPと、GPU間のNVLink経由のアクセスが別途必要になる。[モデルのドキュメント](https://docs.nvidia.com/bionemo-recipes/latest/main/recipes/models/mixtral/)

MXFP8は32個の数値ごとに1つのスケーリング係数を共有し、8ビット表現のダイナミックレンジを改善する。数値には4ビットの指数部と3ビットの仮数部を使い、スケーリング係数は2のべき乗を8ビットで格納する。転置はビットの並べ替えだけでは済まない。方向が異なるブロックでは再量子化が必要になるため、Transformer Engineは元の高精度入力から2つの方向のコピーを作成する。この点からも、8ビット演算を使えば訓練全体のメモリ使用量がそのまま半減するわけではないことが分かる。[精度と転置の説明](https://nvidia.github.io/TransformerEngine/examples/fp8_primer.html)

モデルのドキュメントでは、層ごとに精度を選択できるほか、重みの形式変換と出力比較の手順も提供されている。研究者はまず変換前後の結果を確認し、その後、低精度演算とGPU間の機能を段階的に有効にすることで、数値のずれの原因を調べる範囲を絞り込める。[変換と検証](https://docs.nvidia.com/bionemo-recipes/latest/main/recipes/models/mixtral/)

公開リポジトリでは、これらの手順を変更可能なリファレンス実装と位置付けている。ハードウェアの対応範囲にも制限があり、MXFP8の対応表にはCompute Capability 10.0と10.3が記載され、12.0への対応はまだ追加されていない。開発チームが再検証する際は、コミット、パッケージ、シーケンス長、バッチ設定、ルーティング設定を固定し、損失曲線、検証時の品質、通信が占める割合を追跡する必要がある。スループットと目標品質への到達に必要なステップ数を併せて測定して初めて、低精度演算とエキスパートのシャード化が訓練全体のコストを削減するかどうかを判断できる。[リポジトリと対応表](https://github.com/NVIDIA-BioNeMo/bionemo-recipes)