Nemo Mbridge Perf Cuda Graphs

NVIDIAOfficialPublished by NVIDIA

Skill
Tool integration
Software development
Apache-2.0
Medium risk

From the author

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.

— the description by NVIDIA

AutoGineer summary

Written by AutoGineer from the skill’s own files, not by the author

This skill provides guidance on implementing CUDA graphs within the Megatron Bridge framework to minimize host-driver overhead. It covers two primary implementations: 'local' for full-iteration capture and 'transformer_engine' for scoped capture of specific modules like attention, MLP, and MoE. The documentation details configuration requirements, mandatory RNG tracker settings, performance harness usage, and critical constraints such as static tensor shapes and memory overhead considerations.

What it does

Enables CUDA graph capture for reduced host-driver overhead.
Supports local full-iteration graphs and Transformer Engine scoped graphs.
Provides configuration templates for dense and MoE models.
Includes a performance harness CLI for testing graph implementations.
Enforces mandatory constraints like TE RNG tracker requirements.
Offers verification steps via unit and functional tests.

When to use it

When aiming to reduce host-driver overhead in training workloads.
When troubleshooting performance regressions related to CUDA graph configurations.
When optimizing dense models or dropless MoE architectures.
When transitioning from eager execution to graphed callables.

Advantages

Reduces CPU-side overhead by capturing and replaying GPU operations.
Provides flexible scoping options to target specific performance bottlenecks.
Integrates directly with Megatron Bridge configuration and training loops.
Includes built-in validation for complex constraints like recompute and RNG tracking.

Usage examples

“How do I enable TE-scoped CUDA graphs for a dense model?”: Set cfg.model.cuda_graph_impl to 'transformer_engine', define the scope (e.g., ['attn', 'mlp']), and ensure use_te_rng_tracker is enabled.
“What are the requirements for using full-iteration CUDA graphs?”: Use 'local' implementation, set scope to ['full_iteration'], and disable check_for_nan_in_loss in the configuration.

Install

Download the skill as a ZIP:
https://mp-staging.autogineer.com/agent-skills/nvidia/nemo-mbridge-perf-cuda-graphs/download

Unzip it into the cross-client skills folder of your project:

mkdir -p .agents/skills && unzip nemo-mbridge-perf-cuda-graphs.zip -d .agents/skills

Documentation by NVIDIA

CUDA Graphs

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

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

Stabilize the eager run first.
Fix sequence length and micro-batch size.
Enable the narrowest useful graph scope.
Confirm replay is active and memory is still acceptable.
Compare eager against graph replay iterations after warmup and capture; do not include the capture step in steady-state timing.
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_backward_func = FullCudaGraphWrapper(
            forward_backward_func, cuda_graph_warmup_steps=config.model.cuda_graph_warmup_steps
        )

TE graph capture after warmup

        # 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

        _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,
        )

Delayed wgrad + CUDA graph interaction

            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

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

Graph cleanup

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/full_cuda_graph.py
CudaGraphScope enum: megatron/core/transformer/enums.py

Positive recipe anchors

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

File	Coverage
`tests/unit_tests/training/test_config.py`	`full_iteration` NaN-check constraint
`tests/unit_tests/training/test_comm_overlap.py`	`delay_wgrad` + CUDA graph interaction
`tests/unit_tests/models/test_gpt_full_te_layer_autocast_spec.py`	TE autocast with CUDA graphs
`tests/functional_tests/test_groups/recipes/test_llama_recipes_pretrain_cuda_graphs.py`	End-to-end local and TE graph smoke tests
`tests/unit_tests/recipes/kimi/test_kimi_k2.py`	TE + CUDA graph recipe config
`tests/unit_tests/recipes/gpt/test_gpt3_175b.py`	TE + CUDA graph recipe config
`tests/unit_tests/recipes/qwen_vl/test_qwen25_vl_recipes.py`	VLM CUDA graph settings

Pitfalls

