How to Optimize Transformer-Based Models for Low-Precision Training｜低精度トレーニングにおけるトランスフォーマーベースモデルの最適化手法 | AIニュース最前線

Transformer アーキテクチャは、多くの現代の大規模言語モデルや生成 AI モデルの基盤となっています。これらのモデルが大きくなるにつれて、トレーニングにはより多くの GPU 時間とエンジニアリングの反復作業が必要になります。したがって、Transformer の高速化は単なるパフォーマンス最適化ではなく、チームが実験を迅速に行えるかどうかや、どの程度のサイズのモデルを訓練できるかに直接影響します。[NVIDIA Hopper](https://www.nvidia.com/en-us/data-center/technologies/hopper-architecture/) および [NVIDIA Blackwell](https://www.nvidia.com/en-us/data-center/technologies/blackwell-architecture/) GPU は、FP8 や [NVFP4](https://developer.nvidia.com/blog/introducing-nvfp4-for-efficient-and-accurate-low-precision-inference/) を含む低精度演算子サポートを導入することで、この問題の解決に貢献しています。 Transformer のトレーニング時間の多くは GEMM（一般行列乗算）に費やされており、低精度フォーマットは主にこれらの行列乗算を高速化し、コストを下げることでトレーニングを加速します。ただし、Transformer の設定情報だけでは、モデル内で実際に実行されている GEMM がどれであるかはわかりません。トレーニング時間の内訳を理解したい場合は、Transformer の設定とバッチサイズを変換して、モデルが実行する正確な M×K×N 行列形状を特定し、異なる精度間でそれらの形状をベンチマークする必要があります。これにより、より高コストなトレーニングを実行する前に、アーキテクチャに最適な精度を決定することができます。 [NVIDIA Transformer Engine (TE)](https://github.com/NVIDIA/TransformerEngine) は、量子化とカーネルディスパッチを処理し、低精度フォーマットを解放します。この投稿では、高レベルのモデル設定から具体的な GEMM（行列乗算）ワークロードへ移行する方法、マイクロベンチマークでプロファイルする方法、そしてどこで低精度が実際の速度向上につながるかを見積もる方法について解説し、Transformer ベースのモデルの加速を支援します。このユースケースでは、RNA に焦点を当てた生物学向けの言語モデルである CodonFM を特徴としています。 ## モデル設定とトレーニング入力 [CodonFM](https://github.com/NVIDIA-BioNeMo/CodonFM) 5B のような 50 億パラメータのモデルを扱っていると仮定します。このモデルには以下のような構成があります： hidden_size: 4096 intermediate_size: 16384 num_attention_heads: 32 num_hidden_layers: 24 トレーニング設定は以下の通りです： micro_batch_size: 31 sequence_length: 512 ベンチマークツールはこれらのハイパーパラメータを直接受け取り、単一のコマンドで GEMM の形状を導出し、各種精度間でベンチマークを実行して、完全な速度向上分析を計算できます： python benchmark.py \ --hidden_size 4096 \ --intermediate_size 16384 \ --num_attention_heads 32 \ --num_hidden_layers 24 \ --micro_batch_size 31 \ --sequence_length 512 \ -o ./images/b300_model_config_speedup.png 注：Blackwell 固有のフラグを無効にするには、--no-fp8 --no-fp4 を追加してください。--no-fp8 --no-fp4 は、Hopper で動作する BF16 と 3 つのテンソル単位の FP8 レシピを提供します。 - --no-fp8 は MXFP8 を無効にします - --no-fp4 は NVFP4 を無効にします ## autocast モードの使用と事前量子化の比較 デフォルトでは、このツールは autocast モードで動作します。これは TE（Tensor Engine）がトレーニング中に実行する方式です：各 GEMM（General Matrix Multiply）の前に入力が動的にターゲット精度へ量子化されるため、測定された時間には量子化のコストと GEMM カーネル自体の両方が含まれます。これにより、トレーニングステップ中の現実的な 1 GEMM あたりのパフォーマンス像を提供します。 このツールは M = 31 × 512 = 15,872 トークンを計算し、すべての 12 種類の GEMM 形状を導出し、有効な精度ごとに各形状をベンチマークして完全な結果を出力します。Fprop（順伝播）、Dgrad（誤差逆伝播の勾配）、Wgrad（重み更新用勾配）の形状はすべて個別にベンチマークされ、異なる行列のアスペクト比がカーネル選択に与える影響を捉えます。 デフォルトでは、このツールは autocast モードで動作します。これは TE（Tensor Engine）がトレーニング中に実行する方式です：各 GEMM（General Matrix Multiply）の前に入力が動的にターゲット精度へ量子化されるため、測定された時間には量子化のコストと GEMM カーネル自体の両方が含まれます。これにより、トレーニングステップ中の現実的な 1 GEMM あたりのパフォーマンス像を提供します。 このツールは M = 31 × 512 = 15,872 トークンを計算し、すべての 12 種類の GEMM 形状を導出し、有効な精度ごとに各形状をベンチマークして完全な結果を出力します。Fprop（順伝播）、Dgrad（誤差逆伝播の勾配）、Wgrad（重み更新用勾配）の形状はすべて個別にベンチマークされ、異なる行列のアスペクト比がカーネル選択に与える影響を捉えます。 ![image](https://developer-blogs.nvidia.com/wp-content/uploads/2026/06/per-layer-gemm-time-nvidia-b300-autocast-mode-1.webp)*図 1. NVIDIA B300 SXM6 AC におけるオートキャストモードのレイヤー別 GEMM（行列乗算）時間。精度（BF16、FP8 Current、FP8 Delayed、MXFP8、NVFP4）およびステージ（順伝播＋逆伝播勾配計算と重み勾配計算）別に内訳した結果* 純粋な GEMM カーネルの性能を評価するためには、`--pre-quantize` オプションを追加してください。このオプションは、計測ループの実行前にすべての入力を一度前処理で量子化します。これにより、測定される時間は GEMM カーネルの実行時間のみを反映し、計測領域内での動的な量子化やブロックごとのスケーリング計算、フォーマット変換は一切含まれません。 なお、FP8 DelayedScaling は `--pre-quantize` を指定した場合でも常にオートキャストモードで実行されます。これは、動的な量子化を必要とする amax（最大絶対値）の履歴データに依存しているためです。したがって、その計測時間は前処理量子化モードにおける他の精度の時間とは直接比較できません。 python benchmark.py \ --hidden_size 4096 \ --intermediate_size 16384 \ --num_attention_heads 32 \ --num_hidden_layers 24 \ --micro_batch_size 31 \ --sequence_length 512 \ --pre-quantize \ -o ./images/b300_model_config_speedup_prequant.png ![image](https://developer-blogs.nvidia.com/wp-content/uploads/2026/06/per-layer-gemm-time-nvidia-b300-prequantized-mode.webp)*図 2. NVIDIA B300 SXM6 AC における前処理量子化モードのレイヤー別 GEMM 時間。動的な量子化オーバーヘッドを排除し、純粋なカーネルのスループットのみを抽出した結果* autocast と事前量子化による速度向上を比較することで、量子化のオーバーヘッドが実際にどの程度のコストを発生させるかが明確になります：NVFP4 から BF16 への移行は、autocast では 1.98 倍ですが、カーネルのみでは 3.48 倍に達します。この 2 つの数値間の差は、各トレーニングステップで発生する動的量子化、アダマール変換、およびブロックスケーリングによるオーバーヘッドを表しています。 実際のトレーニング速度向上を予測するには autocast の結果を使用してください。これが TE（Tensor Engine）がトレーニング中に実際に実行する方法です。一方、事前量子化の結果は、量子化のオーバーヘッドがボトルネックとなっているかどうかを理解するため、あるいは量子化の実装に依存しない精度間の生テンソルコアスループットを比較するために使用します。 ## 実モデルにおける結果の解釈 このセクションでは、実モデルに対してこれらの結果をどのように解釈するかを解説します。同じ CodonFM 5B の設定を用いて、NVIDIA B300 上でフルモデル構成ベンチマークを実行しました。Fprop（順伝播）の結果から得られる形状ごとの NVFP4 と MXFP8 の速度向上比は以下の通りです： QKV proj: 0.579 / 0.392 = 1.48x Attn out: 0.269 / 0.256 = 1.05x （わずかに高速 — オーバーヘッドが GEMM の利得とほぼ同程度） MLP up: 0.924 / 0.635 = 1.46x MLP down: 1.076 / 0.649 = 1.66x 以下の点に注意してください： - アテンション出力の GEMM は、低精度化による恩恵をほとんど受けません。MXFP8 ベースラインと比較すると、速度向上はわずか 1.05 倍です。これはレイヤー内でも最も小さな重み行列（4096×4096）であり、低精度がオーバーヘッドを上回るには barely 十分な大きさではありません。一方、はるかに大きな MLP Down GEMM では、同じハードウェア上で MXFP8 に対して NVFP4 が 1.66 倍の速度向上を実現しています。MLP Down GEMM は量子化のオーバーヘッドを相殺するのに十分な大きさですが、アテンション出力ではそれができません。 - 大規模な GEMM では、理論値には届かないが実用的な性能向上が見られます。FP4 テンサーコアは、大規模な GEMM において MXFP8 よりも 1.46 倍から 1.66 倍の速度を示します。これはハードウェア仕様で示される理論的な 2 倍から 3 倍という値には遠く及びません。アテンション出力の GEMM も含めると、混合された順伝播（Fprop）の速度向上は 1.47 倍に低下します。Wgrad の計算時間、非 GEMM オーバーヘッド、および NVFP4 に固有な量子化コストを加味すると、トレーニングにおける NVFP4 と MXFP8 のエンドツーエンドでの差は、これらのカーネルレベルの数値と整合しています。 - FP8 DelayedScaling は、NVIDIA Blackwell において驚くほど競争力があります。autocast モードではレイヤーあたり 7.80 ms で、FP8 CurrentScaling（9.15 ms）および MXFP8（8.98 ms）の両方を上回ります。一方、事前量子化モードでは FP8 CurrentScaling が先行し（6.81 ms 対 8.12 ms）、DelayedScaling の amax-ヒストリアプローチは量子化オーバーヘッドが小さいものの、生カーネルのスループットは同程度であることを示唆しています。これは、量子化税を含めるか含めないかで勝者が異なることを浮き彫りにする、autocast と事前量子化の比較における良い例です。 - 事前量子化の結果は、真のカーネルポテンシャルを明らかにします。--pre-quantize オプションを実行すると量子化オーバーヘッドが完全に除去され、NVFP4 と BF16 の比較では、autocast モードでの 1.98 倍から、カーネルのみでの 3.48 倍へと跳ね上がります。これは FP4 テンサーコアが実際の速度向上を実現していることを示しています。autocast モードにおける量子化オーバーヘッドこそが、この差を縮めている要因です。 - Fprop（順伝播）と Dgrad（逆伝播勾配計算）の比較から、2 倍近似は量子化フォーマットにおいては不正確であることがわかります。BF16 の Dgrad は Fprop と比べて 2% 以内ですが、量子化フォーマットでは Dgrad の合計が 5–13% 遅くなります。特に QKV プロジェクションの Dgrad は著しく非対称で、FP8/FP4 では Fprop よりも 33–51% 遅くなります。これは K（4096）と N（12288）を交換することで行列のアスペクト比が劇的に変化し、カーネルの選択が変わるためです。まさにこの理由から、ツールは Fprop の時間を単に 2 倍カウントするのではなく、Fprop と Dgrad を別々にベンチマークします。 推定された GEMM（行列乗算）のみでの速度向上率が得られたら、それを観測されたエンドツーエンドのトレーニング速度向上率と比較してください： - GEMM の速度向上 ≈ トレーニングの速度向上：GEMM がステップを支配しており、すべてが期待通り動作している - GEMM の速度向上 >> トレーニングの速度向上：GEMM 以外のオーバーヘッドが利益を食いつぶしています。特に NVFP4 の場合、このオーバーヘッドには Wgrad 入力に対するランダム・ハダマード変換、勾配に対する確率的丸め、重みに対する 2D ブロックスケーリング、およびテンソルごとの amax（最大絶対値）計算のための追加メモリパスが含まれます。これらはすべて MXFP8 では不要な追加演算であり、生身の FP4 GEMM がはるかに高速であっても、これらのオーバーヘッドが差を大幅に縮める可能性があります - GEMM の高速化率はマイクロベンチマークでも約 1.0 です。FP4 カーネルはこれらの形状において実際には高速ではなく、あるいは FP8 に静かにフォールバックしている可能性があります。 最後のケースは特に確認する価値があります。NVTE_LOG_LEVEL=1 を設定するか、[NVIDIA Nsight Systems](https://developer.nvidia.com/nsight-systems) で検査して、TE（Transformer Engine）が実際に FP4 カーネルをディスパッチしているかを確認してください。TE はまだ FP4 をサポートしていないレイヤーや演算に対して、FP8 や BF16 に静かにフォールバックすることがあり、これが他の症状なしに同等のパフォーマンスを示す理由となります。また、MXFP8 と NVFP4 の実行時の GPU メモリ使用量を比較することもできます。メモリ使用量がほぼ同一である場合、それは FP4 重みが実際に保存されていないという強力なシグナルです。 ## モデルの低精度トレーニング向けベンチマークを開始する 低精度トレーニングによる高速化は、モデルが実行する実際の GEMM（行列乗算）形状に大きく依存しており、低精度で実行することが自動的にエンドツーエンドのトレーニング性能向上につながるわけではありません。特に量子化オーバーヘッド、カーネル選択、および非 GEMM 演算が含まれる場合です。Transformer 設定を具体的な M×K×N のワークロードに変換することで、フルトレーニングを実行する前に、モデルにとって重要な形状上で BF16、MXFP8、NVFP4 をベンチマークすることができます。 GEMM をベンチマークして、ご自身にとって最適な精度を確認してください。まずは [ベンチマークスクリプト](https://github.com/NVIDIA/TransformerEngine/blob/main/benchmarks/gemm/benchmark_gemm.py) をご覧ください。完全なドキュメントやこれらの形状がどのように導き出されるかについては、Transformer Engine のドキュメントにある [GEMM プロファイリングチュートリアル](https://docs.nvidia.com/deeplearning/transformer-engine/user-guide/examples/gemm_profiling/gemm_profiling.html) を参照してください。 このベンチマークは以下の目的で使用できます: - 自動キャストの結果を確認し、現実的なトレーニングの高速化効果を予測する - 事前量子化の結果を確認し、ボトルネックがカーネルにあるのか量子化にあるのかを把握する - トレーニング実行に着手する前に候補となるモデル構成をツールで検証する。このツールはアーキテクチャ共設計のための有用な手段です

低精度トレーニングにおけるトランスフォーマーベースモデルの最適化手法

背景や根拠まで確認しますか？

関連記事

調べる

選ぶ

サイト