Support old and new MCore CUDA graph APIs

src/megatron/bridge/models/gpt_full_te_layer_autocast_spec.py

@@ -20,14 +20,15 @@
 from megatron.core import tensor_parallel
 from megatron.core.fusions.fused_layer_norm import FusedLayerNorm
 from megatron.core.transformer.cuda_graphs import CudaGraphManager
-from megatron.core.transformer.enums import CudaGraphScope
 from megatron.core.transformer.module import MegatronModule
 from megatron.core.transformer.spec_utils import ModuleSpec
 from megatron.core.transformer.transformer_block import TransformerBlockSubmodules, get_num_layers_to_build
 from megatron.core.transformer.transformer_layer import BaseTransformerLayer
 from megatron.core.transformer.utils import make_sharded_tensors_for_checkpoint
 from transformer_engine.pytorch import TransformerLayer
 
+from megatron.bridge.utils.cuda_graph import uses_local_cuda_graph_manager
+
 
 # Copied from  nemo/collections/nlp/models/language_modeling/megatron/gpt_full_te_layer_autocast_spec.py
 class AutocastTransformerLayer(TransformerLayer):
@@ -226,11 +227,7 @@ def __init__(self, config, layer_number=1, hidden_dropout=None, **kwargs):
             transformer_layer_args["ub_atomic_gemm_rs"] = config.tp_comm_atomic_rs
         self.transformer_layer = AutocastTransformerLayer(**transformer_layer_args)
 
-        if (
-            self.config.cuda_graph_impl == "local"
-            and self.training
-            and CudaGraphScope.full_iteration not in self.config.cuda_graph_scope
-        ):
+        if uses_local_cuda_graph_manager(self.config) and self.training:
             assert not config.cpu_offloading and config.recompute_granularity is None, "Cudagraphs not supported"
             self.add_module("cudagraph_manager", CudaGraphManager(config))
 

src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/transformer_config.py

@@ -20,6 +20,8 @@
 from megatron.core.transformer.transformer_config import TransformerConfig
 from transformers.models.qwen3_vl.configuration_qwen3_vl import Qwen3VLTextConfig
 
+from megatron.bridge.utils.cuda_graph import clear_cuda_graph_modules, set_cuda_graph_modules
+
 
 @dataclass
 class Qwen3VLTransformerConfig(TransformerConfig):
@@ -80,8 +82,6 @@ def get_vision_model_config(hf_config, megatron_config=None):
     config.cuda_graph_retain_backward_graph = megatron_config.cuda_graph_retain_backward_graph
     config.cuda_graph_warmup_steps = megatron_config.cuda_graph_warmup_steps
     config.external_cuda_graph = megatron_config.external_cuda_graph
-    config.cuda_graph_impl = megatron_config.cuda_graph_impl
-    config.cuda_graph_scope = megatron_config.cuda_graph_scope
 
     config.num_moe_experts = None
     config.expert_model_parallel_size = 1
@@ -134,19 +134,12 @@ def get_vision_model_config(hf_config, megatron_config=None):
     ):
         config.cuda_graph_impl = megatron_config.vision_cuda_graph_impl
         if hasattr(megatron_config, "vision_cuda_graph_scope") and megatron_config.vision_cuda_graph_scope:
-            # Convert string scope list to CudaGraphScope enums if needed
-            from megatron.core.transformer.cuda_graphs import CudaGraphScope
-
-            scope_list = megatron_config.vision_cuda_graph_scope
-            if scope_list and isinstance(scope_list[0], str):
-                config.cuda_graph_scope = [CudaGraphScope[scope] for scope in scope_list]
-            else:
-                config.cuda_graph_scope = scope_list
+            set_cuda_graph_modules(config, megatron_config.vision_cuda_graph_scope)
         else:
-            config.cuda_graph_scope = []
+            clear_cuda_graph_modules(config)
     else:
         config.cuda_graph_impl = "none"
-        config.cuda_graph_scope = []
+        clear_cuda_graph_modules(config)
     # Propagate max vision CUDA graph sequence length from provider
     if megatron_config is not None and hasattr(megatron_config, "max_vision_cuda_graph_seq_length"):
         config.max_vision_cuda_graph_seq_length = megatron_config.max_vision_cuda_graph_seq_length

src/megatron/bridge/training/comm_overlap.py

