NVIDIA、MoE学習を最大3.7倍高速化

発表の要点

import1行で3.4〜3.7倍高速化
GPUメモリ最大32%削減
Transformers v5を土台に拡張
HF互換APIで既存コード不変

技術と適用範囲

Expert Parallelismで専門家を分散
DeepEPが通信と計算を融合
550Bモデルの全層調整も実現
詳細を読む

NVIDIAは6月24日、HuggingFace Transformersの上に構築するオープンライブラリ「NeMo AutoModel」を公開しました。import文を1行変えるだけで、MoE(混合専門家)モデルのファインチューニングTransformers v5比で3.4〜3.7倍高速化し、GPUメモリを29〜32%削減します。from_pretrained()など既存APIはそのまま使え、コード改変は不要です。

MoEモデルの学習には固有の難しさがあります。数百の専門家へトークンを振り分け、行列積を一つのカーネルに融合し、重みをGPU間で分割し、通信と計算を重ね合わせる処理が必要だからです。Transformers v5は専門家バックエンドや動的な重み読み込みでこれに対応しましたが、通信と計算を重ねるDeepEPは未実装でした。

NeMo AutoModelはこの欠けた部分を補います。AutoModelForCausalLMを継承し、Expert Parallelism(EP)、DeepEPによる全対全ディスパッチTransformerEngineカーネルを追加しました。EPは専門家の重みをGPU間で物理的に分割し、8GPUなら各GPU専門家の8分の1だけを保持します。これにより、従来は約55GiB必要だった専門家の重みが1GPUあたり約6.8GiBに収まります。

性能評価は2つの規模で実施されました。8GPU単一ノードのQwen3-30B-A3Bでは、v5比でスループットが3.69倍、ピークメモリは29%減。Nemotron 3 Nano 30Bでも3.36倍、メモリ32%減を記録しました。高速化の源はEPによるメモリ削減、DeepEPの通信融合、TransformerEngineの最適化カーネルの3点です。

大規模側では、550BパラメータのNemotron 3 Ultraの全層ファインチューニング16ノード128GPUで実行しました。Transformers v5はこの規模でメモリ不足になり動作しませんが、EPが専門家を分散することで学習が可能になります。EPが本領を発揮するのは、まさにこの大規模領域です。

NeMo AutoModelの出力は標準的なHF形式のsafetensorsであるため、save_pretrained()で保存した重みはvLLMやSGLangといった推論基盤にそのまま載せられます。NVIDIAは、Transformers v5を使うユーザーにとって本ライブラリが摩擦のない次の一歩になると位置づけています。