nemo-mbridge-perf-cuda-graphs

CUDA Graphs

稳定版文档：@docs/training/cuda-graphs.md 卡片：@skills/nemo-mbridge-perf-cuda-graphs/card.yaml

简介

CUDA 图会捕获 GPU 操作，并在之后重放，以尽可能减少主机驱动开销。Bridge 支持两种实现：

cuda_graph_impl	机制	支持的范围
"local"	MCore FullCudaGraphWrapper，包装整个前向传播和反向传播	full_iteration
"transformer_engine"	对每一层调用 TE make_graphed_callables()	attn、mlp、moe、moe_router、moe_preprocess、mamba

快速决策

对于大多数训练工作负载，先使用 TE 范围图，然后在相同的 dispatcher、布局和容器中，将重放耗时与 eager 模式进行对比：

• 稠密模型：使用 attn，然后可选地使用 mlp
• 不丢弃 token 的 MoE：使用 attn moe_router moe_preprocess
• VLM：使用相同的不丢弃 token 的 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 范围图（稠密模型）

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

性能测试工具命令行

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

有效的命令行值位于 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 中强制执行）
• 只有 cuda_graph_impl = "local" 时才能使用 full_iteration 范围
• full_iteration 范围要求 check_for_nan_in_loss = False
• 不要同时使用 moe 范围和 moe_router 范围
• 张量形状必须固定（固定的 seq_length、固定的 micro_batch_size）
• MoE 的 token 不丢弃路由会将可捕获的范围限制在稠密模块
• 使用 PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True 时，设置 NCCL_GRAPH_REGISTER=0（MCore 会针对架构低于 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 枚举：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 跟踪器是必需的：设置 cuda_graph_impl 时，如果未设置 use_te_rng_tracker=True 和 rng.te_rng_tracker=True，provider 中会触发断言。

2. full_iteration 要求禁用 NaN 检查：捕获范围涵盖整个 fwd+bwd，因此无法通过损失 NaN 检查来检查中间值。

3. MoE 范围限制：moe 范围与 moe_router 范围互斥。无丢弃 Token 的 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 架构：在计算能力 < 10.0 的 GPU（Blackwell 之前的架构）上，如果使用 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_granularity="full" 和 recompute_num_layers（重计算 N 个完整 Transformer 层）时，与 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，在 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 错误或非法内存访问。