feat(nemotron): Add GTP, Partial CG and NCCL-EP support for NT3 Ultra NVFP4 on GB200/GB300/VRNVL72 (#6468)

Signed-off-by: rghadia <rghadia@nvidia.com>

Author: ghadiaravi13

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

@@ -74,10 +74,6 @@ def _nemotron_3_super_nvfp4_precision() -> MixedPrecisionConfig:
 def _nemotron_3_ultra_nvfp4_precision() -> MixedPrecisionConfig:
     """Return the NVFP4 precision config used by Nemotron 3 Ultra perf recipes."""
     cfg = nemotron_3_ultra_bf16_with_nvfp4_mixed()
-    # Although MCore PR 4358 is merged,
-    # Megatron-FSDP's ParamAndGradBuffer has no NVFP4 packed-storage support, unlike the
-    # legacy DDP buffer, so FP4 primary weights fault during buffer init.
-    cfg.fp4_param_gather = False
     return cfg
 
 

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

@@ -20,7 +20,6 @@
     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,
@@ -293,29 +292,25 @@ def nemotron_3_ultra_pretrain_256gpu_gb200_fp8mx_config(*, num_gpus: int = 256)
 
 
 def nemotron_3_ultra_pretrain_256gpu_gb200_nvfp4_config(*, num_gpus: int = 256) -> ConfigContainer:
-    """Nemotron 3 Ultra (550B-A55B LatentMoE) pretrain: 256× GB200, NVFP4, Megatron-FSDP (HSDP).
+    """Nemotron 3 Ultra (550B-A55B LatentMoE) pretrain: 256× GB200, NVFP4, GTP.
 
     TP2 + SP (due to smaller GB200 HBM) / PP1 / CP1 / EP64 / ETP1, GBS 256 / MBS 1, seq 8192, BF16 + NVFP4 mixed
-    precision, HybridEP flex dispatcher, CuteDSL fused grouped MLP, selective recompute +
-    fine-grained activation offload of the expert MLP, MTP=2.
+    precision, dense GTP32, NCCLEP flex dispatcher with local attention/Mamba CUDA graphs, 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.
-            Defaults to the canonical 256-GPU benchmark allocation.
+        num_gpus: Retained for compatibility; sharding is rack-local and the
+            distributed-optimizer instance count is fixed at one.
     """
 
     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 = _nemotron_3_ultra_nvfp4_precision()
 
     _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)
-
     # Parallelism
     cfg.model.tensor_model_parallel_size = 2
     cfg.model.pipeline_model_parallel_size = 1
@@ -349,6 +344,21 @@ def nemotron_3_ultra_pretrain_256gpu_gb200_nvfp4_config(*, num_gpus: int = 256)
     cfg.model.recompute_granularity = "selective"
     cfg.model.recompute_modules = ["moe_act"]
 
+    # Shard dense weights with GTP within each 64-GPU rack. GTP currently
+    # requires a single distributed-optimizer instance.
+    _apply_nemotron_3_ultra_gtp(cfg)
+
+    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,
@@ -363,11 +373,8 @@ def nemotron_3_ultra_pretrain_256gpu_gb200_nvfp4_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,

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

@@ -20,7 +20,6 @@
     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,
@@ -320,15 +319,16 @@ def nemotron_3_ultra_pretrain_256gpu_gb300_fp8mx_config(*, num_gpus: int = 256)
 
 
 def nemotron_3_ultra_pretrain_256gpu_gb300_nvfp4_config(*, num_gpus: int = 256) -> ConfigContainer:
-    """Nemotron 3 Ultra (550B-A55B LatentMoE) pretrain: 256× GB300, NVFP4, Megatron-FSDP (HSDP).
+    """Nemotron 3 Ultra (550B-A55B LatentMoE) pretrain: 256× GB300, NVFP4, GTP.
 
     TP1 / PP1 / CP1 / EP64 / ETP1, GBS 256 / MBS 1, seq 8192, BF16 + NVFP4 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 with local attention/Mamba CUDA
+    graphs, 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.
-            Defaults to the canonical 256-GPU benchmark allocation.
+        num_gpus: Retained for compatibility; sharding is rack-local and the
+            distributed-optimizer instance count is fixed at one.
     """
 
     expert_model_parallel_size = 64
@@ -371,8 +371,20 @@ def nemotron_3_ultra_pretrain_256gpu_gb300_nvfp4_config(*, num_gpus: int = 256)
 
     _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)