@@ -18,12 +18,12 @@
 
 from megatron.core.distributed import DistributedDataParallelConfig
 from megatron.core.optimizer import OptimizerConfig
-from megatron.core.transformer.enums import CudaGraphScope
 from megatron.core.utils import get_te_version, is_te_min_version, is_torch_min_version
 
 from megatron.bridge.models import GPTModelProvider, T5ModelProvider
 from megatron.bridge.models.gpt.gpt_builder import GPTModelConfig
 from megatron.bridge.models.mamba.mamba_builder import MambaModelConfig
+from megatron.bridge.utils.cuda_graph import has_cuda_graph_module
 
 
 try:
@@ -521,57 +521,9 @@ def _get_model_comm_overlap_cfgs(
                 or self.user_comm_overlap_cfg.overlap_moe_expert_parallel_comm
             ), "overlap_moe_expert_parallel_comm is required for delay_wgrad_compute"
 
-            # CUDA graph scope-specific validations for delayed wgrad.
-            cuda_graph_scope = getattr(model_cfg, "cuda_graph_scope", []) or []
-            if isinstance(cuda_graph_scope, str):
-                cuda_graph_scope = cuda_graph_scope.split(",") if cuda_graph_scope else []
-            elif not isinstance(cuda_graph_scope, list):
-                cuda_graph_scope = [cuda_graph_scope]
-            attn_scope_enabled = (
-                CudaGraphScope.attn in cuda_graph_scope
-                or CudaGraphScope.attn.value in cuda_graph_scope
-                or f"CudaGraphScope.{CudaGraphScope.attn.value}" in cuda_graph_scope
-            )
-            moe_router_scope_enabled = (
-                CudaGraphScope.moe_router in cuda_graph_scope
-                or CudaGraphScope.moe_router.value in cuda_graph_scope
-                or f"CudaGraphScope.{CudaGraphScope.moe_router.value}" in cuda_graph_scope
-            )
-            wgrad_in_graph_scope = attn_scope_enabled or (
-                moe_router_scope_enabled
-                and getattr(model_cfg, "moe_shared_expert_intermediate_size", None) is not None
-                and not getattr(model_cfg, "moe_shared_expert_overlap", False)
-            )
-            if wgrad_in_graph_scope:
-                assert is_te_min_version("2.12.0"), (
-                    "CUDA graph with delay_wgrad_compute requires TE version >= 2.12.0."
-                )
-                assert model_cfg.gradient_accumulation_fusion, (
-                    "CUDA graph with delay_wgrad_compute requires gradient_accumulation_fusion "
-                    "to be enabled. This is because default gradient accumulation does not use "
-                    "static memory addresses, which breaks CUDA graph requirements."
-                )
-                if attn_scope_enabled:
-                    assert not model_cfg.add_bias_linear and not model_cfg.add_qkv_bias, (
-                        "CUDA graph with delay_wgrad_compute does not support attention bias for now."
-                    )
-
-            # CUDA graph scope-specific validations for delayed wgrad.
-            cuda_graph_scope = getattr(model_cfg, "cuda_graph_scope", None)
-            if cuda_graph_scope is None or cuda_graph_scope == "full":
-                cuda_graph_scope = []
-            elif isinstance(cuda_graph_scope, (str, CudaGraphScope)):
-                cuda_graph_scope = [cuda_graph_scope]
-            attn_scope_enabled = (
-                CudaGraphScope.attn in cuda_graph_scope
-                or CudaGraphScope.attn.value in cuda_graph_scope
-                or f"CudaGraphScope.{CudaGraphScope.attn.value}" in cuda_graph_scope
-            )
-            moe_router_scope_enabled = (
-                CudaGraphScope.moe_router in cuda_graph_scope
-                or CudaGraphScope.moe_router.value in cuda_graph_scope
-                or f"CudaGraphScope.{CudaGraphScope.moe_router.value}" in cuda_graph_scope
-            )
+            # CUDA graph module-specific validations for delayed wgrad.
+            attn_scope_enabled = has_cuda_graph_module(model_cfg, "attn")
+            moe_router_scope_enabled = has_cuda_graph_module(model_cfg, "moe_router")
             wgrad_in_graph_scope = attn_scope_enabled or (
                 moe_router_scope_enabled
                 and getattr(model_cfg, "moe_shared_expert_intermediate_size", None) is not None

src/megatron/bridge/training/config.py

@@ -29,7 +29,7 @@
     ParamKey,
 )
 from megatron.core.process_groups_config import ProcessGroupCollection
