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

Back to Articles

HuggingFace TransformersはオープンソースAIエコシステムの基盤となっており、最近の Transformers v5 リリースでは、フロンティアモデルの主要アーキテクチャであるMixture-of-Experts (MoE) モデルへのファーストクラスサポートが強化されました。v5はMoEの基盤を提供します：エキスパートバックエンド、動的重みロード、および分散実行により、MoEを拡張可能で構築しやすくします。

NVIDIA NeMo AutoModelは、NVIDIA NeMoフレームワークの一部であるオープンライブラリで、大規模なカスタム生成AIモデルを構築するためのものです。NeMo AutoModelはv5の上に構築され、Expert Parallelism、DeepEP融合オールトゥオールディスパッチ、TransformerEngineカーネルを追加し、v5の動的重みロードを活用して、これらの最適化を幅広いモデルファミリーに提供します。その結果、3.4〜3.7倍のトレーニングスループット向上と29〜32%のGPUメモリ削減を実現し、ネイティブのTransformers v5と比較して、from_pretrained() APIをそのまま使用できます：1行のインポート変更だけで、他のコード変更は不要です。

このブログでは、この組み合わせの仕組みと、ユーザーがAPIを変更せずにMoEモデルを高速にファインチューニングする方法について詳しく説明します。

Background

MoEモデルの台頭により、効率的なトレーニングに新たな課題が生まれました：数百のエキスパート間でトークンをルーティングし、エキスパートmatmulを単一のカーネルに融合し、重みをGPU間でシャーディングし、通信と計算をオーバーラップさせるには、汎用ライブラリが標準で提供するインフラストラクチャを超えるものが必要です。

Transformers v5（「v5」）は、エキスパートバックエンド、動的重みロード、および分散実行のためのテンソル並列プランなど、ファーストクラスのMoEサポートを導入しました。さらに、v5はPyTorchのDeviceMeshをfrom_pretrained()に直接統合することで、分散トレーニングをファーストクラスにしました。

NeMo AutoModelはv5の上に構築され、AutoModelForCausalLMをサブクラス化し、Expert Parallelism (EP)、DeepEP融合オールトゥオールディスパッチ、TransformerEngineカーネルを追加します。DeepEPはv5がまだ持っていない部分です：通信とエキスパート計算をオーバーラップさせます。NeMo AutoModelはv5のリバーシブル重み変換を利用して各モデルをロードするため、モデルごとのチェックポイント処理ではなく、これらの再利用可能なコア演算にエンジニアリングを集中でき、save_pretrained()は依然としてvLLMやSGLangなどのツールがロードできる標準のHFチェックポイントを出力します。

次のセクションでは、2つの仕組みと測定されたパフォーマンス向上について説明します。16ノードにわたるフルファインチューニングのNVIDIA Nemotron 3 Ultra 550B A55Bから、Qwen3-30B-A3BやNemotron 3 Nano 30B A3Bなどのシングルノードモデルまでをカバーします。

NeMo AutoModel: Same API, More Performance

NeMo AutoModelの目標の1つは、HuggingFace TransformersとのAPI互換性によりオープンソースコミュニティを可能にすることです。NeMoAutoModelForCausalLMはAutoModelForCausalLMをサブクラス化しているため、HFモデルで動作するコードはAutoModelでも動作します。

以下は、両方でモデルをロードする方法です。インポートのみが変更されます：

この1行のインポートで多くの作業が行われます。Qwen3、NVIDIA Nemotron、GPT-OSS、DeepSeek V3などの人気のMoEアーキテクチャの場合、NeMo AutoModelはTransformerEngineアテンション、融合線形レイヤー、カスタムエキスパートカーネルを備えた手動チューニングされた実装を提供します。その他のモデルについては、バニラHFにフォールバックしつつ、Liger kernelパッチなどの最適化を適用します。どちらのパスでも、結果のモデルはスケーリング可能です：device_meshを渡すことで、さらなる書き換えなしにマルチGPUトレーニングが可能になります。

