NVIDIA Blackwell で NVFP4 を使用し、JAX と MaxText でモデルの学習を高速化

フロンティア LLM の事前学習は、スループットに帰着します。数千のアクセラレータにわたって兆単位のトークンをトレーニングする際、ステップ時間のわずかなパーセントポイントの差が、数日間のトレーニング時間と膨大な計算コストの増大につながります。数値精度は利用可能な最も効果的な調整項目の一つですが、低ビット混合精度事前学習を正しく実装するのは困難です。

これに対処するため、TransformerEngine の NVFP4 トレーニングレシピでは、JAX 事前学習にサブバイト精度（subbyte precision）を採用しています。エンドツーエンドの例については、高性能でスケーラブルな LLM フレームワークライブラリである MaxText のレシピをご覧ください。その結果、NVIDIA Blackwell 上で FP8 ベースラインと比較して測定可能な精度損失なしに、高スループットかつ 4 ビットの混合精度事前学習を実現しています。

本記事では、NVFP4 形式とその超低精度において高性能と高精度を達成するためにどのように設計されているかを解説します。また、MaxText の NVFP4 事前学習レシピの適用方法と、パフォーマンス向上を示す性能データの収集方法についても紹介します。手法の詳細については、NVFP4 事前学習論文 を参照してください。

NVFP4 形式とその利点

この NVFP4 入門記事 は、そのフォーマットと、2 レベルのマイクロスケーリングが他のマイクロスケーリング形式よりも少ない誤差で高い信号を符号化する方法について説明しています。また、NVIDIA GB300 Grace Blackwell Ultra Superchip における NVFP4 のネイティブハードウェアサポートが、NVIDIA Hopper 上のネイティブ FP8 精度と比較して GEMM（行列乗算）スループットを 7 倍に向上させる方法についても解説しています。この高いスループットと NVFP4 プリトレーニングレシピにより、精度の低下を無視できるレベルでトレーニングステップ時間を短縮できます。これにより、AI ファクトリーは同じ時間予算内でより多くの、より大規模なモデルをトレーニングしたり、より短い時間予算でモデルを高速にトレーニングしたりすることが可能になります。

NVFP4 プリトレーニングレシピ

NVFP4 レシピは、収束性を維持しつつ NVIDIA Blackwell および NVIDIA Rubin プラットフォームの NVFP4 スループットを引き出すために組み合わされた複数の要素から成り立っています。効率的な狭帯域幅精度でのトレーニングを可能にするため、このプリトレーニングレシピでは、パフォーマンスと精度に基づいて選択されたいくつかの主要な技術を採用しています。

4 ビットプリトレーニングに必要な精度を維持しながら、5 つの主要な要素が連携して機能します:

• マイクロブロックスケーリングは 16 要素ブロックを使用し、MXFP4 の 32 要素ブロックの半分サイズです。これにより、単一の外れ値が共有スケールに与える影響が小さくなります。

• E4M3 ブロックスケールファクターは、MXFP4 の累乗 2 進数である E8M0 スケーリングではなく、マントissa ビットを使用し、その下にテンソルごとの FP32 スケールが層状に配置されています。8B パラメータ、1T トークンの実験において、MXFP4 は NVFP4 の最終損失と同等の性能を出すために約 36% 多くのトークンを必要とします。

• ランダム・ハドマード変換は、外れ値をガウス分布化するために WGRAD GEMM の入力に対してのみ適用されます。FPROP と DGRAD のパスでは変換が不要です。なぜなら、これらのパスを変換すると重みも変換する必要が生じ、2D スケールの一貫性が損なわれるからです。

• 2D 重みスケーリングは、16×16 の重みブロックごとに FP8 スケールを 1 つ使用するため、FPROP とその転置である DGRAD は同じスケールを使用します。一方、活性化関数と勾配は、オーバーヘッドの低い 1×16 スケーリングを維持します。

• 確率的丸め（Stochastic Rounding）は、バイアスのない丸めを用いて微小な更新値がゼロに潰れてしまうのを防ぎます。重みと活性化関数は「偶数への丸め」方式を維持しており、ここで確率的丸めを使用すると誤差が増幅されてしまいます。両方のモードは Blackwell の FP4 変換命令でネイティブにサポートされています。

図 1 は、1 つの線形層内部における NVFP4 のデータフローを示しています。

3 つの GEMM（FPROP：順伝播、DGRAD：活性化関数の勾配、WGRAD：重みの勾配）は、Transformer の MLP（フィードフォワード）層に対してのみ NVFP4 に量子化されます。アテンションブロック内部の GEMM（QKV 投影、アテンション出力投影、およびスコア/コンテキスト行列積）は、より高い精度を維持します。

NVFP4 はまず MLP レイヤーに適用されます。これは、アテンションのソフトマックス関数が QK^T スコアにおける量子化ノイズを指数関数的に増幅するためです。また、アテンション活性化値には 4 ビット精度では適切に表現できない集中した外れ値も含まれています。MLP はトレーニング FLOPs の大部分を占めるため、このアプローチにより収束のリスクを冒すことなく、速度向上の大半を実現できます。

3 つの MLP GEMM はすべて NVFP4 入力を消費し、BF16 出力を生成します。これらは最終的にオプティマイザステップで FP32 のマスターウェイトに統合されます。同じパスにより、レシピが収束性を維持するための選択が明確になります：重みに対する 2D ブロック量子化（転置Acrossして FPROP/DGRAD 値の一貫性を保つ）、WGRAD 入力に対するランダム・アダマール変換（4 ビット量子化の前に外れ値を平坦化する）、および勾配量子化器における確率的丸め（小さな更新値の不偏性を維持する）。

MaxText での NVFP4 の有効化