TE RNG tracker is mandatory: Setting cuda_graph_impl without use_te_rng_tracker=True and rng.te_rng_tracker=True will assert in the provider.
full_iteration requires NaN checks disabled: The entire fwd+bwd is captured, so loss-NaN checking cannot inspect intermediate values.
MoE scope restrictions: moe scope and moe_router scope are mutually exclusive. Token-dropless MoE can only graph moe_router and moe_preprocess, not the full expert dispatch.
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.
Delayed wgrad interaction: When delay_wgrad_compute=True and attention or MoE router is in cuda_graph_scope, additional constraints apply: TE >= 2.12.0, gradient_accumulation_fusion=True, and no attention bias.
Variable-length sequences break graphs: Sequence lengths must be constant across steps. Use padded packed sequences if packing is needed.
Graph cleanup is required: CUDA graph objects hold NCCL buffer references. Bridge handles this in _delete_cuda_graphs() at the end of training, but early exits must call it explicitly.
Older GPU architectures: On GPUs with compute capability < 10.0 (pre-Blackwell), set NCCL_GRAPH_REGISTER=0 when using PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True. Enforced in MCore CudaGraphManager (cuda_graphs.py:1428) and TECudaGraphHelper (cuda_graphs.py:1697). The TE impl asserts unconditionally regardless of arch.
CPU offloading incompatible: CUDA graphs cannot be used with CPU offloading. Enforced in MCore transformer_config.py:1907.
MoE recompute + moe_router scope: MoE recompute is not supported with moe_router CUDA graph scope when using cuda_graph_impl = "transformer_engine". Enforced in MCore transformer_config.py:1977.
Layer-level recompute requires full_iteration scope: Using recompute_granularity="full" with recompute_num_layers (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, moe_router, 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. LLAMA3_70B_SFT_CONFIG_H100_FP8_CS_V1 uses cuda_graph_impl= "transformer_engine", cuda_graph_scope="mlp"). Fix: use submodule recompute (recompute_granularity="selective" + recompute_modules), disable CUDA graphs, or switch to local + full_iteration. Enforced in MCore transformer_config.py:2001-2005. See also @skills/nemo-mbridge-perf-activation-recompute/SKILL.md.
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.
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,moe_router,moe_preprocess 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.

Files

ATTRIBUTION.mdAdded by AutoGineer857 B
BENCHMARK.md3.9 KiB
LICENSELicence file added by AutoGineer11.2 KiB
SKILL.md13.8 KiB
card.yaml13.5 KiB
evals/evals.json1.4 KiB
skill-card.md4.1 KiB
skill.oms.sig4.7 KiB

card.yaml

title: cuda_graphs
validated_on: "2026-05-18"
summary: >
  Megatron Bridge supports CUDA graph capture through two implementations:
  local full-iteration graphs (MCore FullCudaGraphWrapper) and Transformer
  Engine scoped graphs (TECudaGraphHelper) for fine-grained capture of
  attention, MLP, and MoE router modules. Requires static tensor shapes
  and TE RNG tracker. Verified on Qwen3-30B-A3B and GPT-OSS-20B MoE pretrain
  with TE-scoped graphs, showing roughly 15-25% lower iteration time and
  20-33% higher throughput when the eager baseline is launch-bound. A later
  H100 Qwen3-30B-A3B all-to-all short run captured successfully but was
  neutral/slightly slower in replay, so target-stack validation is required.
  Large-model rollout remains memory-sensitive (GPT-OSS-120B blocked OOM),
  and packed-sequence finetuning remains fragile: Qwen3 SFT hits the
  non-Tensor `packed_seq_params` graph assertion, while GPT-OSS SFT/LoRA were
  blocked earlier by missing TE packed-sequence attention backends in the
  tested container.
validation_status:
  config_validation:
    - code_verified  # config.py:1524-1531
  te_rng_tracker_requirement:
    - code_verified  # gpt_provider.py:213-217
  full_iteration_nan_check_constraint:
    - code_verified  # config.py:1525-1529
  te_scoped_graph_capture:
    - code_verified  # train.py:231-255, 338-350
  delayed_wgrad_cuda_graph_interaction:
    - code_verified  # comm_overlap.py:522-555
  moe_moe_router_mutual_exclusion:
    - code_verified  # MCore transformer_config.py:1928-1931
  cpu_offloading_incompatible:
    - code_verified  # MCore transformer_config.py:1907-1908
  nccl_graph_register_env:
    - code_verified  # MCore cuda_graphs.py:1428-1435, 1697-1703
  graph_cleanup:
    - code_verified  # train.py:1414-1443
  perf_harness_override:
    - code_verified  # overrides.py:102-124
  test_functions_exist:
    - code_verified  # grep confirmed test function names in test files
  end_to_end_functional_smoke:
    - unclear  # test file exists but tests were not executed
  qwen3_moe_pretrain_te_scoped:
    - measured  # Qwen3-30B-A3B, TP2 PP2 EP4, 2x H100 nodes
  qwen3_moe_h100_alltoall_te_scoped_neutral:
    - measured  # Qwen3-30B-A3B H100 BF16, alltoall, replay 42.00s vs 41.36s eager
  qwen3_sft_packed_sequence_block:
    - measured  # packed_seq_params is not a Tensor for TE-scoped graphs
  gpt_oss_20b_pretrain_te_scoped:
    - measured  # GPT-OSS-20B, TP2 PP4 EP4, 2x H100 nodes, job 10111169
  gpt_oss_20b_sft_lora_te_backend_block:
    - measured  # baseline and graph blocked by missing packed-sequence TE attention backends
  gpt_oss_120b_pretrain_oom:
    - measured  # GPT-OSS-120B, TP2 PP4 EP8, 4x H100 nodes, job 10111518
