feat(nemotron): Add GTP + partial CG + cuDNN LN support for NT3 Ultra MXFP8 GB200/GB300/VRNVL72 recipes

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.
 

src/megatron/bridge/perf_recipes/nemotronh/gb200/nemotronh.py

@@ -13,13 +13,15 @@
 # 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,
     ConfigContainer,
     _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,

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,

src/megatron/bridge/perf_recipes/nemotronh/vr200/nemotronh.py

@@ -240,11 +240,8 @@ def nemotron_3_ultra_pretrain_256gpu_vr200_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,

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()
 

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."
         )
 
 

tests/unit_tests/recipes/nemotronh/test_nemotron_3_ultra.py

@@ -161,7 +161,7 @@ def test_gb200_ep16_pretrain_preserves_functional_policy() -> None:
 
 @pytest.mark.unit
 def test_vr200_perf_recipe_uses_nvl72_ultra_topology() -> None:
-    """VR200 Ultra preserves the GB300 execution layout with explicit NVL72 settings."""
+    """VR200 Ultra preserves the GB300 MXFP8 GTP + NCCL EP execution layout."""
     cfg = nemotron_3_ultra_pretrain_256gpu_vr200_fp8mx_config()
 
     assert cfg.model.tensor_model_parallel_size == 1
@@ -171,17 +171,17 @@ def test_vr200_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

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 VR200
-aliases built on the shared HybridEP bases stay on HybridEP."""
+aliases built on the shared HybridEP bases stay on HybridEP. VR200 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"),
 )
 
+# VR200 aliases that share the GB300 NCCL EP builder instead of a HybridEP base.
+_VR200_NCCLEP_RECIPES = (("nemotronh.vr200.nemotronh", "nemotron_3_ultra_pretrain_256gpu_vr200_fp8mx_config"),)
+
 # VR200 aliases of the same recipes: they reuse the shared HybridEP builders and must not pick up NCCL EP.
 _VR200_ALIASES = (
     ("deepseek.vr200.deepseek_v3", "deepseek_v3_pretrain_256gpu_vr200_bf16_config"),
@@ -63,7 +67,6 @@
     ("nemotronh.vr200.nemotronh", "nemotron_3_nano_pretrain_8gpu_vr200_bf16_config"),
     ("nemotronh.vr200.nemotronh", "nemotron_3_super_pretrain_64gpu_vr200_bf16_config"),
     ("nemotronh.vr200.nemotronh", "nemotron_3_super_pretrain_64gpu_vr200_nvfp4_config"),
-    ("nemotronh.vr200.nemotronh", "nemotron_3_ultra_pretrain_256gpu_vr200_fp8mx_config"),
     ("nemotronh.vr200.nemotronh", "nemotron_3_5_lightning_pretrain_8gpu_vr200_bf16_config"),
     ("nemotronh.vr200.nemotronh", "nemotron_3_5_lightning_pretrain_8gpu_vr200_fp8mx_config"),
     ("deepseek.vr200.deepseek_v3", "deepseek_v3_pretrain_64gpu_vr200_fp8mx_proxy_config"),
@@ -81,7 +84,7 @@
 # every recipe module that is already loaded, so a lazy import inside a test would escape the patch.
 _RECIPE_MODULES = {
     module_name: importlib.import_module(f"megatron.bridge.perf_recipes.{module_name}")
-    for module_name in sorted({name for name, _ in (*_GB300_NCCLEP_RECIPES, *_VR200_ALIASES)})
+    for module_name in sorted({name for name, _ in (*_GB300_NCCLEP_RECIPES, *_VR200_NCCLEP_RECIPES, *_VR200_ALIASES)})
 }
 
 
@@ -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, *_VR200_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)
 

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