-from megatron.core.transformer.enums import AttnBackend, CudaGraphScope
+from megatron.core.transformer.enums import AttnBackend
 from megatron.core.transformer.module import MegatronModule
 from megatron.core.transformer.transformer_config import MLATransformerConfig as MCoreMLATransformerConfig
 from megatron.core.transformer.transformer_config import TransformerConfig as MCoreTransformerConfig
@@ -61,6 +61,7 @@
     print_rank_0,
     warn_rank_0,
 )
+from megatron.bridge.utils.cuda_graph import clear_cuda_graph_modules, is_full_iteration_cuda_graph
 
 
 @dataclass
@@ -1207,13 +1208,13 @@ def validate(self) -> None:
         _validate_fine_grained_activation_offloading(self)
 
         # 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:
+        if is_full_iteration_cuda_graph(self.model):
             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 = []
+            clear_cuda_graph_modules(self.model)
 
         # ModelOpt/Quantization checks
         if getattr(self.model, "restore_modelopt_state", False):

src/megatron/bridge/training/eval.py

@@ -25,7 +25,6 @@
 from megatron.core.process_groups_config import MultiModuleProcessGroupCollection, ProcessGroupCollection
 from megatron.core.rerun_state_machine import RerunDataIterator, RerunMode, get_rerun_state_machine
 from megatron.core.transformer import MegatronModule
-from megatron.core.transformer.enums import CudaGraphScope
 from megatron.core.utils import get_model_config
 from modelopt.torch.distill.plugins.megatron import get_tensor_shapes_adjust_fn_for_distillation
 
@@ -40,6 +39,7 @@
 from megatron.bridge.training.utils.pg_utils import get_pg_collection
 from megatron.bridge.training.utils.train_utils import prepare_forward_step_func
 from megatron.bridge.utils.common_utils import is_last_rank, print_rank_0, print_rank_last
+from megatron.bridge.utils.cuda_graph import is_full_iteration_cuda_graph
 
 
 # For Paged Stashing support
@@ -156,10 +156,7 @@ def evaluate(
             )
 
             forward_backward_func = forward_backward_pipelining_without_interleaving