NeMo AutoModelが真に輝くのは、MoEモデルをマルチGPUトレーニングにスケーリングする場合です。Nemotron 3 Nano 30B A3Bを8 GPUにわたるExpert Parallelismでトレーニングするには、分散メッシュ設定を追加します：

import os
import torch
import torch.distributed as dist
from nemo_automodel import NeMoAutoModelForCausalLM
from nemo_automodel.recipes._dist_utils import create_distributed_setup_from_config

dist.init_process_group(backend="nccl")
torch.manual_seed(0)
torch.cuda.set_device(int(os.environ.get("LOCAL_RANK", 0)))

dist_setup = create_distributed_setup_from_config(
    {
        "strategy": "fsdp2",
        "ep_size": 8,
    },
)

model = NeMoAutoModelForCausalLM.from_pretrained(
    "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16",
    dtype=torch.bfloat16,
    distributed_setup=dist_setup,
)

dist.destroy_process_group()

これにより、from_pretrained()呼び出しだけで、FSDP2、Expert Parallelism、TransformerEngineカーネル、DeepEPディスパッチによる速度、スケーラビリティ、メモリ最適化が得られます。

Performance Comparison

NeMo AutoModelを2つのレジームで評価しました：16ノードにわたるフロンティア規模の550Bモデルのフルファインチューニングと、シングルノードでの2つの30B MoEモデルのトレーニングです。550Bの結果は大規模でのExpert Parallelismの必要性を示し、30Bの結果はTransformers v5に対するGPUあたりの高速化を定量化します。

Nemotron 3 Ultra 550B A55B (full fine-tune, multi-node)

Nemotron 3 Ultra 550B A55Bは、Mamba2、LatentMoE、Multi-Token Prediction (MTP)を搭載した550Bパラメータのハイブリッドモデルです。フルファインチューニングをベンチマークします：すべてのパラメータが更新され、Adamオプティマイザ状態が実体化されます。この規模では16 H100ノード（128 GPU）に及びます。

Methodology:

Why there is no Transformers v5 column. Transformers v5はこの規模でメモリ不足になるため、v5の数値は報告されません。AutoModelのExpert ParallelismはエキスパートをGPU間でシャーディングしてフットプリントを予算内に収めるため、フルファインチューニングが実行可能になります。以下の30B比較では、v5が収まる場合でも同じ利点を示します。

Single-node 30B MoE benchmarks

8x H100 80GB GPUを搭載したシングルノードで、3つのアプローチをベンチマークしました：HF Transformers v4（hubコード）、HF Transformers v5（利用可能な最適化を適用）、およびNeMo AutoModel（EP=8 + カスタムカーネル）です。

Methodology:

A note on the routing gate. 以下のNeMo AutoModelの数値は、トークンをエキスパート間で均等に分散させるバランスド・ルーティングゲートを使用しています。これはMoEがトレーニングされる理想的な動作点をエミュレートします：十分にトレーニングされたモデルのロードバランシング損失は、エキスパート利用率をほぼ均一に駆動するため、バランスド・ルーティングは実際のワークロードが収束する定常状態を反映し（ランダムなダミートークンがエキスパート並列性に注入するストラグラーノイズを除去します）。v4/v5は同じダミートークンでネイティブルーターを実行します。したがって、バランスドゲートはターゲットMoE動作点でのNeMo AutoModelを測定し、v4/v5の列はそのままの動作を反映します。

Qwen3-30B-A3B

Why v4 deadlocks: Transformers v4はQwen3 MoEエキスパートを128個の個別MLPモジュールのModuleListとして保存し、それぞれが個別にFSDPラップされます。フォワードパスは、トークンを受け取ったエキスパートのみを反復するデータ依存のループを使用します。ランクごとにデータが異なるため、異なるランクが異なるエキスパートをスキップし、FSDP AllGather/ReduceScatter集団操作のミスマッチを引き起こし、無期限にハングします。Transformers v5は、エキスパートを融合3Dパラメータテンソル（エキスパートごとのモジュールなし、エキスパートごとのFSDP集団操作なし）として保存することでこれを修正します。

