/nemo-mbridge-perf-cuda-graphs

name: nemo-mbridge-perf-cuda-graphs
description: Validate and use CUDA graph capture in Megatron Bridge, including local full-iteration graphs and Transformer Engine scoped graphs for attention, MLP, and MoE modules.
license: Apache-2.0
when_to_use: Reducing host-driver overhead via CUDA graphs, or tracing a crash or regression to a CUDA graph config change; 'cuda_graph_impl', 'full iteration graph', 'TE scoped graph', 'graphed callables', 'CUDA graph capture'.

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:

| `cuda_graph_impl` | Mechanism | Scope support | |---|---|---| | `"local"` | MCore `FullCudaGraphWrapper` wrapping entire fwd+bwd | `full_iteration` | | `"transformer_engine"` | TE `make_graphed_callables()` per layer | `attn`, `mlp`, `moe`, `moe_router`, `moe_preprocess`, `mamba` |

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 moe_router moe_preprocess`
• 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`:

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

The performance harness uses a comma-separated `--cuda_graph_scope` value and auto-enables `model.use_te_rng_tracker` plus `rng.te_rng_tracker` when `--cuda_graph_impl` is not `none`.

Required constraints

• `use_te_rng_tracker = True` (enforced in `gpt_provider.py`)
• `full_iteration` scope only with `cuda_graph_impl = "local"`
• `full_iteration` scope requires `check_for_nan_in_loss = False`
• Do not combine `moe` scope and `moe_router` scope
• Tensor shapes must be static (fixed seq_length, fixed micro_batch_size)
• MoE token-dropless routing limits graphable scope to dense modules
• With `PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True`, set

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

• CPU offloading is incompatible with CUDA graphs
• `moe_preprocess` scope requires `moe_router` 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. 6. Only then widen scope or combine with overlap features.

Code Anchors

Bridge config and validation

        # 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 tracker requirement

        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

    # 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_backwa