diff --git a/src/megatron/bridge/perf_recipes/nemotronh/common.py b/src/megatron/bridge/perf_recipes/nemotronh/common.py index 7a94a4b572..fb90df2363 100644 --- a/src/megatron/bridge/perf_recipes/nemotronh/common.py +++ b/src/megatron/bridge/perf_recipes/nemotronh/common.py @@ -206,6 +206,27 @@ def _apply_nemotron_3_ultra_perf_defaults(cfg: ConfigContainer) -> None: cfg.ddp.num_buckets = 48 +def _apply_nemotron_3_ultra_gtp(cfg: ConfigContainer) -> None: + """Apply rack-local dense GTP settings for Nemotron 3 Ultra on GB300. + + Dense weights are sharded across one 64-GPU NVLink domain. Routed-expert + weights remain unsharded by GTP because EP64 already spans that domain; + this avoids expert-weight all-gathers across racks. + """ + cfg.dist.use_megatron_fsdp = False + cfg.ddp.use_megatron_fsdp = False + + # GTP is controlled through the user-facing total weight-shard counts. + # TP1 x GTP64 gives one dense-weight shard per GPU within a GB300 rack. + cfg.model.tensor_parallel_num_weight_shards = _GB300_NVLINK_DOMAIN_GPUS + cfg.model.expert_tensor_parallel_num_weight_shards = cfg.model.expert_tensor_parallel_size + + # GTP currently supports only one distributed-optimizer instance. + cfg.ddp.num_distributed_optimizer_instances = 1 + cfg.ddp.data_parallel_sharding_strategy = "no_shard" + cfg.ddp.outer_dp_sharding_strategy = "no_shard" + + def _apply_nemotron_3_ultra_fsdp_hsdp(cfg: ConfigContainer, num_gpus: int) -> None: """Apply Megatron-FSDP (HSDP) settings for Nemotron 3 Ultra on GB300. diff --git a/src/megatron/bridge/perf_recipes/nemotronh/gb200/nemotronh.py b/src/megatron/bridge/perf_recipes/nemotronh/gb200/nemotronh.py index b690e0bd74..3261eb50c7 100644 --- a/src/megatron/bridge/perf_recipes/nemotronh/gb200/nemotronh.py +++ b/src/megatron/bridge/perf_recipes/nemotronh/gb200/nemotronh.py @@ -13,6 +13,7 @@ # limitations under the License. """GB200 performance recipes for Nemotron 3 and 3.5.""" +from megatron.bridge.perf_recipes._common import _enable_ncclep from megatron.bridge.perf_recipes.environment import COMMON_PERF_ENV_VARS from megatron.bridge.perf_recipes.nemotronh.common import ( _TE_QUANT_CFG_PATH, @@ -20,6 +21,7 @@ _apply_nemotron_3_nano_perf_defaults, _apply_nemotron_3_super_perf_defaults, _apply_nemotron_3_ultra_fsdp_hsdp, + _apply_nemotron_3_ultra_gtp, _apply_nemotron_3_ultra_perf_defaults, _benchmark_common, _enable_nemotron_3_super_full_iteration, @@ -185,10 +187,10 @@ def nemotron_3_super_pretrain_64gpu_gb200_nvfp4_config() -> ConfigContainer: def nemotron_3_ultra_pretrain_256gpu_gb200_fp8mx_config(*, num_gpus: int = 256) -> ConfigContainer: - """Nemotron 3 Ultra (550B-A55B LatentMoE) pretrain: 256× GB200, MXFP8, Megatron-FSDP (HSDP). + """Nemotron 3 Ultra (550B-A55B LatentMoE) pretrain: 256× GB200, MXFP8, GTP. TP2 + SP (due to smaller GB200 HBM) / PP1 / CP1 / EP64 / ETP1, GBS 256 / MBS 1, seq 8192, BF16 + MXFP8 mixed - precision, HybridEP flex dispatcher, CuteDSL fused grouped MLP, selective recompute + + precision, dense GTP32, NCCLEP flex dispatcher, CuteDSL fused grouped MLP, selective recompute + fine-grained activation offload of the expert MLP, MTP=2. Args: @@ -198,15 +200,17 @@ def nemotron_3_ultra_pretrain_256gpu_gb200_fp8mx_config(*, num_gpus: int = 256) expert_model_parallel_size = 64 global_batch_size = 256 - hybrid_ep_ranks_per_nvlink_domain = 64 cfg = nemotron_3_ultra_pretrain_config() cfg.mixed_precision = _perf_precision("fp8_mx") _apply_nemotron_3_ultra_perf_defaults(cfg) - # Apply HSDP / FSDP dtype overrides last so they win over the generic defaults. - _apply_nemotron_3_ultra_fsdp_hsdp(cfg, num_gpus=num_gpus) + # Shard dense weights with GTP within each 64-GPU rack. GTP currently + # requires a single distributed-optimizer instance. + _apply_nemotron_3_ultra_gtp(cfg) + # Also, enable GTP sharding for the moe_latent_proj module due to memory pressure on GB200 + cfg.model.gtp_remat_opt_in_modules = ["moe_latent_proj"] # Parallelism cfg.model.tensor_model_parallel_size = 2 @@ -241,6 +245,20 @@ def nemotron_3_ultra_pretrain_256gpu_gb200_fp8mx_config(*, num_gpus: int = 256) cfg.model.recompute_granularity = "selective" cfg.model.recompute_modules = ["moe_act"] + # Enable GTP sharding for the moe_latent_proj module + cfg.model.gtp_remat_opt_in_modules = ["moe_latent_proj"] + + cfg.model.cuda_graph_impl = "local" + set_cuda_graph_modules(cfg.model, ["attn", "mamba"]) + # TODO: "moe_preprocess", "moe_router" needs to be disabled due to MCore PR 7046, performing a CPU-to-GPU copy for per token loss + # TE attention requires TE's RNG tracker while MCore local graphs are capturing. + cfg.rng.te_rng_tracker = cfg.model.use_te_rng_tracker = True + + _enable_ncclep(cfg) + # Device-side expert token counts: the legacy grouped MLP path syncs tokens_per_expert to the + # host every layer, which serializes the CPU behind the GPU when dispatch is fast. + cfg.model.moe_use_grouped_tensor = True + # Keep process settings next to the recipe so users can see the exact benchmark environment. cfg.env_vars = { **COMMON_PERF_ENV_VARS, @@ -255,14 +273,14 @@ def nemotron_3_ultra_pretrain_256gpu_gb200_fp8mx_config(*, num_gpus: int = 256) "TORCH_NCCL_AVOID_RECORD_STREAMS": 1, # NCCL user-buffer and launch settings. "NCCL_NVLS_ENABLE": 0, - # HybridEP topology for the target system. - "NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN": hybrid_ep_ranks_per_nvlink_domain, - "NUM_OF_TOKENS_PER_CHUNK_COMBINE_API": 128, - "NVLINK_DOMAIN_SIZE": 72, - "USE_MNNVL": 1, + # NCCL EP dispatcher mode and one GPU per rank. + "NCCL_EP_HT_EM_PULL_PUSH": 1, # Transformer Engine overlap settings for this model. "NVTE_BWD_LAYERNORM_SM_MARGIN": 20, "NVTE_FWD_LAYERNORM_SM_MARGIN": 20, + # Use cuDNN normalization kernels for the VR200 performance candidate. + "NVTE_NORM_BWD_USE_CUDNN": 1, + "NVTE_NORM_FWD_USE_CUDNN": 1, # Required by fine_grained_activation_offloading (TE >= 2.10.0) to avoid # offloading weights; "NVTE_CPU_OFFLOAD_V1": 1, diff --git a/src/megatron/bridge/perf_recipes/nemotronh/gb300/nemotronh.py b/src/megatron/bridge/perf_recipes/nemotronh/gb300/nemotronh.py index 481814c27f..695658c629 100644 --- a/src/megatron/bridge/perf_recipes/nemotronh/gb300/nemotronh.py +++ b/src/megatron/bridge/perf_recipes/nemotronh/gb300/nemotronh.py @@ -21,6 +21,7 @@ _apply_nemotron_3_nano_perf_defaults, _apply_nemotron_3_super_perf_defaults, _apply_nemotron_3_ultra_fsdp_hsdp, + _apply_nemotron_3_ultra_gtp, _apply_nemotron_3_ultra_perf_defaults, _benchmark_common, _enable_nemotron_3_5_lightning_full_iteration, @@ -209,13 +210,15 @@ def nemotron_3_super_pretrain_64gpu_gb300_nvfp4_config() -> ConfigContainer: def _nemotron_3_ultra_gb300_fp8mx_config( *, num_gpus: int, expert_model_parallel_size: int, global_batch_size: int ) -> ConfigContainer: - """Shared builder for Nemotron 3 Ultra GB300 MXFP8 Megatron-FSDP perf recipes.""" + """Shared builder for Nemotron 3 Ultra GB300 MXFP8 GTP perf recipes.""" cfg = nemotron_3_ultra_pretrain_config() cfg.mixed_precision = _perf_precision("fp8_mx") + _apply_nemotron_3_ultra_perf_defaults(cfg) + """ Uses TP1 / PP1 / CP1 / EP64 / ETP1, GBS 256 / MBS 1, - seq 8192, HybridEP flex dispatcher, CuteDSL fused grouped MLP, selective + seq 8192, NCCLEP flex dispatcher, CuteDSL fused grouped MLP, selective recompute + activation offload of the expert MLP, MTP=2. The MoE architecture (512 experts, latent MoE, MTP, squared-relu, hybrid Mamba/attention pattern, ...) is inherited from the base recipe via ``AutoBridge``. @@ -252,20 +255,30 @@ def _nemotron_3_ultra_gb300_fp8mx_config( cfg.model.recompute_granularity = "selective" cfg.model.recompute_modules = ["moe_act"] - _apply_nemotron_3_ultra_perf_defaults(cfg) + # Shard dense weights with GTP within each 64-GPU rack. GTP currently + # requires a single distributed-optimizer instance. + _apply_nemotron_3_ultra_gtp(cfg) - # Apply HSDP / FSDP dtype overrides last so they win over the generic defaults. - _apply_nemotron_3_ultra_fsdp_hsdp(cfg, num_gpus=num_gpus) + cfg.model.cuda_graph_impl = "local" + set_cuda_graph_modules(cfg.model, ["attn", "mamba"]) + # TODO: "moe_preprocess", "moe_router" needs to be disabled due to MCore PR 7046, performing a CPU-to-GPU copy for per token loss + # TE attention requires TE's RNG tracker while MCore local graphs are capturing. + cfg.rng.te_rng_tracker = cfg.model.use_te_rng_tracker = True + + _enable_ncclep(cfg) + # Device-side expert token counts: the legacy grouped MLP path syncs tokens_per_expert to the + # host every layer, which serializes the CPU behind the GPU when dispatch is fast. + cfg.model.moe_use_grouped_tensor = True return cfg def nemotron_3_ultra_pretrain_256gpu_gb300_fp8mx_config(*, num_gpus: int = 256) -> ConfigContainer: - """Nemotron 3 Ultra (550B-A55B LatentMoE) pretrain: 256× GB300, MXFP8, Megatron-FSDP (HSDP). + """Nemotron 3 Ultra (550B-A55B LatentMoE) pretrain: 256× GB300, MXFP8, GTP. TP1 / PP1 / CP1 / EP64 / ETP1, GBS 256 / MBS 1, seq 8192, BF16 + MXFP8 mixed - precision, HybridEP flex dispatcher, CuteDSL fused grouped MLP, selective - recompute + fine-grained activation offload of the expert MLP, MTP=2. + precision, dense GTP64, NCCLEP flex dispatcher, CuteDSL fused grouped MLP, + selective recompute + fine-grained activation offload of the expert MLP, MTP=2. Args: num_gpus: Launch GPU count used to size HSDP optimizer instances. @@ -276,10 +289,6 @@ def nemotron_3_ultra_pretrain_256gpu_gb300_fp8mx_config(*, num_gpus: int = 256) expert_model_parallel_size=64, global_batch_size=256, ) - _enable_ncclep(cfg) - # Device-side expert token counts: the legacy grouped MLP path syncs tokens_per_expert to the - # host every layer, which serializes the CPU behind the GPU when dispatch is fast. - cfg.model.moe_use_grouped_tensor = True # Keep process settings next to the recipe so users can see the exact benchmark environment. cfg.env_vars = { **COMMON_PERF_ENV_VARS, @@ -296,6 +305,9 @@ def nemotron_3_ultra_pretrain_256gpu_gb300_fp8mx_config(*, num_gpus: int = 256) # Transformer Engine overlap settings for this model. "NVTE_BWD_LAYERNORM_SM_MARGIN": 20, "NVTE_FWD_LAYERNORM_SM_MARGIN": 20, + # Use cuDNN normalization kernels for the VR200 performance candidate. + "NVTE_NORM_BWD_USE_CUDNN": 1, + "NVTE_NORM_FWD_USE_CUDNN": 1, # Required by fine_grained_activation_offloading (TE >= 2.10.0) to avoid # offloading weights; "NVTE_CPU_OFFLOAD_V1": 1, diff --git a/src/megatron/bridge/perf_recipes/nemotronh/vrnvl72/nemotronh.py b/src/megatron/bridge/perf_recipes/nemotronh/vrnvl72/nemotronh.py index bfb03ff6bb..29211824af 100644 --- a/src/megatron/bridge/perf_recipes/nemotronh/vrnvl72/nemotronh.py +++ b/src/megatron/bridge/perf_recipes/nemotronh/vrnvl72/nemotronh.py @@ -240,11 +240,8 @@ def nemotron_3_ultra_pretrain_256gpu_vrnvl72_fp8mx_config(*, num_gpus: int = 256 "TORCH_NCCL_AVOID_RECORD_STREAMS": 1, # NCCL user-buffer and launch settings. "NCCL_NVLS_ENABLE": 0, - # Use 64 HybridEP ranks within the 72-GPU NVLink domain. - "NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN": 64, - "NUM_OF_TOKENS_PER_CHUNK_COMBINE_API": 128, - "NVLINK_DOMAIN_SIZE": 72, - "USE_MNNVL": 1, + # NCCL EP dispatcher mode and one GPU per rank. + "NCCL_EP_HT_EM_PULL_PUSH": 1, # Transformer Engine overlap settings for this model. "NVTE_BWD_LAYERNORM_SM_MARGIN": 20, "NVTE_FWD_LAYERNORM_SM_MARGIN": 20, diff --git a/src/megatron/bridge/training/train.py b/src/megatron/bridge/training/train.py index dae17e0620..1cec4b13aa 100644 --- a/src/megatron/bridge/training/train.py +++ b/src/megatron/bridge/training/train.py @@ -57,6 +57,9 @@ VisionTECudaGraphHelper, get_vision_cuda_graph_seq_length, ) +from megatron.core.transformer.cuda_graphs import ( + delete_cuda_graphs as delete_local_cuda_graphs, +) from megatron.core.utils import ( check_param_hashes_across_dp_replicas, get_attr_wrapped_model, @@ -1786,6 +1789,9 @@ def _delete_cuda_graphs(cuda_graph_helper: TECudaGraphHelper | None): if cuda_graph_helper is not None and cuda_graph_helper.graphs_created(): cuda_graph_helper.delete_cuda_graphs() + # Cleanup per-layer CUDA graphs owned by MCore CudaGraphManager (cuda_graph_impl="local"). + delete_local_cuda_graphs() + # Run GC to collect the freshed object gc.collect() diff --git a/src/megatron/bridge/utils/cuda_graph.py b/src/megatron/bridge/utils/cuda_graph.py index d0e2c43c47..2abe01013b 100644 --- a/src/megatron/bridge/utils/cuda_graph.py +++ b/src/megatron/bridge/utils/cuda_graph.py @@ -140,16 +140,21 @@ def validate_cuda_graph_configuration(config: Any, *, config_name: str = "model" graph_modules = cuda_graph_module_names(config) inference_scopes = _member_names(getattr(config, "inference_cuda_graph_scope", None)) + is_local_per_layer_graph = cuda_graph_impl == "local" and bool(graph_modules) is_local_inference_graph = ( cuda_graph_impl == "local" and not graph_modules and bool(inference_scopes & {"layer", "block"}) ) - if cuda_graph_impl == "local" and not is_full_iteration_cuda_graph(config) and not is_local_inference_graph: - modules_msg = f" with scopes {graph_modules}" if graph_modules else "" + if ( + cuda_graph_impl == "local" + and not is_full_iteration_cuda_graph(config) + and not is_local_per_layer_graph + and not is_local_inference_graph + ): raise ValueError( f'Megatron Bridge supports {config_name}.cuda_graph_impl="local" only for ' - f"full-iteration CUDA graphs{modules_msg}. Use " - f'{config_name}.cuda_graph_impl="transformer_engine" for scoped CUDA graph ' - 'capture such as "attn" or "mlp".' + "full-iteration CUDA graphs or per-layer CUDA graphs with explicit " + f'{config_name}.cuda_graph_modules such as ["attn", "mamba"]. Whole-layer local ' + "capture with empty cuda_graph_modules is not supported." ) diff --git a/tests/unit_tests/recipes/nemotronh/test_nemotron_3_ultra.py b/tests/unit_tests/recipes/nemotronh/test_nemotron_3_ultra.py index efe243ebdb..63e7023041 100644 --- a/tests/unit_tests/recipes/nemotronh/test_nemotron_3_ultra.py +++ b/tests/unit_tests/recipes/nemotronh/test_nemotron_3_ultra.py @@ -177,17 +177,17 @@ def test_vrnvl72_perf_recipe_uses_nvl72_ultra_topology() -> None: assert cfg.model.expert_model_parallel_size == 64 assert cfg.model.expert_tensor_parallel_size == 1 assert cfg.model.moe_token_dispatcher_type == "flex" - assert cfg.model.moe_flex_dispatcher_backend == "hybridep" - assert cfg.model.cuda_graph_impl == "none" + assert cfg.model.moe_flex_dispatcher_backend == "ncclep" + assert cfg.model.cuda_graph_impl == "local" assert cfg.train.global_batch_size == 256 assert cfg.train.micro_batch_size == 1 - assert cfg.ddp.use_megatron_fsdp is True - assert cfg.ddp.num_distributed_optimizer_instances == 4 - assert cfg.env_vars["NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN"] == 64 - assert cfg.env_vars["NVLINK_DOMAIN_SIZE"] == 72 + assert cfg.ddp.use_megatron_fsdp is False + assert cfg.model.tensor_parallel_num_weight_shards == 64 + assert cfg.model.expert_tensor_parallel_num_weight_shards == 1 + assert cfg.ddp.num_distributed_optimizer_instances == 1 + assert cfg.env_vars["NCCL_EP_HT_EM_PULL_PUSH"] == 1 assert cfg.env_vars["NVTE_NORM_BWD_USE_CUDNN"] == 1 assert cfg.env_vars["NVTE_NORM_FWD_USE_CUDNN"] == 1 - assert cfg.env_vars["USE_MNNVL"] == 1 @pytest.mark.unit diff --git a/tests/unit_tests/recipes/test_gb300_ncclep_perf_recipes.py b/tests/unit_tests/recipes/test_gb300_ncclep_perf_recipes.py index ce024cb80e..64054808df 100644 --- a/tests/unit_tests/recipes/test_gb300_ncclep_perf_recipes.py +++ b/tests/unit_tests/recipes/test_gb300_ncclep_perf_recipes.py @@ -12,7 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. """Every GB300 MoE perf recipe switched to NCCL EP carries the full NCCL EP stack, and the VRNVL72 -aliases built on the shared HybridEP bases stay on HybridEP.""" +aliases built on the shared HybridEP bases stay on HybridEP. VRNVL72 aliases built on an NCCL EP base +carry the same NCCL EP stack.""" import importlib @@ -53,6 +54,9 @@ ("qwen_vl.gb300.qwen35_vl", "qwen35_vl_35b_a3b_pretrain_8gpu_gb300_bf16_config"), ) +# VRNVL72 aliases that share the GB300 NCCL EP builder instead of a HybridEP base. +_VRNVL72_NCCLEP_RECIPES = (("nemotronh.vrnvl72.nemotronh", "nemotron_3_ultra_pretrain_256gpu_vrnvl72_fp8mx_config"),) + # VRNVL72 aliases of the same recipes: they reuse the shared HybridEP builders and must not pick up NCCL EP. _VRNVL72_ALIASES = ( ("deepseek.vrnvl72.deepseek_v3", "deepseek_v3_pretrain_256gpu_vrnvl72_bf16_config"), @@ -63,7 +67,6 @@ ("nemotronh.vrnvl72.nemotronh", "nemotron_3_nano_pretrain_8gpu_vrnvl72_bf16_config"), ("nemotronh.vrnvl72.nemotronh", "nemotron_3_super_pretrain_64gpu_vrnvl72_bf16_config"), ("nemotronh.vrnvl72.nemotronh", "nemotron_3_super_pretrain_64gpu_vrnvl72_nvfp4_config"), - ("nemotronh.vrnvl72.nemotronh", "nemotron_3_ultra_pretrain_256gpu_vrnvl72_fp8mx_config"), ("nemotronh.vrnvl72.nemotronh", "nemotron_3_5_lightning_pretrain_8gpu_vrnvl72_bf16_config"), ("nemotronh.vrnvl72.nemotronh", "nemotron_3_5_lightning_pretrain_8gpu_vrnvl72_fp8mx_config"), ("deepseek.vrnvl72.deepseek_v3", "deepseek_v3_pretrain_64gpu_vrnvl72_fp8mx_proxy_config"), @@ -94,7 +97,9 @@ def _keep_recipe_construction_offline(monkeypatch: pytest.MonkeyPatch) -> None: patch_recipe_construction_dependencies(monkeypatch) -@pytest.mark.parametrize(("module_name", "factory_name"), _GB300_NCCLEP_RECIPES, ids=lambda value: value) +@pytest.mark.parametrize( + ("module_name", "factory_name"), (*_GB300_NCCLEP_RECIPES, *_VRNVL72_NCCLEP_RECIPES), ids=lambda value: value +) def test_gb300_moe_perf_recipes_default_to_the_nccl_ep_stack(module_name: str, factory_name: str) -> None: cfg = _build(module_name, factory_name) diff --git a/tests/unit_tests/recipes/test_override_precedence.py b/tests/unit_tests/recipes/test_override_precedence.py index 156ad53333..b0a9f1d563 100644 --- a/tests/unit_tests/recipes/test_override_precedence.py +++ b/tests/unit_tests/recipes/test_override_precedence.py @@ -435,19 +435,23 @@ def test_local_full_iteration_is_normalized(self): assert recipe.rng.te_rng_tracker is True assert recipe.rerun_state_machine.check_for_nan_in_loss is False - def test_local_scoped_graph_is_rejected(self): - """Local graphs cannot capture individual layer modules in Bridge.""" + def test_local_scoped_graph_is_normalized(self): + """Local per-layer graphs populate current MCore module values.""" from utils.overrides import _set_cuda_graph_overrides - with pytest.raises( - ValueError, - match='cuda_graph_impl="local".*cuda_graph_impl="transformer_engine"', - ): - _set_cuda_graph_overrides( - _minimal_cuda_graph_recipe(), - cuda_graph_impl="local", - cuda_graph_scope="mlp", - ) + from megatron.bridge.utils.cuda_graph import cuda_graph_module_names, is_full_iteration_cuda_graph + + recipe = _set_cuda_graph_overrides( + _minimal_cuda_graph_recipe(), + cuda_graph_impl="local", + cuda_graph_scope="mlp", + ) + + assert recipe.model.cuda_graph_impl == "local" + assert cuda_graph_module_names(recipe.model) == ["mlp"] + assert not is_full_iteration_cuda_graph(recipe.model) + assert recipe.model.use_te_rng_tracker is True + assert recipe.rng.te_rng_tracker is True def test_transformer_engine_scoped_graph_is_normalized(self): """TE scoped graphs populate current MCore module values.""" diff --git a/tests/unit_tests/training/test_config.py b/tests/unit_tests/training/test_config.py index 80eac53e54..3d65673f83 100644 --- a/tests/unit_tests/training/test_config.py +++ b/tests/unit_tests/training/test_config.py @@ -2595,8 +2595,8 @@ def test_cuda_graph_local_inference_scope_allows_validation(self, inference_scop finally: restore_get_world_size_safe(og_ws, cfg_mod) - def test_cuda_graph_local_inference_scope_rejects_training_modules(self, monkeypatch): - """Test that an inference scope does not bypass local training-scope validation.""" + def test_cuda_graph_local_inference_scope_allows_training_modules(self, monkeypatch): + """Test that local per-layer training modules can be combined with an inference scope.""" gpt_model_cfg = create_test_gpt_config( cuda_graph_impl="local", inference_cuda_graph_scope="block", @@ -2610,14 +2610,15 @@ def test_cuda_graph_local_inference_scope_rejects_training_modules(self, monkeyp ) try: - with pytest.raises(ValueError, match='cuda_graph_impl="local"'): - container.validate() + container.validate() + assert container.model.cuda_graph_impl == "local" + assert cuda_graph_module_names(container.model) == ["mlp"] finally: restore_get_world_size_safe(og_ws, cfg_mod) - @pytest.mark.parametrize("graph_modules", [["mlp"], ["moe_router"]]) - def test_cuda_graph_local_scoped_modules_raise_clear_error(self, graph_modules, monkeypatch): - """Test that Bridge rejects local scoped graphs before MCore layer construction.""" + @pytest.mark.parametrize("graph_modules", [["attn"], ["attn", "mlp"]]) + def test_cuda_graph_local_scoped_modules_allow_validation(self, graph_modules, monkeypatch): + """Test that Bridge passes local per-layer CUDA graph modules through to MCore.""" gpt_model_cfg = create_test_gpt_config( cuda_graph_impl="local", use_te_rng_tracker=True, @@ -2630,16 +2631,14 @@ def test_cuda_graph_local_scoped_modules_raise_clear_error(self, graph_modules, ) try: - with pytest.raises( - ValueError, - match='cuda_graph_impl="local".*cuda_graph_impl="transformer_engine"', - ): - container.validate() + container.validate() + assert container.model.cuda_graph_impl == "local" + assert cuda_graph_module_names(container.model) == graph_modules finally: restore_get_world_size_safe(og_ws, cfg_mod) - def test_cuda_graph_local_deprecated_scope_raise_clear_error(self, monkeypatch): - """Test post-construction cuda_graph_scope overrides fail clearly for local scoped graphs.""" + def test_cuda_graph_local_deprecated_scope_allows_validation(self, monkeypatch): + """Test post-construction cuda_graph_scope overrides are accepted for local per-layer graphs.""" gpt_model_cfg = create_test_gpt_config(use_te_rng_tracker=True) gpt_model_cfg.cuda_graph_impl = "local" gpt_model_cfg.cuda_graph_scope = ["mlp"] @@ -2649,24 +2648,42 @@ def test_cuda_graph_local_deprecated_scope_raise_clear_error(self, monkeypatch): model_config=gpt_model_cfg, ) + try: + container.validate() + assert container.model.cuda_graph_impl == "local" + assert cuda_graph_module_names(container.model) == ["mlp"] + finally: + restore_get_world_size_safe(og_ws, cfg_mod) + + def test_cuda_graph_local_whole_layer_raise_clear_error(self, monkeypatch): + """Test that Bridge rejects local whole-layer capture before MCore layer construction.""" + gpt_model_cfg = create_test_gpt_config( + cuda_graph_impl="local", + use_te_rng_tracker=True, + ) + + container, og_ws, cfg_mod = create_test_config_container( + world_size_override=1, + model_config=gpt_model_cfg, + ) + try: with pytest.raises( ValueError, - match='cuda_graph_impl="local".*cuda_graph_impl="transformer_engine"', + match='cuda_graph_impl="local".*cuda_graph_modules', ): container.validate() finally: restore_get_world_size_safe(og_ws, cfg_mod) - def test_cuda_graph_local_deprecated_scope_direct_provider_raise_clear_error(self): - """Test direct provider construction validates local scoped graphs before MCore build.""" + def test_cuda_graph_local_whole_layer_direct_provider_raise_clear_error(self): + """Test direct provider construction validates local whole-layer capture before MCore build.""" gpt_model_cfg = create_test_gpt_config(use_te_rng_tracker=True, vocab_size=128) gpt_model_cfg.cuda_graph_impl = "local" - gpt_model_cfg.cuda_graph_scope = ["mlp"] with pytest.raises( ValueError, - match='cuda_graph_impl="local".*cuda_graph_impl="transformer_engine"', + match='cuda_graph_impl="local".*cuda_graph_modules', ): gpt_model_cfg.provide() diff --git a/tests/unit_tests/training/test_train.py b/tests/unit_tests/training/test_train.py index 53870e3968..4ab84fadc0 100644 --- a/tests/unit_tests/training/test_train.py +++ b/tests/unit_tests/training/test_train.py @@ -273,8 +273,9 @@ def assert_graph_deleted(): ): _finish_train(state, checkpoint_manager) + @patch("megatron.bridge.training.train.delete_local_cuda_graphs") @patch("megatron.bridge.training.train.gc.collect") - def test_delete_cuda_graphs_restores_full_graph_state(self, mock_collect): + def test_delete_cuda_graphs_restores_full_graph_state(self, mock_collect, mock_delete_local): """Cleanup must leave reusable full-graph mappings for the next training run.""" from megatron.core.full_cuda_graph import FullCudaGraphWrapper from megatron.core.optimizer.optimizer_cuda_graph import OptimizerCudaGraphWrapper @@ -301,6 +302,7 @@ def test_delete_cuda_graphs_restores_full_graph_state(self, mock_collect): assert OptimizerCudaGraphWrapper.result is None assert OptimizerCudaGraphWrapper.curr_iteration == 0 helper.delete_cuda_graphs.assert_not_called() + mock_delete_local.assert_called_once_with() mock_collect.assert_called_once()