From 6815c35c5037dcbbcf1a9090011cc54e95cee5e5 Mon Sep 17 00:00:00 2001 From: Robin Zhang Date: Mon, 27 Apr 2026 19:55:27 -0700 Subject: [PATCH] Support old and new MCore CUDA graph APIs Signed-off-by: Robin Zhang --- .../models/gpt_full_te_layer_autocast_spec.py | 9 +- .../modelling_qwen3_vl/transformer_config.py | 17 +-- src/megatron/bridge/training/comm_overlap.py | 56 +------- src/megatron/bridge/training/config.py | 7 +- src/megatron/bridge/training/eval.py | 12 +- src/megatron/bridge/training/train.py | 4 +- src/megatron/bridge/utils/cuda_graph.py | 132 ++++++++++++++++++ .../test_qwen3_vl_transformer_config.py | 45 ++++-- .../test_gpt_full_te_layer_autocast_spec.py | 19 ++- .../unit_tests/training/test_comm_overlap.py | 8 +- tests/unit_tests/training/test_config.py | 36 +++-- 11 files changed, 225 insertions(+), 120 deletions(-) create mode 100644 src/megatron/bridge/utils/cuda_graph.py diff --git a/src/megatron/bridge/models/gpt_full_te_layer_autocast_spec.py b/src/megatron/bridge/models/gpt_full_te_layer_autocast_spec.py index 83724f25b9..b9281f2558 100644 --- a/src/megatron/bridge/models/gpt_full_te_layer_autocast_spec.py +++ b/src/megatron/bridge/models/gpt_full_te_layer_autocast_spec.py @@ -20,7 +20,6 @@ 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 @@ -28,6 +27,8 @@ 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)) diff --git a/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/transformer_config.py b/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/transformer_config.py index 7b43fb58c2..5133dd5544 100644 --- a/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/transformer_config.py +++ b/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 diff --git a/src/megatron/bridge/training/comm_overlap.py b/src/megatron/bridge/training/comm_overlap.py index 5958eb743a..47ac1e69c9 100644 --- a/src/megatron/bridge/training/comm_overlap.py +++ b/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 diff --git a/src/megatron/bridge/training/config.py b/src/megatron/bridge/training/config.py index 31afdc64af..6217f282c0 100644 --- a/src/megatron/bridge/training/config.py +++ b/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): diff --git a/src/megatron/bridge/training/eval.py b/src/megatron/bridge/training/eval.py index 5c68a39a83..e440cb8a53 100644 --- a/src/megatron/bridge/training/eval.py +++ b/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): diff --git a/src/megatron/bridge/training/train.py b/src/megatron/bridge/training/train.py index 58197e00e8..f16f3a59e3 100644 --- a/src/megatron/bridge/training/train.py +++ b/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 ) diff --git a/src/megatron/bridge/utils/cuda_graph.py b/src/megatron/bridge/utils/cuda_graph.py new file mode 100644 index 0000000000..f96493892d --- /dev/null +++ b/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) diff --git a/tests/unit_tests/models/qwen_vl/modelling_qwen3_vl/test_qwen3_vl_transformer_config.py b/tests/unit_tests/models/qwen_vl/modelling_qwen3_vl/test_qwen3_vl_transformer_config.py index 1e39657178..ae81954aad 100644 --- a/tests/unit_tests/models/qwen_vl/modelling_qwen3_vl/test_qwen3_vl_transformer_config.py +++ b/tests/unit_tests/models/qwen_vl/modelling_qwen3_vl/test_qwen3_vl_transformer_config.py @@ -14,12 +14,13 @@ """Unit tests for Qwen3-VL vision transformer_config (get_vision_model_config).""" +from enum import Enum from types import SimpleNamespace import pytest -from megatron.core.transformer.cuda_graphs import CudaGraphScope from megatron.bridge.models.qwen_vl.modelling_qwen3_vl.transformer_config import get_vision_model_config +from megatron.bridge.utils.cuda_graph import cuda_graph_module_names def _hf_config(): @@ -57,6 +58,14 @@ def _megatron_base(**overrides): return SimpleNamespace(**base) +def _cuda_graph_storage(config): + """Return the concrete CUDA graph field used by this MCore version.""" + modules = getattr(config, "cuda_graph_modules", None) + if modules is not None: + return modules + return getattr(config, "cuda_graph_scope", []) + + class TestGetVisionModelConfigVisionCudaGraph: """Vision encoder CUDA graph propagation from megatron_config (provider).""" @@ -64,44 +73,56 @@ def test_vision_cuda_graph_defaults_when_impl_none(self): megatron = _megatron_base( vision_cuda_graph_impl="none", cuda_graph_impl="local_transformer_engine", - cuda_graph_scope=[CudaGraphScope.attn], + cuda_graph_scope=["attn"], ) cfg = get_vision_model_config(_hf_config(), megatron) assert cfg.cuda_graph_impl == "none" - assert cfg.cuda_graph_scope == [] + assert cuda_graph_module_names(cfg) == [] def test_vision_cuda_graph_defaults_when_attr_missing(self): megatron = _megatron_base( cuda_graph_impl="local_transformer_engine", - cuda_graph_scope=[CudaGraphScope.attn], + cuda_graph_scope=["attn"], + ) + cfg = get_vision_model_config(_hf_config(), megatron) + assert cfg.cuda_graph_impl == "none" + assert cuda_graph_module_names(cfg) == [] + + def test_vision_cuda_graph_ignores_language_legacy_full_iteration_when_disabled(self): + megatron = _megatron_base( + cuda_graph_impl="local", + cuda_graph_scope=["full_iteration"], ) cfg = get_vision_model_config(_hf_config(), megatron) assert cfg.cuda_graph_impl == "none" - assert cfg.cuda_graph_scope == [] + assert cuda_graph_module_names(cfg) == [] def test_vision_cuda_graph_impl_propagated_scope_empty_without_scope_attr(self): megatron = _megatron_base(vision_cuda_graph_impl="local_transformer_engine") cfg = get_vision_model_config(_hf_config(), megatron) assert cfg.cuda_graph_impl == "local_transformer_engine" - assert cfg.cuda_graph_scope == [] + assert cuda_graph_module_names(cfg) == [] - def test_vision_cuda_graph_scope_string_list_converted_to_enum(self): + def test_vision_cuda_graph_scope_string_list_converted_to_enums(self): megatron = _megatron_base( vision_cuda_graph_impl="local_transformer_engine", vision_cuda_graph_scope=["attn", "mlp"], ) cfg = get_vision_model_config(_hf_config(), megatron) + cuda_graph_values = _cuda_graph_storage(cfg) + assert cfg.cuda_graph_impl == "local_transformer_engine" - assert cfg.cuda_graph_scope == [CudaGraphScope.attn, CudaGraphScope.mlp] + assert cuda_graph_module_names(cfg) == ["attn", "mlp"] + assert all(isinstance(value, Enum) for value in cuda_graph_values) - def test_vision_cuda_graph_scope_enum_list_passed_through(self): - scopes = [CudaGraphScope.attn] + def test_vision_cuda_graph_scope_list_propagated(self): + scopes = ["attn"] megatron = _megatron_base( vision_cuda_graph_impl="local_transformer_engine", vision_cuda_graph_scope=scopes, ) cfg = get_vision_model_config(_hf_config(), megatron) - assert cfg.cuda_graph_scope is scopes + assert cuda_graph_module_names(cfg) == scopes def test_vision_cuda_graph_scope_empty_list_clears_scope(self): megatron = _megatron_base( @@ -109,7 +130,7 @@ def test_vision_cuda_graph_scope_empty_list_clears_scope(self): vision_cuda_graph_scope=[], ) cfg = get_vision_model_config(_hf_config(), megatron) - assert cfg.cuda_graph_scope == [] + assert cuda_graph_module_names(cfg) == [] def test_max_vision_cuda_graph_seq_length_propagated(self): megatron = _megatron_base( diff --git a/tests/unit_tests/models/test_gpt_full_te_layer_autocast_spec.py b/tests/unit_tests/models/test_gpt_full_te_layer_autocast_spec.py index ae031667d1..602274cfeb 100644 --- a/tests/unit_tests/models/test_gpt_full_te_layer_autocast_spec.py +++ b/tests/unit_tests/models/test_gpt_full_te_layer_autocast_spec.py @@ -23,7 +23,6 @@ import pytest import torch -from megatron.core.transformer.enums import CudaGraphScope from megatron.bridge.models.gpt_full_te_layer_autocast_spec import ( AutocastTransformerLayer, @@ -31,6 +30,7 @@ get_gpt_full_te_layer_autocast_spec, torch_dtype_from_precision, ) +from megatron.bridge.utils.cuda_graph import set_cuda_graph_modules, set_full_iteration_cuda_graph class TestTorchDtypeFromPrecision: @@ -214,7 +214,7 @@ def mock_config(self): config.bf16 = False config.num_layers = 12 config.cuda_graph_impl = "none" - config.cuda_graph_scope = [] + set_cuda_graph_modules(config, []) config.cpu_offloading = False config.recompute_granularity = None config.virtual_pipeline_model_parallel_size = None @@ -305,7 +305,7 @@ def rank(self): mock_config._pg_collection = type("PGC", (), {"pp": _PG()})() mock_config.cuda_graph_impl = "local" - mock_config.cuda_graph_scope = [] # Empty list means layerwise graph + set_cuda_graph_modules(mock_config, []) with patch("megatron.bridge.models.gpt_full_te_layer_autocast_spec.AutocastTransformerLayer"): with patch("megatron.bridge.models.gpt_full_te_layer_autocast_spec.CudaGraphManager") as mock_cuda_manager: @@ -324,8 +324,8 @@ def __init__(self, config): def test_te_transformer_layer_autocast_with_full_iteration_cuda_graph(self, mock_config): """Test TETransformerLayerAutocast with full_iteration CUDA graph (cudagraph_manager should NOT be created). - Note: MCore's TransformerConfig.__post_init__ converts string scope values to CudaGraphScope enums. - This test uses the enum directly to match the runtime behavior after config finalization. + MCore represents full-iteration graph capture with either the new + cuda_graph_impl="full_iteration" form or the old cuda_graph_scope value. """ mock_config.tensor_model_parallel_size = 1 mock_config.pipeline_model_parallel_size = 1 @@ -335,8 +335,7 @@ def rank(self): return 0 mock_config._pg_collection = type("PGC", (), {"pp": _PG()})() - mock_config.cuda_graph_impl = "local" - mock_config.cuda_graph_scope = [CudaGraphScope.full_iteration] # Full iteration graph (enum) + set_full_iteration_cuda_graph(mock_config) with patch("megatron.bridge.models.gpt_full_te_layer_autocast_spec.AutocastTransformerLayer"): with patch("megatron.bridge.models.gpt_full_te_layer_autocast_spec.CudaGraphManager") as mock_cuda_manager: @@ -349,8 +348,8 @@ def rank(self): def test_te_transformer_layer_autocast_external_cuda_graph(self, mock_config): """Test TETransformerLayerAutocast with external CUDA graph. - Note: MCore's TransformerConfig.__post_init__ converts string scope values to CudaGraphScope enums. - This test uses the enum directly to match the runtime behavior after config finalization. + MCore represents TE graph scopes as either cuda_graph_modules or + cuda_graph_scope depending on version. """ mock_config.tensor_model_parallel_size = 1 mock_config.pipeline_model_parallel_size = 1 @@ -361,7 +360,7 @@ def rank(self): mock_config._pg_collection = type("PGC", (), {"pp": _PG()})() mock_config.cuda_graph_impl = "transformer_engine" - mock_config.cuda_graph_scope = [CudaGraphScope.attn, CudaGraphScope.mlp] # TE supports multi-scope (enum) + set_cuda_graph_modules(mock_config, ["attn", "mlp"]) with patch("megatron.bridge.models.gpt_full_te_layer_autocast_spec.AutocastTransformerLayer") as mock_autocast: mock_transformer = Mock() diff --git a/tests/unit_tests/training/test_comm_overlap.py b/tests/unit_tests/training/test_comm_overlap.py index 7e01862f9d..c2350e5981 100644 --- a/tests/unit_tests/training/test_comm_overlap.py +++ b/tests/unit_tests/training/test_comm_overlap.py @@ -15,7 +15,6 @@ from unittest.mock import MagicMock, patch import pytest -from megatron.core.transformer.enums import CudaGraphScope from megatron.bridge.models.gpt.gpt_builder import GPTModelConfig from megatron.bridge.models.gpt_provider import GPTModelProvider @@ -28,6 +27,7 @@ userbuffers_bf16_h100_h8192_tp4_mbs1_seqlen8192, ) from megatron.bridge.training.config import DistributedDataParallelConfig, OptimizerConfig +from megatron.bridge.utils.cuda_graph import set_cuda_graph_modules def create_gpt_config(**kwargs): @@ -654,8 +654,8 @@ def test_delay_wgrad_cuda_graph_attn_requires_grad_accum_fusion(self): add_bias_linear=False, add_qkv_bias=False, gradient_accumulation_fusion=False, - cuda_graph_scope=[CudaGraphScope.attn], ) + set_cuda_graph_modules(model_cfg, ["attn"]) ddp_cfg = DistributedDataParallelConfig(use_distributed_optimizer=False) with ( @@ -687,8 +687,8 @@ def test_delay_wgrad_cuda_graph_attn_rejects_attention_bias(self): add_bias_linear=True, add_qkv_bias=False, gradient_accumulation_fusion=True, - cuda_graph_scope=[CudaGraphScope.attn], ) + set_cuda_graph_modules(model_cfg, ["attn"]) ddp_cfg = DistributedDataParallelConfig(use_distributed_optimizer=False) with ( @@ -720,8 +720,8 @@ def test_delay_wgrad_cuda_graph_attn_validation_passes_with_supported_settings(s add_bias_linear=False, add_qkv_bias=False, gradient_accumulation_fusion=True, - cuda_graph_scope=[CudaGraphScope.attn], ) + set_cuda_graph_modules(model_cfg, ["attn"]) ddp_cfg = DistributedDataParallelConfig(use_distributed_optimizer=False) with ( diff --git a/tests/unit_tests/training/test_config.py b/tests/unit_tests/training/test_config.py index 71a3e6f979..564c10fd38 100644 --- a/tests/unit_tests/training/test_config.py +++ b/tests/unit_tests/training/test_config.py @@ -17,7 +17,6 @@ import pytest import torch -from megatron.core.transformer.enums import CudaGraphScope from megatron.bridge.models.gpt_provider import GPTModelProvider from megatron.bridge.models.mla_provider import MLAModelProvider @@ -45,6 +44,11 @@ _validate_and_sync_distributed_optimizer_settings, _validate_mixed_precision_consistency, ) +from megatron.bridge.utils.cuda_graph import ( + cuda_graph_module_names, + set_cuda_graph_modules, + set_full_iteration_cuda_graph, +) def mock_get_world_size_safe(world_size_to_return: int): @@ -1216,11 +1220,10 @@ def test_megatron_fsdp_config_with_dp_last_dim(self, monkeypatch): def test_cuda_graph_full_iteration_requires_check_for_nan_disabled(self, monkeypatch): """Test that full_iteration CUDA graph requires check_for_nan_in_loss=False.""" - # Create config with cuda_graph_impl="local" and TE RNG tracker (required for cuda graphs) gpt_model_cfg = create_test_gpt_config( - cuda_graph_impl="local", use_te_rng_tracker=True, ) + set_full_iteration_cuda_graph(gpt_model_cfg) container, og_ws, cfg_mod = create_test_config_container( world_size_override=1, @@ -1228,10 +1231,6 @@ def test_cuda_graph_full_iteration_requires_check_for_nan_disabled(self, monkeyp ) try: - # Set cuda_graph_scope to include full_iteration after model creation - # (MCore's __post_init__ converts strings to enums during finalize) - container.model.cuda_graph_scope = [CudaGraphScope.full_iteration] - # Default check_for_nan_in_loss is True - should fail validation assert container.rerun_state_machine.check_for_nan_in_loss is True with pytest.raises( @@ -1252,6 +1251,7 @@ def test_cuda_graph_non_full_iteration_allows_check_for_nan(self, monkeypatch): cuda_graph_impl="local", use_te_rng_tracker=True, ) + set_cuda_graph_modules(gpt_model_cfg, ["attn", "mlp"]) container, og_ws, cfg_mod = create_test_config_container( world_size_override=1, @@ -1259,15 +1259,31 @@ def test_cuda_graph_non_full_iteration_allows_check_for_nan(self, monkeypatch): ) try: - # Set cuda_graph_scope to NOT include full_iteration - container.model.cuda_graph_scope = [CudaGraphScope.attn, CudaGraphScope.mlp] - # check_for_nan_in_loss=True should be allowed assert container.rerun_state_machine.check_for_nan_in_loss is True container.validate() # Should pass without error finally: restore_get_world_size_safe(og_ws, cfg_mod) + def test_cuda_graph_impl_none_clears_modules(self, monkeypatch): + """Test that cuda_graph_impl=none clears module-scoped CUDA graph settings.""" + gpt_model_cfg = create_test_gpt_config( + cuda_graph_impl="none", + use_te_rng_tracker=True, + ) + set_cuda_graph_modules(gpt_model_cfg, ["attn", "mlp"]) + + container, og_ws, cfg_mod = create_test_config_container( + world_size_override=1, + model_config=gpt_model_cfg, + ) + + try: + container.validate() + assert cuda_graph_module_names(container.model) == [] + finally: + restore_get_world_size_safe(og_ws, cfg_mod) + @pytest.mark.parametrize("model_factory", [create_test_gpt_config, create_test_deepseek_config]) def test_default_pipeline_dtype(self, model_factory, monkeypatch): """