MaxText NVFP4 レシピ は、JAX-Toolbox GitHub リポジトリで利用可能です。この起動スクリプトは、Blackwell 上で Llama 3 8B を NVFP4 でトレーニングします。有効化するには、MaxText の量子化フラグを設定して NVFP4 パスに切り替えます。2 つのモードが用意されています：

• quantization=te_nvfp4: ランダムアダマール変換を伴う NVFP4。te_nvfp4_no_rht での収束が満足できない場合に推奨されます。

• quantization=te_nvfp4_no_rht: RHT を伴わない NVFP4。オーバーヘッドは最小限ですが、収束の質が低下する可能性があります。

JAX、NVIDIA Transformer Engine、および必要な NVIDIA CUDA/cuDNN ライブラリがインストールされたコンテナ内で、MaxText リポジトリのルートから例のスクリプトを実行してください。パブリックな NVIDIA MaxText コンテナ ghcr.io/nvidia/jax:maxtext の使用を推奨します。

以下は、Transformer Engine を介して nvfp4 引数を宣言する Llama3 8B MaxText NVFP4 学習スクリプトの部分的な例です：

RUN_SETTINGS="-m maxtext.trainers.pre_train.train maxtext/configs/base.yml run_name=debug_run base_output_directory=./debug_logs hardware=gpu dataset_type=synthetic model_name=llama3-8b remat_policy='minimal_with_context_and_quantization' scan_layers=False attention='cudnn_flash_te' steps=50 dtype=bfloat16 max_target_length=8192 per_device_batch_size=4 ici_data_parallelism=${ici_DP} dcn_data_parallelism=${dcn_DP} ici_fsdp_parallelism=${ici_FSDP} dcn_fsdp_parallelism=${dcn_FSDP} profiler=nsys enable_checkpointing=false override_model_config=True gradient_accumulation_steps=1 quantization=te_nvfp4_no_rht max_segments_per_seq=32"

起動後、MaxText はステップ時間、デバイスあたりの TFLOP/s、およびデバイスあたりのトークン数を出力します。検証用の NVIDIA Nsight Systems トレースは、base_output_directory に書き込まれます。

以下に示す比較で用いた FP8 ベースラインを生成するには、quantization=te_fp8_delayedscaling を指定して同じスクリプトを実行してください。

パフォーマンス結果

本ベンチマークでは、MaxText による Llama 3 8B の事前トレーニングを FSDP=4 で実行し、シーケンス長は 8,192、デバイスあたりのバッチサイズ（MBS: per-device batch size）は 4 と設定しました。また、評価は public ghcr.io/nvidia/jax:maxtext コンテナ内で 50 ステップ分行われました。

表 1 は、Llama 3 8B および Llama 3.1 405B モデルにおける NVIDIA GB200 Grace Blackwell Superchip と NVIDIA GB300 Grace Blackwell Ultra Superchip 上での MaxText 事前トレーニング性能を要約したものです。これは、同じハードウェア、並列化設定、およびグローバルバッチサイズ（GBS: global batch size）条件下で NVFP4 レシピと FP8 ベースラインを比較した結果です。数値はシーケンス長 8,192 で測定されたものです。

モデル**ハードウェアGPU 数FSDPMBSGBS**Seq len

Llama3 8BGB200444168,192

Llama3 8BGB300444168,192

Llama 3.1 405BGB20012812811288,192

Llama 3.1 405BGB30012812811288,192

*表 1. GB200 および GB300 における Llama3 8B と Llama3.1 405B モデルの NVFP4 vs FP8 事前トレーニングレシピ構成*

モデル**ハードウェアGPU あたりの FP8 TFLOPsGPU あたりの NVFP4 TFLOPs**FP8 に対する速度向上比

Llama 3 8BGB200149720171.35×

Llama 3 8BGB300175923011.31×

Llama 3.1 405BGB200155722411.44×

Llama 3.1 405BGB300210336331.73×

*表2. GB200 および GB300 で測定された、Llama3 8B および Llama3.1 405B モデルにおける NVFP4 と FP8 の事前学習パフォーマンスベースライン構成の比較*

図2は、4 つのベースライン構成全体にわたる GPU 単位の持続 TFLOP/s を示しています。NVFP4 はすべての構成において、GPU あたり追加で 500〜700 TF/s を達成します。FP8 ベースラインに対する 1.31〜1.73 倍の高速化は、モデル、ハイパーパラメータ、並列処理、およびグローバルバッチサイズを同一に保ちつつ、GEMM（行列乗算）の精度を変更した結果によるものです。

最も大きな相対的な向上が見られるのは 405B 構成です（GB200 で 1.44 倍、GB300 で 1.73 倍）。ここではステップあたりの GEMM の負荷が FSDP（Fully Sharded Data Parallel）の集合通信オーバーヘッドを支配しており、精度レベルでの高速化がそのまま壁時計時間の短縮に直結します。

図3は、他のハイパーパラメータを同一に保ちつつ、10,000 ステップの事前学習全体にわたる FP8 ベースラインと NVFP4 の Llama 3 8B のトレーニング損失を重ねて示しています。両方の実行は、約 12.2 nats から約 3.9 nats へと同じ曲線を描いて低下し、収束領域における平均ギャップはわずか +0.026 nats で、これはステップ間のノイズの範囲内に収まっています。図2で示される NVFP4 の高速化には、測定可能な精度の低下を伴いません。

始め方

MaxText コンテナ をプルし、Blackwell 上で nvfp4_example.sh を実行して開始してください。

謝辞

*JAX、XLA、TE における NVFP4 の有効化への貢献に対し、Jaroslav Sevcik、Ilia Sergachev、Johannes Reifferscheid、Phuong Nguyen、Jeremy Berchtold の皆様に特別なお礼を申し上げます。*