Nemotron 3 Nano 30B A3B

v4 config: trust_remote_code=True（NVIDIAのhubモデリングコード）。hubコードのエキスパートループはFSDPセーフ（トークン割り当てに関係なくすべてのエキスパートを反復）であるため、Qwen3 v4のようにデッドロックしません。

Where the speedup comes from

NeMo AutoModelがTransformers v5に対して達成する3.4〜3.7倍の高速化は、3つの要因によるものです：

1. Expert Parallelismはメモリプレッシャーを軽減します。 EP=8はエキスパート重みをGPU間で分散し、GPUあたりのMoEフットプリントを8倍削減します。Qwen3の場合、ピークメモリを68.2 GiBから48.1 GiBに削減（-29%）。Nemotron Nanoの場合、62.1 GiBから42.5 GiBに削減（-32%）し、より大きなバッチサイズや長いシーケンスのための余裕を生み出します。

2. DeepEPは通信と計算を融合します。 エキスパートルーティングのための個別のAllGather/ReduceScatter集団操作の代わりに、DeepEPはトークンディスパッチとコンバインを最適化されたGPUカーネルに融合し、通信とエキスパート計算をオーバーラップさせます。

3. TransformerEngineカーネルはコア演算を高速化します。 TEの融合アテンション、線形レイヤー、RMSNorm実装は、MoEレイヤーだけでなくすべてのレイヤータイプでPyTorch/Flash Attention相当品に対して一貫した高速化を提供します。

Transformers v5 Features Leveraged by HuggingFace AutoModel

Expert Backends

Transformers v5の最も影響力のある機能の1つは、experts_implementationパラメータで、3つのエキスパートバックエンドを含みます：

grouped_mmバックエンドは主要なトレーニング最適化です：エキスパートを1つずつループする代わりに、トークンを割り当てられたエキスパートでソートし、単一の融合グループ化行列乗算を実行します。

NeMo AutoModelはこれをさらに進めます。カスタム実装を持つモデルの場合、DeepEP融合オールトゥオールディスパッチとグループ化GEMMカーネル、TransformerEngine線形レイヤーを組み合わせて使用します。進化は次のようになります：

v4 (eager for-loop) → v5 (grouped_mm) → NeMo AutoModel (DeepEP + GMM + TE)

NeMo AutoModelでは、エキスパートバックエンドはBackendConfigを通じて設定されます：

from nemo_automodel.components.models.common.utils import BackendConfig

backend = BackendConfig(
    attn="te",           # TransformerEngine attention
    linear="te",         # TransformerEngine linear layers
    experts="torch_mm",  # Grouped expert matmul
    dispatcher="deepep", # DeepEP fused all-to-all
)

Expert Parallelism and DeepEP

Transformers v5はExpert Parallelismパスも提供します。エキスパート重みをGPU間でシャーディングします。GroupedGemmParallelスタイルは各デバイスのローカルエキスパートのみをロードし、RouterParallelはトークンをルーティングしてall_reduceで結果を結合します。v5の既存のテンソル並列機構の上にきれいに構築されています。有効にすると、モデルのtp_planがexpert planを返し、エキスパート並列性はデータ並列性とデバイス予算を共有します（ep × dp = world_size）。ここでのシングルノード30Bベンチマークでは、プレーンなデータ並列v5（dp=8、ep=1）が最速のv5構成であることがわかったため、それが報告するv5設定です。