training_dimensions:
  speed:
    effect: "+/-"
    confidence: medium
    rationale: >
      Measured throughput gains on two MoE pretrain workloads with TE-scoped
      graphs: Qwen3-30B-A3B improves from 623ms to 484ms steady-state and
      214 to 274 TFLOP/s/GPU; GPT-OSS-20B improves from 467-520ms to
      391-399ms and 37.9-42.2 to 49.4-50.4 TFLOP/s/GPU. A 2026-05-18
      Qwen3-30B-A3B H100 BF16 all-to-all run captured 48 TE-scoped graphable
      layers but replay was slightly slower than eager, so speedup depends on
      launch overhead and dispatcher/container shape.
  memory:
    effect: "-"
    confidence: medium
    rationale: >
      No extra memory was observed in the pre-capture reports for
      Qwen3-30B-A3B or GPT-OSS-20B, but GPT-OSS-120B OOMed at iteration 2
      with roughly 69-70 GB already allocated on 79 GB H100s, so rollout is
      clearly headroom-sensitive.
  scale:
    effect: "+/-"
    confidence: medium
    rationale: >
      CUDA graphs can help at scale when launch overhead matters, but larger
      MoE topologies may fail on memory before capture benefits are realized.
  convergence:
    effect: "0"
    confidence: medium
    rationale: >
      Qwen3-30B-A3B pretrain matched baseline within 0.001, but the
      GPT-OSS-20B loss comparison was inconclusive because the short run used
      mock data, GBS=4, and a production LR that made the loss curve noisy.
  stability:
    effect: "-"
    confidence: high
    rationale: >
      Adds shape, scope, environment, and backend constraints that can block
      capture or fail before capture starts.
enable_when:
  - sequence length and micro-batch size are static across steps
  - host overhead is a meaningful part of step time
  - the run has memory headroom for pinned graph buffers
  - throughput improvement is desired without changing training math
  - pretrain or other static-shape workloads can avoid packed-sequence paths
avoid_when:
  - sequence length or batch shapes vary across steps
  - CPU offloading is enabled
  - memory is already tight, especially with pipeline parallelism or large MoE models
  - unsupported runtime checks or scope combinations are required
  - TE-scoped graphs must handle packed_sequence=True workloads
  - the TE/container build lacks the attention backend required by the recipe
  - full activation recompute with TE-scoped graphs is required
interactions:
  required:
    - model.use_te_rng_tracker
    - rng.te_rng_tracker
  conditional:
    - rerun_state_machine.check_for_nan_in_loss must be false for local + full_iteration
    - delay_wgrad_compute adds TE/version/fusion constraints for captured attention or moe_router scopes
    - moe_preprocess requires moe_router
    - expandable_segments may require NCCL_GRAPH_REGISTER=0
    - recompute_granularity=full only works with local full_iteration, not TE scoped
    - packed_sequence=True is incompatible with Qwen3-style TE-scoped graph capture because packed_seq_params is not a Tensor
    - older TE/container builds may fail packed-sequence attention before any graph-specific behavior is reached
  incompatible:
    - cpu_offloading
    - moe with moe_router
    - recompute_granularity=full with transformer_engine impl
feature_meaning:
  cuda_graph_impl: >
    Which graph capture backend to use: local (MCore), transformer_engine (TE), or none.
  cuda_graph_scope: >
    Which modules to capture: full_iteration, attn, mlp, moe, moe_router, moe_preprocess, mamba.
  cuda_graph_warmup_steps: >
    Number of eager warmup steps before graph capture begins (default 3).
config_keys:
  - model.cuda_graph_impl
  - model.cuda_graph_scope
  - model.cuda_graph_warmup_steps
  - model.use_te_rng_tracker
  - rng.te_rng_tracker
  - rerun_state_machine.check_for_nan_in_loss
  - ddp.check_for_nan_in_grad
recommended_path:
  model.cuda_graph_impl: transformer_engine_for_scoped_or_local_for_full
  model.cuda_graph_scope: attn_plus_moe_router_moe_preprocess_for_moe_models
  rng.te_rng_tracker: true
  model.use_te_rng_tracker: true