+
+    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 = {
@@ -385,11 +397,8 @@ def nemotron_3_ultra_pretrain_256gpu_gb300_nvfp4_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": 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/perf_recipes/nemotronh/vrnvl72/nemotronh.py

@@ -260,8 +260,8 @@ def nemotron_3_ultra_pretrain_256gpu_vrnvl72_nvfp4_config(*, num_gpus: int = 256
     """Nemotron 3 Ultra pretrain: 256× VRNVL72, NVFP4 (alias of GB300).
 
     Args:
-        num_gpus: Launch GPU count used to size HSDP optimizer instances.
-            Defaults to the canonical 256-GPU benchmark allocation.
+        num_gpus: Retained for compatibility; sharding is rack-local and the
+            distributed-optimizer instance count is fixed at one.
     """
     cfg = nemotron_3_ultra_pretrain_256gpu_gb300_nvfp4_config(num_gpus=num_gpus)
     # Keep process settings next to the recipe so users can see the exact benchmark environment.
@@ -275,11 +275,8 @@ def nemotron_3_ultra_pretrain_256gpu_vrnvl72_nvfp4_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,

tests/unit_tests/recipes/nemotronh/test_nemotron_3_ultra.py

@@ -18,7 +18,10 @@
 import torch
 
 import megatron.bridge.recipes as recipes
+from megatron.bridge.perf_recipes.environment import HYBRID_EP_ENV_NAMES
 from megatron.bridge.perf_recipes.nemotronh import (
+    nemotron_3_ultra_pretrain_256gpu_gb200_nvfp4_config,
+    nemotron_3_ultra_pretrain_256gpu_gb300_nvfp4_config,
     nemotron_3_ultra_pretrain_256gpu_vrnvl72_fp8mx_config,
     nemotron_3_ultra_pretrain_256gpu_vrnvl72_nvfp4_config,
 )
@@ -36,6 +39,7 @@
     nemotron_3_ultra_pretrain_config,
     nemotron_3_ultra_sft_openmathinstruct2_packed_config,
 )
+from megatron.bridge.utils.cuda_graph import cuda_graph_module_names
 from tests.unit_tests.recipes.recipe_test_utils import (
     NEMOTRON_3_ULTRA_HYBRID_LAYER_PATTERN,
     patch_recipe_module_global,
@@ -202,23 +206,69 @@ def test_vrnvl72_nvfp4_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.moe_flex_dispatcher_backend == "ncclep"
+    assert cfg.model.moe_use_grouped_tensor is True
     assert cfg.model.moe_router_padding_for_quantization is True
-    assert cfg.model.cuda_graph_impl == "none"
+    assert cfg.model.cuda_graph_impl == "local"
+    assert cuda_graph_module_names(cfg.model) == ["attn", "mamba"]
+    assert cfg.model.use_te_rng_tracker is True
+    assert cfg.rng.te_rng_tracker is True
     assert cfg.mixed_precision.fp4 == "e2m1"
-    assert cfg.mixed_precision.fp4_param_gather is False
+    assert cfg.mixed_precision.fp4_param_gather is True
     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.dist.use_megatron_fsdp is False
+    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.keys().isdisjoint(HYBRID_EP_ENV_NAMES)
     assert cfg.env_vars["NVTE_CPU_OFFLOAD_V1"] == 1
     assert cfg.env_vars["NVTE_CUTEDSL_FUSED_GROUPED_MLP"] == 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["NVTE_USE_FAST_MATH"] == 1
-    assert cfg.env_vars["USE_MNNVL"] == 1
+
+
+@pytest.mark.unit
+@pytest.mark.parametrize(
+    ("recipe_factory", "tensor_model_parallel_size", "sequence_parallel"),
+    [
+        (nemotron_3_ultra_pretrain_256gpu_gb200_nvfp4_config, 2, True),
+        (nemotron_3_ultra_pretrain_256gpu_gb300_nvfp4_config, 1, False),
+    ],
+    ids=["gb200", "gb300"],
+)
+def test_nvfp4_perf_recipes_use_gtp_partial_cuda_graphs_and_nccl_ep(
+    recipe_factory, tensor_model_parallel_size: int, sequence_parallel: bool
+) -> None:
+    """GB200 and GB300 Ultra NVFP4 shard dense weights rack-locally with GTP instead of HSDP."""
+    cfg = recipe_factory()
+
+    assert cfg.model.tensor_model_parallel_size == tensor_model_parallel_size
+    assert cfg.model.sequence_parallel is sequence_parallel
+    assert cfg.model.expert_model_parallel_size == 64
+    assert cfg.model.expert_tensor_parallel_size == 1
+    assert cfg.mixed_precision.fp4 == "e2m1"
+    assert cfg.mixed_precision.fp4_param_gather is True
+
+    assert cfg.dist.use_megatron_fsdp is False
+    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.model.cuda_graph_impl == "local"
+    assert cuda_graph_module_names(cfg.model) == ["attn", "mamba"]
+    assert cfg.model.use_te_rng_tracker is True
+    assert cfg.rng.te_rng_tracker is True
+
+    assert cfg.model.moe_token_dispatcher_type == "flex"
+    assert cfg.model.moe_flex_dispatcher_backend == "ncclep"
+    assert cfg.model.moe_use_grouped_tensor is True
+    assert cfg.env_vars["NCCL_EP_HT_EM_PULL_PUSH"] == 1
+    assert cfg.env_vars.keys().isdisjoint(HYBRID_EP_ENV_NAMES)
 
 
 @pytest.mark.unit

tests/unit_tests/recipes/test_gb300_ncclep_perf_recipes.py

@@ -42,6 +42,7 @@
     ("nemotronh.gb300.nemotronh", "nemotron_3_super_pretrain_64gpu_gb300_fp8mx_config"),
     ("nemotronh.gb300.nemotronh", "nemotron_3_super_pretrain_64gpu_gb300_nvfp4_config"),
     ("nemotronh.gb300.nemotronh", "nemotron_3_ultra_pretrain_256gpu_gb300_fp8mx_config"),
+    ("nemotronh.gb300.nemotronh", "nemotron_3_ultra_pretrain_256gpu_gb300_nvfp4_config"),
     ("nemotronh.gb300.nemotronh", "nemotron_3_nano_pretrain_8gpu_gb300_bf16_config"),
     ("nemotronh.gb300.nemotronh", "nemotron_3_nano_pretrain_8gpu_gb300_fp8mx_config"),
     ("nemotronh.gb300.nemotronh", "nemotron_3_5_lightning_pretrain_8gpu_gb300_bf16_config"),
@@ -57,7 +58,10 @@
 )
 
 # 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_NCCLEP_RECIPES = (
+    ("nemotronh.vrnvl72.nemotronh", "nemotron_3_ultra_pretrain_256gpu_vrnvl72_fp8mx_config"),
+    ("nemotronh.vrnvl72.nemotronh", "nemotron_3_ultra_pretrain_256gpu_vrnvl72_nvfp4_config"),
+)
 
 # VRNVL72 aliases of the same recipes: they reuse the shared HybridEP builders and must not pick up NCCL EP.
 _VRNVL72_ALIASES = (