nemo-mbridge-perf-cuda-graphs

CUDA グラフ

安定版ドキュメント：@docs/training/cuda-graphs.md カード：@skills/nemo-mbridge-perf-cuda-graphs/card.yaml

概要

CUDA グラフは GPU 操作を一度キャプチャし、ホストドライバーのオーバーヘッドを最小限に抑えて再実行します。Bridge は 2 つの実装をサポートしています。

cuda_graph_impl	仕組み	スコープのサポート
"local"	MCore の FullCudaGraphWrapper で fwd+bwd 全体をラップ	full_iteration
"transformer_engine"	レイヤーごとに TE の make_graphed_callables() を使用	attn、mlp、moe、moe_router、moe_preprocess、mamba

すぐに判断する

多くのトレーニングワークロードでは、まず TE スコープのグラフから始め、同じディスパッチャー、レイアウト、コンテナーで eager と再実行の時間を比較してください。

• dense モデル：attn、次に必要に応じて mlp
• dropless MoE：attn moe_router moe_preprocess
• VLM：同じ dropless-MoE スコープ。ただし、実データのパスが安定してから使用します

完全なイテレーションのキャプチャを明確に必要としており、より厳しい制約を満たせる場合に限り、local と full_iteration を使用してください。

再計算の多いワークロードの場合：

• TE スコープのグラフは選択的再計算と自然に組み合わせられます
• 完全な再計算では、通常、local の完全なイテレーショングラフを使うか、グラフの使用を避けることになります

関連ドキュメント：

• @docs/training/cuda-graphs.md
• @docs/training/activation-recomputation.md

有効化

Local の完全なイテレーショングラフ

cfg.model.cuda_graph_impl = "local"
cfg.model.cuda_graph_scope = ["full_iteration"]
cfg.model.cuda_graph_warmup_steps = 3
cfg.model.use_te_rng_tracker = True
cfg.rng.te_rng_tracker = True
cfg.rerun_state_machine.check_for_nan_in_loss = False
cfg.ddp.check_for_nan_in_grad = False

TE スコープのグラフ（dense モデル）

cfg.model.cuda_graph_impl = "transformer_engine"
cfg.model.cuda_graph_scope = ["attn"]           # or ["attn", "mlp"]
cfg.model.cuda_graph_warmup_steps = 3
cfg.model.use_te_rng_tracker = True
cfg.rng.te_rng_tracker = True

TE スコープのグラフ（MoE モデル）

cfg.model.cuda_graph_impl = "transformer_engine"
cfg.model.cuda_graph_scope = ["attn", "moe_router", "moe_preprocess"]
cfg.model.cuda_graph_warmup_steps = 3
cfg.model.use_te_rng_tracker = True
cfg.rng.te_rng_tracker = True

パフォーマンスハーネスの CLI

uv run python scripts/performance/run_script.py \
  -m qwen \
  -mr qwen3_30b_a3b \
  --task pretrain \
  -g h100 \
  -c bf16 \
  -ng 16 \
  --cuda_graph_impl transformer_engine \
  --cuda_graph_scope attn,moe_router,moe_preprocess \
  ...

有効な CLI 値は scripts/performance/argument_parser.py にあります。

• VALID_CUDA_GRAPH_IMPLS：["none", "local", "transformer_engine"]
• VALID_CUDA_GRAPH_SCOPES：["full_iteration", "attn", "mlp", "moe", "moe_router", "moe_preprocess", "mamba"]

パフォーマンスハーネスでは、--cuda_graph_scope にカンマ区切りの値を指定します。また、--cuda_graph_impl が none 以外の場合は、model.use_te_rng_tracker と rng.te_rng_tracker を自動的に有効にします。

必須の制約

