Profile, validate, and use CUDA graph capture in Megatron Bridge, including local full-iteration graphs and Transformer Engine scoped graphs. Covers replay evidence, matched eager A/Bs, model-specific scopes, regressions, and failures.
Profile, validate, and use CUDA graph capture in Megatron Bridge, including local full-iteration graphs and Transformer Engine scoped graphs. Covers replay evidence, matched eager A/Bs, model-specific scopes, regressions, and failures.
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'.
First confirm a profile shows material host/launch gaps. Then test the narrowest
TE-scoped graph candidate and verify replay timing against eager on the same
dispatcher, layout, routing, precision, and container:
dense models: attn, then optionally mlp
dropless MoE: moe_router moe_preprocess, then optionally attn
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
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.
Scope names follow the actual MCore module type. mamba applies to real Mamba
layers; do not use it as a generic label for every linear-attention block.
Qwen3.5 Gated DeltaNet (GDN) lives under self_attention, so its applicable TE
scope is attn, not mamba.
Required constraints
use_te_rng_tracker = True (enforced in gpt_provider.py)
full_iteration scope only with cuda_graph_impl = "local"
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.
Re-run the eager/graph A/B on the final dispatcher, overlap, precision, and
recompute stack. An earlier graph win can disappear after other tuning.
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()
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, ...
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.
Scope names are structural, not architectural nicknames: use mamba
only for real Mamba layers. Hybrid models such as Qwen3.5 place GDN in the
attention path, so graph it with attn when supported.