-        elif (
-            state.cfg.model.cuda_graph_impl == "local"
-            and CudaGraphScope.full_iteration in state.cfg.model.cuda_graph_scope
-        ):
+        elif is_full_iteration_cuda_graph(state.cfg.model):
             forward_backward_func = FullCudaGraphWrapper(
                 get_forward_backward_func(
                     pp_size=pg_collection.pp.size(),
@@ -243,10 +240,7 @@ def evaluate(
             fault_tolerance.on_eval_step_end(state)
 
             # Workaround: for FullIteration CG only. TODO: Filed #2569 to fix this.
-            if (
-                state.cfg.model.cuda_graph_impl == "local"
-                and CudaGraphScope.full_iteration in state.cfg.model.cuda_graph_scope
-            ):
+            if is_full_iteration_cuda_graph(state.cfg.model):
                 torch.cuda.synchronize()
 
             if should_fire(callback_manager, step_end_event):

src/megatron/bridge/training/train.py

@@ -54,7 +54,6 @@
     VisionTECudaGraphHelper,
     get_vision_cuda_graph_seq_length,
 )
-from megatron.core.transformer.enums import CudaGraphScope
 from megatron.core.utils import (
     check_param_hashes_across_dp_replicas,
     get_attr_wrapped_model,
@@ -98,6 +97,7 @@
     training_log,
 )
 from megatron.bridge.utils.common_utils import get_world_size_safe, print_rank_0
+from megatron.bridge.utils.cuda_graph import is_full_iteration_cuda_graph
 
 
 # For Paged Stashing support
@@ -299,7 +299,7 @@ def train(
         pp_size=pg_collection.pp.size(),
         vp_size=config.model.virtual_pipeline_model_parallel_size,
     )
-    if config.model.cuda_graph_impl == "local" and CudaGraphScope.full_iteration in config.model.cuda_graph_scope:
+    if is_full_iteration_cuda_graph(config.model):
         forward_backward_func = FullCudaGraphWrapper(
             forward_backward_func, cuda_graph_warmup_steps=config.model.cuda_graph_warmup_steps
         )

src/megatron/bridge/utils/cuda_graph.py

@@ -0,0 +1,132 @@
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Any
+
+from megatron.core.transformer.enums import CudaGraphScope
+
+
+try:
+    from megatron.core.transformer.enums import CudaGraphModule
+except ImportError:
+    CudaGraphModule = None
+
+
+def _as_list(value: Any) -> list[Any]:
+    if value is None:
+        return []
+    if isinstance(value, str):
+        if not value or value == "full":
+            return []
+        return value.split(",")
+    if isinstance(value, list):
+        return value
+    return [value]
+
+
+def _member_name(value: Any) -> str:
+    if isinstance(value, str):
+        return value.rsplit(".", 1)[-1]
+    name = getattr(value, "name", None)
+    if isinstance(name, str):
+        return name
+    return str(value).rsplit(".", 1)[-1]
+
+
+def _member_name_list(value: Any) -> list[str]:
+    return [_member_name(item) for item in _as_list(value)]
+
+
+def _member_names(value: Any) -> set[str]:
+    return set(_member_name_list(value))
+
+
+def _supports_cuda_graph_modules(config: Any) -> bool:
+    return CudaGraphModule is not None and hasattr(config, "cuda_graph_modules")
+
+
+def _module_value(name: str):
+    if CudaGraphModule is not None:
+        return CudaGraphModule[name]
+    return CudaGraphScope[name]
+
+
+def cuda_graph_module_names(config: Any) -> list[str]:
+    """Return configured per-layer CUDA graph module names."""
+
+    if getattr(config, "cuda_graph_modules", None) is not None:
+        return _member_name_list(getattr(config, "cuda_graph_modules"))
+    names = _member_name_list(getattr(config, "cuda_graph_scope", None))
+    return [name for name in names if name not in ("full_iteration", "full_iteration_inference")]
+
+
+def set_cuda_graph_modules(config: Any, modules: Any) -> None:
+    """Set per-layer CUDA graph modules using the current MCore API when available."""
+
+    module_names = _member_name_list(modules)
+    if _supports_cuda_graph_modules(config):
+        config.cuda_graph_modules = [_module_value(name) for name in module_names]
+        if hasattr(config, "cuda_graph_scope"):
+            config.cuda_graph_scope = None
+    else:
+        config.cuda_graph_scope = [CudaGraphScope[name] for name in module_names]
+
+
+def clear_cuda_graph_modules(config: Any) -> None:
+    """Clear per-layer CUDA graph modules using the active MCore API."""
+
+    set_cuda_graph_modules(config, [])
+
+
+def set_full_iteration_cuda_graph(config: Any) -> None:
+    """Enable full-iteration CUDA graph capture using the current MCore API."""
+
+    if _supports_cuda_graph_modules(config):
+        config.cuda_graph_impl = "full_iteration"
+        config.cuda_graph_modules = []
+        if hasattr(config, "cuda_graph_scope"):
+            config.cuda_graph_scope = None
+    else:
+        config.cuda_graph_impl = "local"
+        config.cuda_graph_scope = [CudaGraphScope.full_iteration]
+
+
+def has_cuda_graph_module(config: Any, module: Any) -> bool:
+    """Return whether a per-layer CUDA graph module is enabled.
+
+    Supports both the current MCore ``cuda_graph_modules`` API and the deprecated
+    ``cuda_graph_scope`` values still present in older Bridge configs.
+    """
+
+    module_name = _member_name(module)
+    module_names = _member_names(getattr(config, "cuda_graph_modules", None))
+    legacy_scope_names = _member_names(getattr(config, "cuda_graph_scope", None))
+    return module_name in module_names or module_name in legacy_scope_names
+
+
+def is_full_iteration_cuda_graph(config: Any) -> bool:
+    """Return whether config enables full-iteration CUDA graph capture."""
+
+    cuda_graph_impl = getattr(config, "cuda_graph_impl", "none")
+    if cuda_graph_impl == "full_iteration":
+        return True
+    if cuda_graph_impl != "local":
+        return False
+    return "full_iteration" in _member_names(getattr(config, "cuda_graph_scope", None))
+
+
+def uses_local_cuda_graph_manager(config: Any) -> bool:
+    """Return whether Bridge should create a local MCore CudaGraphManager."""
+
+    return getattr(config, "cuda_graph_impl", "none") == "local" and not is_full_iteration_cuda_graph(config)