• use_te_rng_tracker = True（gpt_provider.py で強制されます）
• full_iteration スコープは cuda_graph_impl = "local" の場合にのみ使用できます
• full_iteration スコープでは check_for_nan_in_loss = False が必要です
• moe スコープと moe_router スコープを併用しないでください
• テンソル形状は静的である必要があります（seq_length と micro_batch_size を固定）
• MoE のトークン脱落なしルーティングでは、グラフ化できるスコープが dense モジュールに限定されます
• PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True の場合は、NCCL_GRAPH_REGISTER=0 を設定してください（MCore は arch < sm_100 で local 実装に対してこれを強制し、TE 実装では無条件にアサートします）
• CPU オフロードは CUDA グラフと互換性がありません
• moe_preprocess スコープを使用するには、moe_router スコープも設定する必要があります

実用的な導入手順

1. まず eager 実行を安定させます。
2. シーケンス長とマイクロバッチサイズを固定します。
3. 役立つ範囲で最も限定的なグラフスコープを有効にします。
4. 再実行が有効になっていることと、メモリ使用量が許容範囲内であることを確認します。
5. ウォームアップとキャプチャの後に、eager 実行とグラフ再実行のイテレーションを比較します。定常状態の計測にはキャプチャのステップを含めないでください。
6. その後にのみ、スコープを広げるか、オーバーラップ機能と組み合わせます。

コードの参照箇所

Bridge の設定と検証

        # CUDA graph scope validation: check_for_nan_in_loss must be disabled with full_iteration graph
        if self.model.cuda_graph_impl == "local" and CudaGraphScope.full_iteration in self.model.cuda_graph_scope:
            assert not self.rerun_state_machine.check_for_nan_in_loss, (
                "check_for_nan_in_loss must be disabled when using full_iteration CUDA graph. "
                "Set rerun_state_machine.check_for_nan_in_loss=False."
            )
        if self.model.cuda_graph_impl == "none":
            self.model.cuda_graph_scope = []