expected_metric_change:
  - metric: step_time
    direction: down
    magnitude: ~15-25%
    conditions: MoE pretrain, TE scoped (attn+moe_router+moe_preprocess), static shapes
    evidence: measured_qwen3_30b_a3b_and_gpt_oss_20b
  - metric: tokens_per_sec
    direction: up
    magnitude: ~20-33%
    conditions: same as step_time
    evidence: measured_qwen3_30b_a3b_and_gpt_oss_20b
  - metric: peak_memory
    direction: neutral_pre_capture_but_headroom_sensitive
    magnitude: none_observed_pre_capture_on_smaller_runs
    conditions: TE scoped graphs on H100 80GB
    evidence: measured_qwen3_30b_a3b_and_gpt_oss_20b_plus_gpt_oss_120b_oom
measured_results:
  - model: Qwen3-30B-A3B
    task: pretrain
    parallelism: TP2_PP2_EP4
    nodes: 2
    gpu: H100_80GB
    scope: "attn+moe_router+moe_preprocess"
    impl: transformer_engine
    baseline_iter_ms: 623
    graph_iter_ms: 484
    speedup_pct: 22
    baseline_tflops: 214
    graph_tflops: 274
    loss_match: true
    loss_max_delta: 0.001
    graphable_layers_per_pp: 24
    capture_time_s: 5.6
    status: success
  - model: Qwen3-30B-A3B
    task: sft_packed
    parallelism: TP2_PP2_EP4
    nodes: 2
    gpu: H100_80GB
    scope: "attn+moe_router+moe_preprocess"
    impl: transformer_engine
    baseline_iter_ms: 880
    status: blocked
    error: "packed_seq_params not a Tensor; TE-scoped graphs incompatible with packed sequences"
  - model: GPT-OSS-20B
    task: pretrain
    parallelism: TP2_PP4_EP4_CP1
    nodes: 2
    gpu: H100_80GB
    scope: "attn+moe_router+moe_preprocess"
    impl: transformer_engine
    baseline_iter_ms_range: "467-520"
    graph_iter_ms_range: "391-399"
    speedup_pct_range: "16-24"
    baseline_tflops_range: "37.9-42.2"
    graph_tflops_range: "49.4-50.4"
    loss_match: inconclusive
    loss_note: "first ~10 post-capture iterations close; later divergence is noise-dominated"
    graphable_layers_per_pp: 6
    capture_time_s: 0.95
    memory_delta: none_observed_pre_capture
    jobs: "10111169"
    status: success
  - model: GPT-OSS-20B
    task: sft_lora_packed
    parallelism: recipe_default_packed_sequence_paths
    nodes: 1
    gpu: H100_80GB
    scope: "attn+moe_router+moe_preprocess"
    impl: transformer_engine
    status: blocked
    error: "TE attention backend unavailable in tested container; baseline and graph both fail"
    jobs: "10111729/10111730"
  - model: GPT-OSS-120B
    task: pretrain
    parallelism: TP2_PP4_EP8_CP1
    nodes: 4
    gpu: H100_80GB
    scope: "attn+moe_router+moe_preprocess"
    impl: transformer_engine
    status: blocked
    error: "OOM on iteration 2 while allocating 1.54 GiB"
    memory_at_iter1_allocated_gb: "69-70"
    memory_at_iter1_reserved_gb: "72-73"
    jobs: "10111518"
minimal_example:
  docs_example: docs/training/cuda-graphs.md
  skill_example: skills/nemo-mbridge-perf-cuda-graphs/SKILL.md
  test_path: tests/functional_tests/test_groups/recipes/test_llama_recipes_pretrain_cuda_graphs.py
  cli_path: scripts/performance/run_script.py
