feat(nemotron): add NVFP4 mixed precision configurations for Ultra pretraining on GB200 and GB300

- Introduced new mixed precision configuration for Nemotron 3 Ultra using BF16 with NVFP4.
- Added NVFP4 pretraining configurations for 256 GPUs on both GB200 and GB300.
- Updated existing functions to incorporate new precision settings and ensure compatibility with the latest architecture.

Signed-off-by: Ravi Ghadia <ghadiaravi13@users.noreply.github.com>

Author: ghadiaravi13

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

@@ -30,6 +30,7 @@
     nemotron_3_super_pretrain_64gpu_gb200_fp8mx_config,
     nemotron_3_super_pretrain_64gpu_gb200_nvfp4_config,
     nemotron_3_ultra_pretrain_256gpu_gb200_fp8mx_config,
+    nemotron_3_ultra_pretrain_256gpu_gb200_nvfp4_config,
     nemotronh_56b_pretrain_64gpu_gb200_fp8cs_config,
 )
 from megatron.bridge.perf_recipes.nemotronh.gb300.nemotronh import (
@@ -40,6 +41,7 @@
     nemotron_3_super_pretrain_64gpu_gb300_fp8mx_config,
     nemotron_3_super_pretrain_64gpu_gb300_nvfp4_config,
     nemotron_3_ultra_pretrain_256gpu_gb300_fp8mx_config,
+    nemotron_3_ultra_pretrain_256gpu_gb300_nvfp4_config,
     nemotronh_56b_pretrain_64gpu_gb300_fp8cs_config,
     nemotronh_56b_pretrain_256gpu_gb300_bf16_config,
     nemotronh_56b_pretrain_256gpu_gb300_fp8cs_config,

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

@@ -27,7 +27,11 @@
 from megatron.bridge.recipes.nemotronh.nemotron_3_ultra import nemotron_3_ultra_pretrain_config
 from megatron.bridge.recipes.nemotronh.nemotronh import nemotronh_56b_pretrain_config
 from megatron.bridge.training.config import ConfigContainer
-from megatron.bridge.training.mixed_precision import MixedPrecisionConfig, nemotron_3_super_bf16_with_nvfp4_mixed
+from megatron.bridge.training.mixed_precision import (
+    MixedPrecisionConfig,
+    nemotron_3_super_bf16_with_nvfp4_mixed,
+    nemotron_3_ultra_bf16_with_nvfp4_mixed,
+)
 
 
 _TE_QUANT_CFG_PATH = Path(__file__).with_name("te_quant.cfg")
@@ -65,7 +69,19 @@ def _with_global_batch_size(cfg: ConfigContainer, global_batch_size: int) -> Con
 def _nemotron_3_super_nvfp4_precision() -> MixedPrecisionConfig:
     """Return the NVFP4 precision config used by Nemotron 3 Super perf recipes."""
     cfg = nemotron_3_super_bf16_with_nvfp4_mixed()
-    # Disabled until MCore PR 4358 lands.
+    # 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
+
+
+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

@@ -23,6 +23,7 @@
     _apply_nemotron_3_ultra_perf_defaults,
     _benchmark_common,
     _nemotron_3_super_nvfp4_precision,
+    _nemotron_3_ultra_nvfp4_precision,
     _perf_precision,
     load_quantization_recipe,
     nemotron_3_nano_pretrain_config,
@@ -303,6 +304,93 @@ def nemotron_3_ultra_pretrain_256gpu_gb200_fp8mx_config() -> ConfigContainer:
     return cfg
 
 
+def nemotron_3_ultra_pretrain_256gpu_gb200_nvfp4_config() -> ConfigContainer:
+    """Nemotron 3 Ultra (550B-A55B LatentMoE) pretrain: 256× GB200, NVFP4, Megatron-FSDP (HSDP).
+
+    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.
+    """
+
+    num_gpus = 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 = _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
+    cfg.model.virtual_pipeline_model_parallel_size = None
+    cfg.model.context_parallel_size = 1
+    cfg.model.sequence_parallel = True
+    cfg.model.expert_tensor_parallel_size = 1
+    cfg.model.pipeline_model_parallel_layout = None
+    cfg.model.seq_length = 8192
+
+    # Only tensors larger than 350M elements are offloaded, which
+    # approximates offloading the moe_act (pre-activation input) for seq 8192/2 (due to SP) / MBS 1.
+    cfg.model.min_offloaded_tensor_size = 350_000_000
+
+    # MXFP8 requires router padding for quantization.
+    cfg.model.moe_router_padding_for_quantization = True
+
+    # GPU-count specific overrides of the canonical (256-GPU / EP64) defaults.
+    cfg.model.expert_model_parallel_size = expert_model_parallel_size
+    cfg.train.global_batch_size = global_batch_size
+
+    # Fine-grained activation offloading. Requires NVTE_CPU_OFFLOAD_V1=1 in the
+    # launch environment (set in this recipe's env_vars).
+    # NOTE: also requires setting the min_offloaded_tensor_size to selectively offload moe_act of the fused_group_mlp, to avoid CPU OOM issues
+    cfg.model.fine_grained_activation_offloading = True
+    cfg.model.offload_modules = ["fused_group_mlp"]
+    cfg.model.fine_grained_offloading_max_inflight_offloads = 1
+
+    # Selective recompute of the MoE activation
+    # recomputes the activation output of the MoE expert MLP, while FC1 output (activation input) is saved and offloaded to cpu
+    cfg.model.recompute_granularity = "selective"
+    cfg.model.recompute_modules = ["moe_act"]
+
+    # Keep process settings next to the recipe so users can see the exact benchmark environment.
+    cfg.env_vars = {
+        **COMMON_PERF_ENV_VARS,
+        # CUDA stream scheduling for this model and parallel layout.
+        "CUDA_DEVICE_MAX_CONNECTIONS": 32,
+        # CUDA graph and allocator behavior for this recipe.
+        "NCCL_GRAPH_REGISTER": 0,
+        # TODO: graph_capture_record_stream_reuse might be potentially useful
+        # when enabling CG, because it allows the freed-up memory buffers of the offloaded tensors
+        # to be reused during the CG capture, thus keeping the peak memory usage lower.
+        "PYTORCH_CUDA_ALLOC_CONF": ("expandable_segments:True,graph_capture_record_stream_reuse:True"),
+        "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,
+        # Transformer Engine overlap settings for this model.
+        "NVTE_BWD_LAYERNORM_SM_MARGIN": 20,
+        "NVTE_FWD_LAYERNORM_SM_MARGIN": 20,
+        # Required by fine_grained_activation_offloading (TE >= 2.10.0) to avoid
+        # offloading weights;
+        "NVTE_CPU_OFFLOAD_V1": 1,
+        # Enable TE's CuteDSL fused grouped MLP kernel (sm100+). Required by the
+        # op fuser + fused weighted squared-ReLU with moe_act activation recompute
+        # (ScaledSReLU(activation_recompute_in_mlp=True) only runs on this path).
+        "NVTE_CUTEDSL_FUSED_GROUPED_MLP": 1,
+    }
+    return cfg
+
+
 def nemotron_3_nano_pretrain_8gpu_gb200_bf16_config() -> ConfigContainer:
     """Nemotron 3 Nano pretrain: 8× GB200, BF16."""
     cfg = nemotron_3_nano_pretrain_config()

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

@@ -23,6 +23,7 @@
     _apply_nemotron_3_ultra_perf_defaults,
     _benchmark_common,
     _nemotron_3_super_nvfp4_precision,
+    _nemotron_3_ultra_nvfp4_precision,
     _perf_precision,
     _with_global_batch_size,
     load_quantization_recipe,
@@ -308,6 +309,102 @@ def nemotron_3_ultra_pretrain_256gpu_gb300_fp8mx_config() -> ConfigContainer:
     return cfg
 
 
+def _nemotron_3_ultra_gb300_nvfp4_config(
+    *, num_gpus: int, expert_model_parallel_size: int, global_batch_size: int
+) -> ConfigContainer:
+    """Shared builder for Nemotron 3 Ultra GB300 NVFP4 Megatron-FSDP perf recipes."""
+    cfg = nemotron_3_ultra_pretrain_config()
+    cfg.mixed_precision = _nemotron_3_ultra_nvfp4_precision()
+
+    """
+    Uses TP1 / PP1 / CP1 / EP64 / ETP1, GBS 256 / MBS 1,
+    seq 8192, HybridEP 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``.
+    """
+    # Parallelism
+    cfg.model.tensor_model_parallel_size = 1
+    cfg.model.pipeline_model_parallel_size = 1
+    cfg.model.virtual_pipeline_model_parallel_size = None
+    cfg.model.context_parallel_size = 1
+    cfg.model.sequence_parallel = False
+    cfg.model.expert_tensor_parallel_size = 1
+    cfg.model.pipeline_model_parallel_layout = None
+    cfg.model.seq_length = 8192
+
+    # Only tensors larger than 500MB are offloaded, which
+    # approximates offloading the moe_act input for seq 8192 / MBS 1.
+    cfg.model.min_offloaded_tensor_size = 500_000_000
+
+    # MXFP8 requires router padding for quantization.
+    cfg.model.moe_router_padding_for_quantization = True
+
+    # GPU-count specific overrides of the canonical (256-GPU / EP64) defaults.
+    cfg.model.expert_model_parallel_size = expert_model_parallel_size
+    cfg.train.global_batch_size = global_batch_size
+
+    # Fine-grained activation offloading. Requires NVTE_CPU_OFFLOAD_V1=1 in the
+    # launch environment (set in this recipe's env_vars).
+    # NOTE: also requires setting the min_offloaded_tensor_size to avoid CPU OOM issues
+    cfg.model.fine_grained_activation_offloading = True
+    cfg.model.offload_modules = ["fused_group_mlp"]
+
+    # Selective recompute of the MoE activation
+    # recomputes the activation output of the MoE expert MLP, while FC1 output (activation input) is saved and offloaded to cpu
+    cfg.model.recompute_granularity = "selective"
+    cfg.model.recompute_modules = ["moe_act"]
+
+    _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)
+
+    return cfg
+
+
+def nemotron_3_ultra_pretrain_256gpu_gb300_nvfp4_config() -> ConfigContainer:
+    """Nemotron 3 Ultra (550B-A55B LatentMoE) pretrain: 256× GB300, NVFP4, Megatron-FSDP (HSDP).
+
+    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.
+    """
+    cfg = _nemotron_3_ultra_gb300_nvfp4_config(
+        num_gpus=256,
+        expert_model_parallel_size=64,
+        global_batch_size=256,
+    )
+    # Keep process settings next to the recipe so users can see the exact benchmark environment.
+    cfg.env_vars = {
+        **COMMON_PERF_ENV_VARS,
+        # CUDA stream scheduling for this model and parallel layout.
+        "CUDA_DEVICE_MAX_CONNECTIONS": 32,
+        # CUDA graph and allocator behavior for this recipe.
+        "NCCL_GRAPH_REGISTER": 0,
+        "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True",
+        "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,
+        # Transformer Engine overlap settings for this model.
+        "NVTE_BWD_LAYERNORM_SM_MARGIN": 20,
+        "NVTE_FWD_LAYERNORM_SM_MARGIN": 20,
+        # Required by fine_grained_activation_offloading (TE >= 2.10.0) to avoid
+        # offloading weights;
+        "NVTE_CPU_OFFLOAD_V1": 1,
+        # Enable TE's CuteDSL fused grouped MLP kernel (sm100+). Required by the
+        # op fuser + fused weighted squared-ReLU with moe_act activation recompute
+        # (ScaledSReLU(activation_recompute_in_mlp=True) only runs on this path).
+        "NVTE_CUTEDSL_FUSED_GROUPED_MLP": 1,
+    }
+    return cfg
+
+
 def nemotron_3_nano_pretrain_8gpu_gb300_bf16_config() -> ConfigContainer:
     """Nemotron 3 Nano pretrain: 8× GB300, BF16."""
     cfg = nemotron_3_nano_pretrain_config()

src/megatron/bridge/training/mixed_precision.py

@@ -429,6 +429,19 @@ def nemotron_3_super_bf16_with_nvfp4_mixed() -> MixedPrecisionConfig:
     return cfg
 
 
+@register
+def nemotron_3_ultra_bf16_with_nvfp4_mixed() -> MixedPrecisionConfig:
+    """Create a MixedPrecisionConfig for mixed precision training using BF16 with NVFP4
+    Returns:
+        MixedPrecisionConfig: Configuration for BF16 with NVFP4 mixed precision training
+    """
+    cfg = bf16_with_nvfp4_mixed()
+    cfg.first_last_layers_bf16 = True
+    cfg.num_layers_at_start_in_bf16 = 0
+    cfg.num_layers_at_end_in_bf16 = 16  # last 15% of layers in BF16
+    return cfg
+
+
 def get_mixed_precision_config(name: str | MixedPrecisionConfig) -> MixedPrecisionConfig:
     """Return a :class:`MixedPrecisionConfig` for *name*.