TE RNG トラッカーの要件

        if self.cuda_graph_impl != "none":
            assert getattr(self, "use_te_rng_tracker", False), (
                "Transformer engine's RNG tracker is required for cudagraphs, it can be "
                "enabled with use_te_rng_tracker=True'."

トレーニングループでのグラフ作成とキャプチャ

    # Capture CUDA Graphs.
    cuda_graph_helper = None
    if model_config.cuda_graph_impl == "transformer_engine":
        cuda_graph_helper = TECudaGraphHelper(...)
    # ...
    if config.model.cuda_graph_impl == "local" and CudaGraphScope.full_iteration in config.model.cuda_graph_scope:
        forward_backward_func = FullCudaGraphWrapper(
            forward_backward_func, cuda_graph_warmup_steps=config.model.cuda_graph_warmup_steps
        )

ウォームアップ後の TE グラフキャプチャ

        # Capture CUDA Graphs after warmup.
        if (
            model_config.cuda_graph_impl == "transformer_engine"
            and cuda_graph_helper is not None
            and not cuda_graph_helper.graphs_created()
            and global_state.train_state.step - start_iteration == model_config.cuda_graph_warmup_steps
        ):
            if model_config.cuda_graph_warmup_steps > 0 and should_toggle_forward_pre_hook:
                disable_forward_pre_hook(model, param_sync=False)
            cuda_graph_helper.create_cudagraphs()
            if model_config.cuda_graph_warmup_steps > 0 and should_toggle_forward_pre_hook:
                enable_forward_pre_hook(model)
                cuda_graph_helper.cuda_graph_set_manual_hooks()

RNG の初期化

        _set_random_seed(
            rng_config.seed,
            rng_config.data_parallel_random_init,
            rng_config.te_rng_tracker,
            rng_config.inference_rng_tracker,
            use_cudagraphable_rng=(model_config.cuda_graph_impl != "none"),
            pg_collection=pg_collection,
        )

遅延 wgrad と CUDA グラフの相互作用

            cuda_graph_scope = getattr(model_cfg, "cuda_graph_scope", []) or []
            # ... scope parsing ...
            if wgrad_in_graph_scope:
                assert is_te_min_version("2.12.0"), ...
                assert model_cfg.gradient_accumulation_fusion, ...
                if attn_scope_enabled:
                    assert not model_cfg.add_bias_linear and not model_cfg.add_qkv_bias, ...

パフォーマンスハーネスのオーバーライドヘルパー

def _set_cuda_graph_overrides(
    recipe, cuda_graph_impl=None, cuda_graph_scope=None
):
    # Sets impl, scope, and auto-enables te_rng_tracker

グラフのクリーンアップ

def _delete_cuda_graphs(cuda_graph_helper):
    # Deletes FullCudaGraphWrapper and TE graph objects to free NCCL buffers

MCore クラス（3rdparty/Megatron-LM 内）

• CudaGraphManager：megatron/core/transformer/cuda_graphs.py
• TECudaGraphHelper：megatron/core/transformer/cuda_graphs.py
• FullCudaGraphWrapper：megatron/core/full_cuda_graph.py
• CudaGraphScope enum：megatron/core/transformer/enums.py

推奨レシピの参照箇所

• src/megatron/bridge/perf_recipes/deepseek/gb300/deepseek_v3.py
• src/megatron/bridge/perf_recipes/qwen/gb300/qwen3_moe.py
• src/megatron/bridge/perf_recipes/gpt_oss/gb300/gpt_oss.py

テスト

ファイル	カバレッジ
tests/unit_tests/training/test_config.py	full_iteration の NaN チェック制約
tests/unit_tests/training/test_comm_overlap.py	delay_wgrad と CUDA グラフの相互作用
tests/unit_tests/models/test_gpt_full_te_layer_autocast_spec.py	CUDA グラフでの TE オートキャスト
tests/functional_tests/test_groups/recipes/test_llama_recipes_pretrain_cuda_graphs.py	ローカルおよび TE グラフのエンドツーエンドのスモークテスト
tests/unit_tests/recipes/kimi/test_kimi_k2.py	TE と CUDA グラフのレシピ設定
tests/unit_tests/recipes/gpt/test_gpt3_175b.py	TE と CUDA グラフのレシピ設定
tests/unit_tests/recipes/qwen_vl/test_qwen25_vl_recipes.py	VLM の CUDA グラフ設定

注意点

1. TE RNG トラッカーは必須です：use_te_rng_tracker=True と rng.te_rng_tracker=True を指定せずに cuda_graph_impl を設定すると、 プロバイダーでアサートが発生します。

2. full_iteration では NaN チェックを無効にする必要があります：fwd+bwd 全体を キャプチャするため、損失の NaN チェックでは中間値を検査できません。

3. MoE スコープの制限：moe スコープと moe_router スコープは 併用できません。トークンをドロップしない MoE では、グラフ化できるのは moe_router と moe_preprocess のみで、エキスパートのディスパッチ全体は グラフ化できません。

4. メモリーオーバーヘッド：CUDA グラフはグラフの存続期間中、すべての中間 バッファーを固定します（メモリーは再利用されません）。TE スコープのグラフでは 数 GB が追加され、イテレーション全体のグラフではピークメモリーが 1.5～2 倍に 増えることがあります。PP > 1 の場合、各ステージが独自のグラフを保持するため、 オーバーヘッドがさらに増加します。

5. 遅延 wgrad との相互作用：delay_wgrad_compute=True を設定し、attention または MoE router が cuda_graph_scope に含まれる場合、追加の制約が適用されます： TE >= 2.12.0、gradient_accumulation_fusion=True、attention bias なし。

6. 可変長シーケンスではグラフが動作しません：シーケンス長はステップ間で一定で ある必要があります。パッキングが必要な場合は、パディングしたパック済みシーケンスを 使用してください。

7. グラフのクリーンアップが必要です：CUDA グラフオブジェクトは NCCL バッファーへの 参照を保持します。Bridge はトレーニングの最後に _delete_cuda_graphs() を呼び出して これを処理しますが、早期終了する場合は明示的に呼び出す必要があります。

8. 古い GPU アーキテクチャ：compute capability < 10.0 （Blackwell より前）の GPU で PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True を 使用する場合は、NCCL_GRAPH_REGISTER=0 を設定してください。MCore の CudaGraphManager（cuda_graphs.py:1428）と TECudaGraphHelper （cuda_graphs.py:1697）で適用されます。TE 実装ではアーキテクチャに関係なく 常にアサートが発生します。

9. CPU オフロードとは併用できません：CUDA グラフは CPU オフロードと併用できません。 MCore の transformer_config.py:1907 で適用されます。

10. MoE 再計算と moe_router スコープ：cuda_graph_impl = "transformer_engine" を使用する場合、MoE 再計算は moe_router CUDA グラフ スコープではサポートされません。MCore の transformer_config.py:1977 で 適用されます。

11. レイヤーレベルの再計算には full_iteration スコープが必要です： recompute_num_layers（トランスフォーマーレイヤー全体を N 個再計算する）と recompute_granularity="full" を使用すると、TE スコープのグラフとは併用できません。 MCore では、再計算するレイヤー数を選択している場合でも、この粒度を「full」と呼びます。 この名前はモデル全体ではなく、レイヤー全体を再計算することを指します。TE スコープ （attn、mlp、moe_router など）を使用すると、次のアサートが発生します： AssertionError: full recompute is only supported with full iteration CUDA graph. これは、デフォルトで TE スコープのグラフを使用する FP8 構成（例： LLAMA3_70B_SFT_CONFIG_H100_FP8_CS_V1 は cuda_graph_impl= "transformer_engine"、cuda_graph_scope="mlp" を使用）でよく発生します。対処方法： サブモジュール再計算（recompute_granularity="selective" と recompute_modules）を 使用する、CUDA グラフを無効にする、または local と full_iteration に切り替えます。 MCore の transformer_config.py:2001-2005 で適用されます。詳細は @skills/nemo-mbridge-perf-activation-recompute/SKILL.md も参照してください。

12. ベンチマーク結果はワークロード固有です：ホストのオーバーヘッドが顕著な場合、 グラフ化による性能向上は通常実際に得られますが、具体的な向上幅はバッチ形状、PP の深さ、 再計算、ディスパッチャーのバックエンド、eager のベースラインがすでに最適化されていたか どうかによって異なります。

13. キャプチャに成功しても高速化が保証されるわけではありません：2026-05-18、 all-to-all ディスパッチャーを使用した Qwen3 30B A3B H100 BF16 の事前学習では、 TE スコープの attn,moe_router,moe_preprocess グラフのキャプチャに成功しました （グラフ化可能なレイヤーは 48、rank 0 でのキャプチャ時間は約 6.9 s）。しかし、 再生イテレーション 5～8 の平均は 42.00 s で、eager の 41.36 s を上回りました。 スコープ付きグラフは初期導入時の候補として扱い、対象のスタックで検証してください。

検証

ユニットテスト

uv run python -m pytest \
  tests/unit_tests/training/test_config.py -k "cuda_graph" \
  tests/unit_tests/training/test_comm_overlap.py -k "cuda_graph" \
  tests/unit_tests/models/test_gpt_full_te_layer_autocast_spec.py -k "cuda_graph" -q

機能スモークテスト（GPU が必要）

uv run python -m pytest \
  tests/functional_tests/test_groups/recipes/test_llama_recipes_pretrain_cuda_graphs.py -q

成功基準

• local と transformer_engine の両実装について、設定の検証を網羅したユニットテストが成功します。
• 機能テストで、両方の CUDA グラフ実装を使用したトレーニングステップが完了します。
• ログに NCCL エラーや不正なメモリアクセスがありません。