NVIDIA NeMo AutoModelによるTransformerファインチューニングの高速化

NVIDIA NeMo AutoModelは、大規模なカスタム生成AIモデルを構築するためのNVIDIA NeMoフレームワークの一部であるオープンライブラリです。これはHuggingFace Transformers v5の上に構築されており、エキスパート並列化（Expert Parallelism）、DeepEP融合オールツーオールディスパッチ、TransformerEngineカーネルを追加し、v5の動的ウェイトローディングを活用して、これらの最適化を広範かつ成長し続けるモデルファミリーに適用します。その成果は、MoEモデルのファインチューニングにおいて、ネイティブのTransformers v5と比較してトレーニングスループットが3.4～3.7倍、GPUメモリ使用量が29～32%削減されることです。しかも、同じfrom_pretrained() APIを使用し、たった1行のインポートコードを変更するだけで、他のコード変更は一切不要です。

MoEモデルの台頭により、効率的なトレーニングには新たな課題が生じています。数百のエキスパートにトークンをルーティングする、エキスパートの行列積を単一のカーネルに融合する、GPU間でウェイトをシャーディングする、通信と計算をオーバーラップするなど、これらすべてに汎用ライブラリが標準で提供する以上のインフラストラクチャが必要です。Transformers v5は、エキスパートバックエンド、動的ウェイトローディング、分散実行のためのテンソル並列計画など、MoEのファーストクラスサポートを導入しました。さらに、v5はPyTorchのDeviceMeshをfrom_pretrained()に直接統合することで、分散トレーニングをファーストクラスにしました。

NeMo AutoModelは、AutoModelForCausalLMをサブクラス化し、エキスパート並列化、DeepEP融合オールツーオールディスパッチ、TransformerEngineカーネルを追加することで、v5の上に構築されています。DeepEPはv5にはまだない部分で、通信をエキスパート計算とオーバーラップさせます。また、NeMo AutoModelはv5の可逆ウェイト変換を利用して各モデルをロードするため、モデルごとのチェックポイント配管ではなく、これらの再利用可能なコア演算に工数を集中でき、save_pretrained()は依然としてvLLMやSGLangがロードできる標準的なHFチェックポイントを出力します。

性能評価では、NeMo AutoModelの利点を2つのシナリオで示しました。16ノードにわたる550Bモデルのフルファインチューニングと、単一ノードでの2つの30B MoEモデルのトレーニングです。550Bモデルでは、Transformers v5はメモリ不足で実行できませんでしたが、NeMo AutoModelはエキスパート並列化（EP=64）によりエキスパートウェイトをGPU間でシャーディングし、フルファインチューニングを可能にし、GPUあたり毎秒815トークン、ピークメモリ58.2 GiBを達成しました。単一ノード8×H100 GPUでは、Qwen3-30B-A3Bモデルはv5の3,075 TPS/GPUから11,340 TPS/GPU（3.69倍）に向上し、ピークメモリは68.2 GiBから48.1 GiB（-29%）に減少しました。Nemotron 3 Nano 30B A3Bモデルは4,583 TPS/GPUから15,421 TPS/GPU（3.36倍）に向上し、ピークメモリは62.1 GiBから42.5 GiB（-32%）に減少しました。

速度向上の源泉は3つあります。エキスパート並列化によるメモリ圧力の軽減（EP=8でGPUあたりのMoEメモリフットプリントを8分の1に削減）、DeepEPによる通信と計算の融合（個別の非ブロッキング通信を回避）、TransformerEngineカーネルによる融合アテンション、線形層、RMSNormなどのコア演算の高速化です。NeMo AutoModelはBackendConfigを通じてアテンション、線形層、エキスパート、ディスパッチャーの実装を選択可能にしています。要するに、NeMo AutoModelはHuggingFace TransformersとのAPI互換性を維持しながら、MoEモデルのファインチューニングに顕著な性能向上をもたらします。