CUDA Graphs

Stable documentation: @docs/training/cuda-graphs.md Card: @skills/nemo-mbridge-perf-cuda-graphs/card.yaml

<!-- NVSkills CI refresh: 2026-06-15. No instruction changes. -->

What It Is

CUDA graphs capture GPU operations once and replay them with minimal host-driver overhead. Bridge supports two implementations:

Quick Decision

Start with TE-scoped graphs for most training workloads, then verify replay timing against eager on the same dispatcher, layout, and container:

• dense models: attn, then optionally mlp
• dropless MoE: attn moerouter moepreprocess
• VLMs: the same dropless-MoE scope, but only after the real-data path is stable

Use local + full_iteration only when you specifically want full-iteration capture and can satisfy the tighter constraints.

For recompute-heavy workloads:

• TE-scoped graphs pair naturally with selective recompute
• full recompute usually pushes you toward local full-iteration graphs or away

from graphs entirely

Related docs:

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

Enablement

Local full-iteration graph

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 scoped graph (dense model)

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 scoped graph (MoE model)

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

Performance harness 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 \
  ...

Valid CLI values live in scripts/performance/argument_parser.py:

• VALIDCUDAGRAPHIMPLS: ["none", "local", "transformerengine"]
• VALIDCUDAGRAPHSCOPES: ["fulliteration", "attn", "mlp", "moe", "moerouter", "moepreprocess", "mamba"]

The performance harness uses a comma-separated --cudagraphscope value and auto-enables model.useterngtracker plus rng.terngtracker when --cudagraph_impl is not none.

Required constraints

• useterngtracker = True (enforced in gptprovider.py)
• fulliteration scope only with cudagraph_impl = "local"
• fulliteration scope requires checkfornanin_loss = False
• Do not combine moe scope and moe_router scope
• Tensor shapes must be static (fixed seqlength, fixed microbatch_size)
• MoE token-dropless routing limits graphable scope to dense modules
• With PYTORCHCUDAALLOCCONF=expandablesegments:True, set

NCCLGRAPHREGISTER=0 (MCore enforces for local impl on arch < sm_100; TE impl asserts unconditionally)

• CPU offloading is incompatible with CUDA graphs
• moepreprocess scope requires moerouter scope to also be set

Practical bring-up order

1. Stabilize the eager run first.
2. Fix sequence length and micro-batch size.
3. Enable the narrowest useful graph scope.
4. Confirm replay is active and memory is still acceptable.
5. Compare eager against graph replay iterations after warmup and capture; do

not include the capture step in steady-state timing.

1. Only then widen scope or combine with overlap features.

Code Anchors

Bridge config and validation

```1524:1531:src/megatron/bridge/training/config.py # CUDA graph scope validation: checkfornaninloss must be disabled with fulliteration graph if self.model.cudagraphimpl == "local" and CudaGraphScope.fulliteration in self.model.cudagraphscope: assert not self.rerunstatemachine.checkfornaninloss, ( "checkfornaninloss must be disabled when using fulliteration CUDA graph. " "Set rerunstatemachine.checkfornaninloss=False." ) if self.model.cudagraphimpl == "none": self.model.cudagraph_scope = []

### TE RNG tracker requirement

```213:216:src/megatron/bridge/models/gpt_provider.py
        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'."

Graph creation and capture in training loop

```231:255:src/megatron/bridge/training/train.py # Capture CUDA Graphs. cudagraphhelper = None if modelconfig.cudagraphimpl == "transformerengine": cudagraphhelper = TECudaGraphHelper(...) # ... if config.model.cudagraphimpl == "local" and CudaGraphScope.fulliteration in config.model.cudagraphscope: forwardbackwardfunc = FullCudaGraphWrapper( forwardbackwardfunc, cudagraphwarmupsteps=config.model.cudagraphwarmup_steps )

### TE graph capture after warmup

```338:350:src/megatron/bridge/training/train.py
        # 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 initialization

```199:206:src/megatron/bridge/training/initialize.py setrandomseed( rngconfig.seed, rngconfig.dataparallelrandominit, rngconfig.terngtracker, rngconfig.inferencerngtracker, usecudagraphablerng=(modelconfig.cudagraphimpl != "none"), pgcollection=pg_collection, )

### Delayed wgrad + CUDA graph interaction

```522:555:src/megatron/bridge/training/comm_overlap.py
            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, ...

Perf harness override helper

```102:124:scripts/performance/utils/overrides.py def setcudagraphoverrides( recipe, cudagraphimpl=None, cudagraphscope=None ): # Sets impl, scope, and auto-enables terngtracker

### Graph cleanup

```1414:1441:src/megatron/bridge/training/train.py
def _delete_cuda_graphs(cuda_graph_helper):
    # Deletes FullCudaGraphWrapper and TE graph objects to free NCCL buffers

