低精度学習でトランスフォーマーを最適化
NVIDIAはTransformerアーキテクチャの低精度トレーニング最適化手法を公開した。モデル大規模化に伴うGPU時間と開発コスト増大への対応を狙い、HopperおよびBlackwell GPUがサポートするFP8やNVFP4の活用を推進する。 本手法の核心は、モデル設定とバッチサイズから実動行列形状を算出することである。NVIDIA Transformer Engineを用い、BF16からNVFP4まで複数精度でベンチマークを実施し、トレーニング全体の速度向上に寄与する最適精度を事前特定する。実トレーニングに近い自動量子化モードとカーネル演算のみを測定する事前量子化モードを比較することで、量子化オーバーヘッドを精緻に評価可能だ。 CodonFM 5Bを用いたBlackwell GPU(B300)での検証では、理論値に迫る性能が確認された。事前量子化ではNVFP4がBF16に対し3.48倍のカーネルスループットを記録したが、自動量子化ではオーバーヘッドにより1.98倍となった。実環境におけるFprop速度はNVFP4がMXFP8に対し約1.47倍の改善を示した。また、FP8 DelayedScalingはBlackwellにおいて7.80ms/レイヤーを達成し、他のFP8方式を上回る競争力を示した。 重要な知見として、量子化形式における勾配計算の非対称性が挙げられる。FP8やNVFP4ではDgradがFpropに対し最大51%低速化し、行列のアスペクト比変化によるカーネル選択の違いが原因となる。さらに、Transformer Engineはサポート外のレイヤーでFP4からFP8やBF16に沈黙的にフォールバックするケースがあり、速度向上が確認できない場合はログ監視やGPUメモリ使用量の比較による実Kernelの特定が不可欠だ。 NVIDIAは低精度トレーニングが自動的にエンドツーエンドの高速化につながるとは限らないと指摘する。量子化オーバーヘッドやGEMM以外の演算コストを考慮した事前ベンチマークが、リソース効率とモデルサイズの最適化には必須となる。Transformer Engineのツールを活用し、自モデルの行列形状に合わせた精度設計を徹底するよう業界に呼びかけている。