failure_modes:
  - name: missing_te_rng_tracker
    symptom: provider assertion before training starts
    likely_cause: use_te_rng_tracker or rng.te_rng_tracker not enabled
  - name: illegal_scope_combination
    symptom: config validation assertion before capture
    likely_cause: invalid impl and scope pairing or unsupported scope mix
  - name: full_recompute_with_te_scoped
    symptom: "AssertionError: full recompute is only supported with full iteration CUDA graph"
    likely_cause: recompute_granularity=full with any TE-scoped graph (attn, mlp, moe_router, etc.). Common on FP8 CS configs that default to cuda_graph_impl=transformer_engine + scope=mlp.
    fix: use recompute_granularity=selective with recompute_modules, or disable CUDA graphs (cuda_graph_impl=none), or switch to cuda_graph_impl=local + cuda_graph_scope=full_iteration. See skills/nemo-mbridge-perf-activation-recompute/SKILL.md.
  - name: packed_sequences_with_te_scoped
    symptom: "AssertionError: CUDA graph accepts only Tensor inputs. packed_seq_params excluded"
    likely_cause: packed_sequence=True passes a non-Tensor packed_seq_params input into TE-scoped capture
    fix: disable packed sequences or use local full-iteration graphs
  - name: packed_sequence_attention_backend_missing
    symptom: "Available backends = {FlashAttention=False, FusedAttention=False, UnfusedDotProductAttention=False}"
    likely_cause: TE/container build does not provide the packed-sequence attention backend required by the recipe
    fix: upgrade the container or TE build before evaluating CUDA graphs
  - name: dynamic_shapes
    symptom: capture or replay failure after warmup
    likely_cause: sequence length or micro-batch size changes across steps
  - name: memory_overhead_oom
    symptom: run fits in eager mode but OOMs after enabling graphs or when close to the memory limit
    likely_cause: graph buffers remain pinned and reduce headroom
  - name: allocator_nccl_env_conflict
    symptom: graph registration or NCCL-related runtime failure
    likely_cause: incompatible allocator and NCCL graph register settings
known_constraints:
  - use_te_rng_tracker must be True when cuda_graph_impl is not none (gpt_provider.py:213).
  - full_iteration scope only works with cuda_graph_impl local (MCore transformer_config.py:1704).
  - full_iteration requires check_for_nan_in_loss False (config.py:1525).
  - moe scope and moe_router scope are mutually exclusive (MCore transformer_config.py:1928).
  - moe_preprocess scope requires moe_router scope (MCore transformer_config.py:1934).
  - CPU offloading is incompatible with CUDA graphs (MCore transformer_config.py:1907).
  - With expandable_segments, NCCL_GRAPH_REGISTER=0 required (MCore cuda_graphs.py:1428, 1697).
  - MoE recompute unsupported with moe_router scope + TE impl (MCore transformer_config.py:1977).
  - full recompute (recompute_granularity=full) only works with local full_iteration, not TE scoped.
known_limitations:
  - Most public recipes default to cuda_graph_impl none.
  - Tensor shapes must be static; variable-length sequences break graph replay.
  - TE-scoped graphs and packed-sequence finetuning remain a fragile combination across model families.
  - Large-model MoE rollout is memory-headroom-sensitive even when smaller-model TE-scoped graphs work.
evidence:
  - docs/training/cuda-graphs.md
  - docs/performance-guide.md
  - src/megatron/bridge/training/config.py
  - src/megatron/bridge/training/train.py
  - src/megatron/bridge/models/gpt_provider.py
  - src/megatron/bridge/training/initialize.py
  - src/megatron/bridge/training/comm_overlap.py
  - scripts/performance/utils/overrides.py
  - tests/unit_tests/training/test_config.py
  - tests/functional_tests/test_groups/recipes/test_llama_recipes_pretrain_cuda_graphs.py
follow_up_validation:
  - Re-run GPT-OSS-20B pretrain with lower LR and/or larger GBS for a meaningful loss comparison.
  - Retry GPT-OSS packed-sequence SFT and LoRA on a newer TE/container build.
  - Re-test GPT-OSS-120B with more PP or memory-tuning changes before concluding rollout feasibility.

Requirements and compatibility

Works with: Claude Code, Claude apps, OpenAI Codex, GitHub Copilot, Cursor, Gemini CLI, Kiro, OpenCode, Goose, Amp, Roo Code, JetBrains Junie, Factory Droid, Devin (ex-Windsurf), Mistral Vibe, OpenHands, Letta

“Works with” means the tool reads the open Agent Skills format; a skill written for one tool may use features others lack.

Needs: Needs Python

Security

What our scanner noticed· Medium risk

Large base64-encoded blobs

A static scan is not a guarantee. Skills are third-party content: review a skill before you let an agent use it.

Attribution and licence

Published by

NVIDIA

Publisher

NVIDIA (Official)

Source

NVIDIA/skills · skills/nemo-mbridge-perf-cuda-graphs · commit fce3cca

Licence

Apache License 2.0 (redistribution allowed)

LICENSE (added by AutoGineer from LICENSE-APACHE)
ATTRIBUTION.md (written by AutoGineer)

Synced

2026-11-08

AutoGineer is not affiliated with or endorsed by NVIDIA.

Download ZIP

The ZIP holds the skill’s files, its licence files and an ATTRIBUTION.md. Download nemo-mbridge-perf-cuda-graphs.zip