NeMo AutoModelは、マルチGPU MoEトレーニング向けに調整された補完的なアプローチを取ります。EPを独自の並列性次元とし、データ並列メッシュから切り出されるのではなく、専用moe_meshとして使用し、PyTorchのDTensorとShard(0)を使用します。エキスパートメッシュはデータ並列性と直交するため、両者は同じデバイス上で構成されます。8 GPU上でNeMo AutoModelはep=8とdp=8を同時に実行するため、各GPUは独自のデータシャードをトレーニングしながら、エキスパートの1/8のみを保持します。エキスパート重みはエキスパート次元に沿ってGPU間で物理的にシャーディングされます。

# From nemo_automodel/components/moe/parallelizer.py
from torch.distributed.tensor import Shard, distribute_tensor

# Each GPU holds only 1/ep_size of the expert weights
distribute_tensor(param, device_mesh, [Shard(0)])

8 GPU上でep_size=8の場合、各GPUはエキスパートパラメータの1/8のみを保持します。~55 GiBのエキスパート重みを持つNemotron-3-Nano-30B-A3Bのようなモデルの場合、EPはGPUあたりのエキスパートフットプリントを~55 GiBから~6.8 GiBに削減し、FSDPのみのアプローチでメモリ不足になる状況でトレーニングを可能にします。

EPの上に、NeMo AutoModelはトークンルーティングを最適化されたGPUカーネルに融合するDeepEPを統合し、グループ化エキスパート計算のためのグループ化GEMMと組み合わせることで大幅な高速化を実現します。大規模MoEベンチマークでは、DeepEP + grouped GEMMにより、all-gather + ループエキスパートベースラインと比較して、完全なDeepSeek V3 671Bモデルで反復あたりのコストを47%削減しました。

Dynamic Weight Loading

Transformers v5は、WeightConverterとWeightRenamingを通じて動的重みロードシステムも導入しました。これにより、MoEチェックポイントをより効率的な実行のために融合3Dテンソルとして保存できます。WeightConverterはfrom_pretrained()中にチェックポイントテンソルをオンザフライで変換するコンポーザブルな演算を適用します。

NeMo AutoModelはこのv5 APIの直接的なコンシューマーです。20以上のモデルタイプがMODELS_REQUIRING_TENSOR_MERGINGを通じてこのメカニズムを使用しており、Mixtral、Qwen2 MoE、Qwen3 MoE、DeepSeek V2/V3、OLMoEなどが含まれます。変換は完全にリバーシブルです：save_pretrained()は、任意の下流ツールがロードできる標準のHF形式チェックポイントを生成します。

Getting Started

NeMo AutoModelを試すには、公式ドキュメントページにアクセスしてget startedしてください。

詳細については、以下を参照してください：

• NeMo AutoModel HuggingFace API Compatibility Guide
• NeMo AutoModel Model Coverage
• NeMo AutoModel Performance Summary
• NeMo AutoModel on HuggingFace

Conclusion

NVIDIA NeMo AutoModelは、モデルトレーニングをスケールアップするHuggingFaceユーザーにとって自然な次のステップです。Transformers v5の上に直接構築することで、AutoModelは摩擦ゼロのアップグレードパスを提供します：1行のインポートを変更するだけで、3倍以上高速なモデルインスタンスが得られます。

Qwen3-30B-A3BとNemotron 3 Nano 30B-A3Bでは、Transformers v5の最適構成と比較して3.4〜3.7倍のトレーニングスループット向上と29〜32%のGPUメモリ削減を実現します。また、真のExpert ParallelismはエキスパートをGPU間でシャーディングするため、同じパスでNemotron 3 Ultraのような550Bモデルを16ノードにわたってフルファインチューニングにスケールアップでき、Expert Parallelismがメモリにモデルを収めるために不可欠になるレジームに対応します。NeMo AutoModelチェックポイントは標準のHF形式safetensorsであるため、vLLMやSGLangなどの推論フレームワークにデプロイできます。

コード、設定、ベンチマークスクリプトはすべてNeMo AutoModelリポジトリで入手できます。

Acknowledgements

この作業の主要貢献者（姓のアルファベット順）：Adil Asif, Hemil Desai, Alexandros Koumparoulis, and Huiying Li.