diff --git a/src/megatron/bridge/perf_recipes/nemotronh/__init__.py b/src/megatron/bridge/perf_recipes/nemotronh/__init__.py index 05b9864b96..785860600b 100644 --- a/src/megatron/bridge/perf_recipes/nemotronh/__init__.py +++ b/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 ( diff --git a/src/megatron/bridge/perf_recipes/nemotronh/gb200/nemotronh.py b/src/megatron/bridge/perf_recipes/nemotronh/gb200/nemotronh.py index 37c0e26bdb..2dd44834b3 100644 --- a/src/megatron/bridge/perf_recipes/nemotronh/gb200/nemotronh.py +++ b/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() diff --git a/tests/unit_tests/recipes/test_perf_recipe_environment.py b/tests/unit_tests/recipes/test_perf_recipe_environment.py index 1a6cf1c6a9..201184a38a 100644 --- a/tests/unit_tests/recipes/test_perf_recipe_environment.py +++ b/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