MCore classes (in 3rdparty/Megatron-LM)

• CudaGraphManager: megatron/core/transformer/cuda_graphs.py
• TECudaGraphHelper: megatron/core/transformer/cuda_graphs.py
• FullCudaGraphWrapper: megatron/core/fullcudagraph.py
• CudaGraphScope enum: megatron/core/transformer/enums.py

Positive recipe anchors

• src/megatron/bridge/perfrecipes/deepseek/gb300/deepseekv3.py
• src/megatron/bridge/perfrecipes/qwen/gb300/qwen3moe.py
• src/megatron/bridge/perfrecipes/gptoss/gb300/gpt_oss.py

Tests

Pitfalls

1. TE RNG tracker is mandatory: Setting cudagraphimpl without

useterngtracker=True and rng.terng_tracker=True will assert in the provider.

1. full_iteration requires NaN checks disabled: The entire fwd+bwd is

captured, so loss-NaN checking cannot inspect intermediate values.

1. MoE scope restrictions: moe scope and moe_router scope are

mutually exclusive. Token-dropless MoE can only graph moerouter and moepreprocess, not the full expert dispatch.

1. Memory overhead: CUDA graphs pin all intermediate buffers for the

graph's lifetime (no memory reuse). TE scoped graphs add a few GB; full-iteration graphs can increase peak memory by 1.5–2×. PP > 1 compounds overhead since each stage holds its own graph.

1. Delayed wgrad interaction: When delaywgradcompute=True and

attention or MoE router is in cudagraphscope, additional constraints apply: TE >= 2.12.0, gradientaccumulationfusion=True, and no attention bias.

1. Variable-length sequences break graphs: Sequence lengths must be

constant across steps. Use padded packed sequences if packing is needed.

1. Graph cleanup is required: CUDA graph objects hold NCCL buffer

references. Bridge handles this in deletecuda_graphs() at the end of training, but early exits must call it explicitly.

1. Older GPU architectures: On GPUs with compute capability < 10.0

(pre-Blackwell), set NCCLGRAPHREGISTER=0 when using PYTORCHCUDAALLOCCONF=expandablesegments:True. Enforced in MCore CudaGraphManager (cudagraphs.py:1428) and TECudaGraphHelper (cudagraphs.py:1697). The TE impl asserts unconditionally regardless of arch.

1. CPU offloading incompatible: CUDA graphs cannot be used with CPU

offloading. Enforced in MCore transformer_config.py:1907.

1. MoE recompute + moe_router scope: MoE recompute is not supported

with moerouter CUDA graph scope when using cudagraphimpl = "transformerengine". Enforced in MCore transformer_config.py:1977.

1. Layer-level recompute requires full_iteration scope: Using

recomputegranularity="full" with recomputenumlayers (recompute N whole transformer layers) is incompatible with TE-scoped graphs. MCore calls this "full" granularity even though you're selecting how many layers — the name refers to recomputing the full layer, not full model. Any TE-scoped scope (attn, mlp, moerouter, etc.) will assert: AssertionError: full recompute is only supported with full iteration CUDA graph. This commonly hits FP8 configs that default to TE-scoped graphs (e.g. LLAMA370BSFTCONFIGH100FP8CSV1 uses cudagraphimpl= "transformerengine", cudagraphscope="mlp"). Fix: use submodule recompute (recomputegranularity="selective" + recomputemodules), disable CUDA graphs, or switch to local + fulliteration. Enforced in MCore transformerconfig.py:2001-2005. See also @skills/nemo-mbridge-perf-activation-recompute/SKILL.md.

1. Benchmark numbers are workload-specific: graph wins are usually real

when host overhead is visible, but the exact gain depends on batch shape, PP depth, recompute, dispatcher backend, and whether the eager baseline was already optimized.

1. A successful capture is not a speedup guarantee: On 2026-05-18,

Qwen3 30B A3B H100 BF16 pretrain with the all-to-all dispatcher captured TE-scoped attn,moerouter,moepreprocess graphs successfully (48 graphable layers, about 6.9 s capture time on rank 0), but replay iterations 5-8 averaged 42.00 s versus 41.36 s for eager. Treat scoped graphs as a bring-up candidate and validate on the target stack.

Verification

Unit tests

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

Functional smoke test (requires GPU)

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

Success criteria

• Unit tests pass, covering config validation for both local and

transformer_engine implementations.

• Functional test completes training steps with both CUDA graph

implementations.

• No NCCL errors or illegal memory access in logs.