⚡

nemo-mbridge-perf-cuda-graphsSkill nemo-mbridge-perf-cuda-graphs

本技能用于在Megatron Bridge中验证和使用CUDA图捕获，包括局部全迭代图和Transformer Engine（TE）作用域图，以降低主机驱动开销并提升训练性能。涵盖CUDA图配置、启用步骤、约束条件、代码集成、性能测试及常见陷阱排查，适用于注意力、MLP、MoE等模块的性能优化。关键词：CUDA图、Megatron Bridge、性能优化、图捕获、Transformer Engine、训练加速、全迭代图、MoE、重计算。

Megatron-Core训练 0 次安装 0 次浏览 更新于 9/7/2026

名称: nemo-mbridge-perf-cuda-graphs 描述: 在Megatron Bridge中验证和使用CUDA图捕获，包括局部全迭代图和Transformer Engine作用域图，用于注意力、MLP和MoE模块。 许可证: Apache-2.0 使用时机: 通过CUDA图减少主机驱动开销，或将崩溃或回归追溯到CUDA图配置更改；‘cuda_graph_impl’，‘full iteration graph’，‘TE scoped graph’，‘graphed callables’，‘CUDA graph capture’。

CUDA图

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

<!-- NVSkills CI刷新：2026-06-15。无指令更改。 -->

它是什么

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作用域图开始，然后在相同的调度器、布局和容器上验证重放计时与eager对比：

  • 稠密模型：attn，然后可选mlp
  • 无drop MoE：attn moe_router moe_preprocess
  • VLM：与无drop MoE相同的作用域，但仅在真实数据路径稳定后

当您特别需要全迭代捕获且能满足更严格的约束时，使用local + full_iteration。

对于重计算繁重的工作负载：

  • TE作用域图自然与选择性重计算配合
  • 完全重计算通常促使您使用local全迭代图或完全不用图

相关文档：

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

启用

局部全迭代图

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"]           # 或 ["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

性能测试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 \
  ...

有效的CLI值位于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作用域
  • 张量形状必须静态（固定的序列长度、固定的微批大小）
  • MoE token无drop路由将可图化作用域限制在稠密模块
  • 当使用PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True时，设置NCCL_GRAPH_REGISTER=0（MCore对local实现在sm_100以下的架构强制要求；TE实现无条件断言）
  • CPU卸载与CUDA图不兼容
  • moe_preprocess作用域要求同时设置moe_router作用域

实际启动顺序

  1. 先稳定eager运行。
  2. 固定序列长度和微批大小。
  3. 启用最窄的有用图作用域。
  4. 确认重放已激活且内存仍可接受。
  5. 在预热和捕获后比较eager与图重放迭代；不要将捕获步骤计入稳态计时。
  6. 然后才扩大作用域或与重叠功能组合。

代码锚点

Bridge配置和验证

        # CUDA图作用域验证：使用full_iteration图时必须禁用check_for_nan_in_loss
        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'."

训练循环中的图创建和捕获

    # 捕获CUDA图。
    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图捕获

        # 在预热后捕获CUDA图。
        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 []
            # ... 作用域解析 ...
            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
):
    # 设置impl、作用域，并自动启用te_rng_tracker

图清理

def _delete_cuda_graphs(cuda_graph_helper):
    # 删除FullCudaGraphWrapper和TE图对象以释放NCCL缓冲区

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 TE autocast与CUDA图
tests/functional_tests/test_groups/recipes/test_llama_recipes_pretrain_cuda_graphs.py 端到端local和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检查：整个前向+反向被捕获，因此loss-NaN检查无法检查中间值。

  3. MoE作用域限制：moe作用域和moe_router作用域互斥。无drop MoE只能图化moe_router和moe_preprocess，不能图化完整的专家分发。

  4. 内存开销：CUDA图在图的生存期内固定所有中间缓冲区（无内存重用）。TE作用域图增加几GB；全迭代图可能使峰值内存增加1.5–2倍。PP > 1会叠加开销，因为每个stage都有自己的图。

  5. 延迟wgrad交互：当delay_wgrad_compute=True且注意力或MoE路由器在cuda_graph_scope中时，有额外约束：TE >= 2.12.0，gradient_accumulation_fusion=True，且无注意力偏置。

  6. 可变长度序列破坏图：序列长度必须在不同步骤间恒定。如果需要打包，请使用填充的打包序列。

  7. 需要图清理：CUDA图对象持有NCCL缓冲区引用。Bridge在训练结束时在_delete_cuda_graphs()中处理此问题，但提前退出必须显式调用它。

  8. 较旧的GPU架构：在计算能力 < 10.0（pre-Blackwell）的GPU上，当使用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. 这通常影响FP8配置，这些配置默认为TE作用域图（例如，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预训练使用all-to-all调度器成功捕获了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错误或非法内存访问。