feat (recipe): add Nemotron 3 Ultra MXFP8 GB200 pretraining perf-recipe

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

@@ -29,6 +29,7 @@
     nemotron_3_super_pretrain_64gpu_gb200_bf16_config,
     nemotron_3_super_pretrain_64gpu_gb200_fp8mx_config,
     nemotron_3_super_pretrain_64gpu_gb200_nvfp4_config,
+    nemotron_3_ultra_pretrain_256gpu_gb200_fp8mx_config,
     nemotronh_56b_pretrain_64gpu_gb200_fp8cs_config,
 )
 from megatron.bridge.perf_recipes.nemotronh.gb300.nemotronh import (

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

@@ -19,12 +19,15 @@
     ConfigContainer,
     _apply_nemotron_3_nano_perf_defaults,
     _apply_nemotron_3_super_perf_defaults,
+    _apply_nemotron_3_ultra_fsdp_hsdp,
+    _apply_nemotron_3_ultra_perf_defaults,
     _benchmark_common,
     _nemotron_3_super_nvfp4_precision,
     _perf_precision,
     load_quantization_recipe,
     nemotron_3_nano_pretrain_config,
     nemotron_3_super_pretrain_config,
+    nemotron_3_ultra_pretrain_config,
     nemotronh_56b_pretrain_config,
 )
 from megatron.bridge.utils.cuda_graph import set_cuda_graph_modules
@@ -213,6 +216,93 @@ def nemotron_3_super_pretrain_64gpu_gb200_nvfp4_config() -> ConfigContainer:
     return cfg
 
 
+def nemotron_3_ultra_pretrain_256gpu_gb200_fp8mx_config() -> ConfigContainer:
+    """Nemotron 3 Ultra (550B-A55B LatentMoE) pretrain: 256× GB200, MXFP8, Megatron-FSDP (HSDP).
+
+    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 +
+    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 = _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)
+
+    # 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()

tests/unit_tests/recipes/test_perf_recipe_environment.py

@@ -72,6 +72,15 @@ def _function(path: Path, function_name: str) -> ast.FunctionDef:
 def _explicit_environment(path: Path, function_name: str) -> dict[str, str | int | float | bool]:
     """Read the literal env mapping written in a flat recipe builder."""
     function = _function(path, function_name)
+    local_constants = {
+        node.targets[0].id: node.value.value
+        for node in function.body
+        if isinstance(node, ast.Assign)
+        and len(node.targets) == 1
+        and isinstance(node.targets[0], ast.Name)
+        and isinstance(node.value, ast.Constant)
+        and isinstance(node.value.value, (str, int, float, bool))
+    }
     assignments = [
         node
         for node in function.body
@@ -93,7 +102,12 @@ def _explicit_environment(path: Path, function_name: str) -> dict[str, str | int
             assert isinstance(value, ast.Name) and value.id == "COMMON_PERF_ENV_VARS"
             common_expansions += 1
             continue
-        result[ast.literal_eval(key)] = ast.literal_eval(value)
+        if isinstance(value, ast.Name):
+            assert value.id in local_constants
+            env_value = local_constants[value.id]
+        else:
+            env_value = ast.literal_eval(value)
+        result[ast.literal_eval(key)] = env_value
     assert common_expansions == 1
     return result