diff --git a/scripts/training/README.md b/scripts/training/README.md index 158c1c0c2e..0286e9e693 100644 --- a/scripts/training/README.md +++ b/scripts/training/README.md @@ -88,7 +88,7 @@ Text SFT/PEFT benchmark recipes retain the flat runner's mock-data default. Qwen model-specific dataset configuration. Exported benchmark PEFT recipes are fixed LoRA configs; DoRA remains available through configurable library recipes. Explicit benchmark dataset replacement remains on the compatibility launcher. -Five legacy duplicate names resolve to the benchmark definition; their library workloads remain available through +Known duplicate names resolve to the benchmark definition; their library workloads remain available through the corresponding generic recipe aliases. New recipe names should be unique across both packages. Text recipes default to `llm_step`; all recipes infer their modality-specific forward step from the same registry. diff --git a/scripts/training/recipe_metadata.py b/scripts/training/recipe_metadata.py index 29c3cd9e2a..5fe0d5f3ed 100644 --- a/scripts/training/recipe_metadata.py +++ b/scripts/training/recipe_metadata.py @@ -55,6 +55,7 @@ "gpt_oss_120b_pretrain_64gpu_h100_bf16_config", "llama3_70b_peft_8gpu_h100_bf16_config", "llama3_70b_sft_32gpu_h100_bf16_config", + "nemotron_3_nano_pretrain_8gpu_gb200_bf16_config", "qwen3_235b_a22b_pretrain_256gpu_h100_bf16_config", "qwen3_30b_a3b_pretrain_16gpu_h100_bf16_config", } diff --git a/src/megatron/bridge/perf_recipes/nemotronh/common.py b/src/megatron/bridge/perf_recipes/nemotronh/common.py index 8baf41f4ce..460662fa98 100644 --- a/src/megatron/bridge/perf_recipes/nemotronh/common.py +++ b/src/megatron/bridge/perf_recipes/nemotronh/common.py @@ -20,7 +20,9 @@ from megatron.core.quantization.utils import load_quantization_recipe from megatron.bridge.perf_recipes._common import _benchmark_common, _perf_precision -from megatron.bridge.recipes.nemotronh.nemotron_3_nano import nemotron_3_nano_pretrain_config +from megatron.bridge.recipes.nemotronh._nemotron_3_nano import ( + _nemotron_3_nano_pretrain_reference_config as nemotron_3_nano_pretrain_config, +) from megatron.bridge.recipes.nemotronh.nemotron_3_super import nemotron_3_super_pretrain_config 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 diff --git a/src/megatron/bridge/recipes/nemotronh/__init__.py b/src/megatron/bridge/recipes/nemotronh/__init__.py index 1540e438a8..de6b77c096 100644 --- a/src/megatron/bridge/recipes/nemotronh/__init__.py +++ b/src/megatron/bridge/recipes/nemotronh/__init__.py @@ -19,7 +19,9 @@ nemotron_3_5_nano_pretrain_8k_fsdp_config, nemotron_3_5_nano_sft_openmathinstruct2_packed_tp1_config, nemotron_3_nano_gb200_pretrain_config, + nemotron_3_nano_peft_8gpu_gb200_bf16_config, nemotron_3_nano_pretrain_8gpu_gb200_bf16_config, + nemotron_3_nano_sft_8gpu_gb200_bf16_config, ) from megatron.bridge.recipes.nemotronh.nemotron_3_nano import ( nemotron_3_5_nano_peft_config, @@ -107,7 +109,9 @@ "nemotron_3_nano_sft_config", "nemotron_3_nano_peft_config", "nemotron_3_nano_gb200_pretrain_config", + "nemotron_3_nano_peft_8gpu_gb200_bf16_config", "nemotron_3_nano_pretrain_8gpu_gb200_bf16_config", + "nemotron_3_nano_sft_8gpu_gb200_bf16_config", # Nemotron 3 Nano 4B model "nemotron_3_nano_4b_pretrain_config", "nemotron_3_nano_4b_sft_config", diff --git a/src/megatron/bridge/recipes/nemotronh/_nemotron_3_nano.py b/src/megatron/bridge/recipes/nemotronh/_nemotron_3_nano.py new file mode 100644 index 0000000000..f56a476fc0 --- /dev/null +++ b/src/megatron/bridge/recipes/nemotronh/_nemotron_3_nano.py @@ -0,0 +1,310 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Hardware-neutral Nemotron 3 Nano recipe builders.""" + +import torch +from megatron.core.activations import squared_relu + +from megatron.bridge.models.hybrid.hybrid_provider import HybridModelProvider +from megatron.bridge.peft.base import PEFT +from megatron.bridge.peft.lora import LoRA +from megatron.bridge.recipes.common import _peft_common, _pretrain_common, _sft_common +from megatron.bridge.recipes.utils.dataset_utils import default_peft_config +from megatron.bridge.training.comm_overlap import CommOverlapConfig +from megatron.bridge.training.config import ConfigContainer +from megatron.bridge.utils.cuda_graph import clear_cuda_graph_modules + + +def _nemotron_3_nano_pretrain_reference_config() -> ConfigContainer: + """Build the shared Nemotron 3 Nano pretraining contract. + + Hardware-specific library and performance recipes apply execution settings + to this config without changing its training semantics. + + Returns: + Pretraining configuration for Nemotron 3 Nano. + """ + cfg = _pretrain_common() + + cfg.model = HybridModelProvider( + hybrid_layer_pattern="MEMEM*EMEMEM*EMEMEM*EMEMEM*EMEMEM*EMEMEMEM*EMEMEMEME", + num_layers=52, + hidden_size=2688, + mamba_num_heads=64, + kv_channels=128, + mamba_state_dim=128, + ffn_hidden_size=1856, + num_attention_heads=32, + mamba_head_dim=64, + seq_length=8192, + num_query_groups=2, + num_moe_experts=128, + moe_ffn_hidden_size=1856, + moe_shared_expert_intermediate_size=3712, + moe_router_topk=6, + moe_router_topk_scaling_factor=2.5, + moe_router_num_groups=1, + moe_router_group_topk=1, + mamba_num_groups=8, + make_vocab_size_divisible_by=128, + activation_func=squared_relu, + masked_softmax_fusion=True, + apply_query_key_layer_scaling=False, + persist_layer_norm=True, + attention_softmax_in_fp32=False, + first_last_layers_bf16=True, + is_hybrid_model=True, + moe_aux_loss_coeff=0.0001, + moe_router_score_function="sigmoid", + moe_router_enable_expert_bias=True, + moe_router_load_balancing_type="seq_aux_loss", + moe_router_dtype="fp32", + moe_grouped_gemm=True, + moe_token_dispatcher_type="alltoall", + moe_permute_fusion=True, + moe_shared_expert_overlap=True, + tensor_model_parallel_size=4, + pipeline_model_parallel_size=1, + pipeline_dtype=torch.bfloat16, + virtual_pipeline_model_parallel_size=None, + context_parallel_size=1, + sequence_parallel=True, + expert_tensor_parallel_size=1, + expert_model_parallel_size=8, + ) + + cfg.tokenizer.tokenizer_model = "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16" + + cfg.dataset.seq_length = 8192 + cfg.dataset.blend = None + cfg.dataset.num_workers = 8 + cfg.dataset.mmap_bin_files = False + + cfg.model.pipeline_model_parallel_layout = None + cfg.model.moe_token_dispatcher_type = "flex" + cfg.model.moe_flex_dispatcher_backend = "deepep" + cfg.model.moe_hybridep_num_sms = 16 + + cfg.train.train_iters = 39735 + cfg.train.global_batch_size = 3072 + cfg.train.micro_batch_size = 2 + cfg.train.manual_gc = False + cfg.train.manual_gc_interval = 0 + + cfg.model.transformer_impl = "transformer_engine" + cfg.model.cuda_graph_impl = "none" + cfg.model.cuda_graph_scope = "full" + cfg.model.cuda_graph_warmup_steps = 3 + + cfg.model.attention_backend = "fused" + cfg.model.moe_router_fusion = False + cfg.model.moe_permute_fusion = True + cfg.model.moe_grouped_gemm = True + cfg.model.cross_entropy_loss_fusion = True + cfg.model.cross_entropy_fusion_impl = "native" + + cfg.model.recompute_granularity = None + cfg.model.recompute_modules = None + cfg.model.fine_grained_activation_offloading = False + cfg.model.offload_modules = None + cfg.model.moe_router_padding_for_fp8 = False + + cfg.optimizer.use_precision_aware_optimizer = False + cfg.optimizer.main_grads_dtype = torch.float32 + cfg.optimizer.main_params_dtype = torch.float32 + cfg.optimizer.exp_avg_dtype = torch.float32 + cfg.optimizer.exp_avg_sq_dtype = torch.float32 + cfg.optimizer.lr = 1.6e-3 + cfg.optimizer.weight_decay = 0.1 + cfg.optimizer.min_lr = 1.6e-5 + cfg.scheduler.lr_warmup_iters = 333 + + cfg.comm_overlap = CommOverlapConfig( + tp_comm_bootstrap_backend="nccl", + tp_comm_overlap=True, + ) + cfg.comm_overlap.delay_wgrad_compute = False + cfg.comm_overlap.overlap_moe_expert_parallel_comm = False + cfg.model.moe_shared_expert_overlap = False + + cfg.checkpoint.save_interval = 200 + cfg.checkpoint.ckpt_assume_constant_structure = True + cfg.checkpoint.dist_ckpt_strictness = "log_all" + + cfg.ddp.overlap_grad_reduce = True + cfg.ddp.overlap_param_gather = True + cfg.ddp.check_for_nan_in_grad = True + cfg.ddp.use_distributed_optimizer = True + + cfg.model.moe_router_force_load_balancing = False + cfg.model.init_method_std = 0.0173 + cfg.model.apply_rope_fusion = False + cfg.model.use_fused_weighted_squared_relu = True + + return cfg + + +def _nemotron_3_nano_finetune_model() -> HybridModelProvider: + """Build the shared Nemotron 3 Nano model contract for SFT and PEFT.""" + return HybridModelProvider( + hybrid_layer_pattern="MEMEM*EMEMEM*EMEMEM*EMEMEM*EMEMEM*EMEMEMEM*EMEMEMEME", + num_layers=52, + hidden_size=2688, + mamba_num_heads=64, + kv_channels=128, + mamba_state_dim=128, + ffn_hidden_size=1856, + num_attention_heads=32, + mamba_head_dim=64, + seq_length=2048, + num_query_groups=2, + num_moe_experts=128, + moe_ffn_hidden_size=1856, + moe_shared_expert_intermediate_size=3712, + moe_router_topk=6, + moe_router_topk_scaling_factor=2.5, + moe_router_num_groups=1, + moe_router_group_topk=1, + mamba_num_groups=8, + make_vocab_size_divisible_by=128, + activation_func=squared_relu, + masked_softmax_fusion=True, + apply_query_key_layer_scaling=False, + persist_layer_norm=True, + attention_softmax_in_fp32=False, + first_last_layers_bf16=True, + is_hybrid_model=True, + moe_aux_loss_coeff=0.0001, + moe_router_score_function="sigmoid", + moe_router_enable_expert_bias=True, + moe_router_load_balancing_type="seq_aux_loss", + moe_router_dtype="fp32", + moe_grouped_gemm=True, + moe_token_dispatcher_type="alltoall", + moe_permute_fusion=True, + moe_shared_expert_overlap=True, + apply_rope_fusion=False, + attention_backend="fused", + init_method_std=0.0173, + use_fused_weighted_squared_relu=True, + calculate_per_token_loss=True, + hf_model_id="nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16", + tensor_model_parallel_size=1, + pipeline_model_parallel_size=1, + pipeline_dtype=torch.bfloat16, + virtual_pipeline_model_parallel_size=None, + context_parallel_size=1, + sequence_parallel=False, + expert_tensor_parallel_size=1, + expert_model_parallel_size=8, + ) + + +def _apply_nemotron_3_nano_finetune_defaults(cfg: ConfigContainer) -> None: + """Apply hardware-neutral SFT/PEFT semantics and safe execution defaults.""" + cfg.model = _nemotron_3_nano_finetune_model() + cfg.model.pipeline_model_parallel_layout = None + + cfg.model.moe_token_dispatcher_type = "flex" + cfg.model.moe_flex_dispatcher_backend = "deepep" + cfg.model.moe_shared_expert_overlap = False + cfg.model.moe_hybridep_num_sms = None + cfg.model.moe_flex_dispatcher_num_sms = None + + cfg.model.transformer_impl = "transformer_engine" + cfg.model.cuda_graph_impl = "none" + clear_cuda_graph_modules(cfg.model) + cfg.model.cuda_graph_warmup_steps = 3 + + cfg.model.attention_backend = "fused" + cfg.model.moe_router_fusion = False + cfg.model.moe_permute_fusion = True + cfg.model.moe_grouped_gemm = True + cfg.model.cross_entropy_loss_fusion = True + cfg.model.cross_entropy_fusion_impl = "native" + + cfg.model.recompute_granularity = None + cfg.model.recompute_modules = None + cfg.model.fine_grained_activation_offloading = False + cfg.model.offload_modules = None + + cfg.optimizer.use_precision_aware_optimizer = False + cfg.optimizer.main_grads_dtype = torch.float32 + cfg.optimizer.main_params_dtype = torch.float32 + cfg.optimizer.exp_avg_dtype = torch.float32 + cfg.optimizer.exp_avg_sq_dtype = torch.float32 + cfg.model.moe_router_padding_for_fp8 = False + cfg.model.moe_router_force_load_balancing = False + + cfg.validation.eval_interval = 500 + if cfg.model.context_parallel_size > 1: + cfg.dataset.offline_packing_specs.pad_seq_to_mult = cfg.model.context_parallel_size * 2 + + cfg.optimizer.adam_beta2 = 0.95 + cfg.optimizer.adam_eps = 1e-8 + cfg.optimizer.weight_decay = 0.1 + cfg.scheduler.start_weight_decay = 0.1 + cfg.scheduler.end_weight_decay = 0.1 + cfg.scheduler.lr_decay_style = "cosine" + + cfg.tokenizer.tokenizer_model = "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16" + + cfg.checkpoint.save_interval = 200 + cfg.checkpoint.ckpt_format = "torch_dist" + cfg.checkpoint.dist_ckpt_strictness = "log_all" + cfg.checkpoint.ckpt_assume_constant_structure = True + + cfg.logger.log_interval = 10 + cfg.logger.log_timers_to_tensorboard = False + cfg.rng.seed = 1234 + + cfg.ddp.check_for_nan_in_grad = True + cfg.ddp.grad_reduce_in_fp32 = True + cfg.ddp.overlap_grad_reduce = True + cfg.ddp.overlap_param_gather = True + cfg.ddp.use_distributed_optimizer = True + + +def _nemotron_3_nano_sft_reference_config() -> ConfigContainer: + """Build the hardware-neutral Nemotron 3 Nano SFT contract.""" + cfg = _sft_common() + _apply_nemotron_3_nano_finetune_defaults(cfg) + return cfg + + +def _nemotron_3_nano_peft_reference_config( + peft_scheme: str | PEFT = "lora", +) -> ConfigContainer: + """Build the hardware-neutral Nemotron 3 Nano PEFT contract.""" + cfg = _peft_common() + _apply_nemotron_3_nano_finetune_defaults(cfg) + + target_modules = ["linear_qkv", "linear_proj", "linear_fc1", "linear_fc2", "in_proj", "out_proj"] + if isinstance(peft_scheme, str) and peft_scheme.lower() in ["lora", "dora"]: + cfg.peft = default_peft_config(peft_scheme, target_modules=target_modules) + elif isinstance(peft_scheme, PEFT): + cfg.peft = peft_scheme + else: + cfg.peft = LoRA( + target_modules=target_modules, + dim=32, + alpha=32, + dropout=0.0, + dropout_position="pre", + lora_A_init_method="xavier", + lora_B_init_method="zero", + ) + + return cfg diff --git a/src/megatron/bridge/recipes/nemotronh/gb200/__init__.py b/src/megatron/bridge/recipes/nemotronh/gb200/__init__.py index 2fc637f14f..ef6934d908 100644 --- a/src/megatron/bridge/recipes/nemotronh/gb200/__init__.py +++ b/src/megatron/bridge/recipes/nemotronh/gb200/__init__.py @@ -17,7 +17,9 @@ nemotron_3_5_nano_pretrain_8k_fsdp_config, nemotron_3_5_nano_sft_openmathinstruct2_packed_tp1_config, nemotron_3_nano_gb200_pretrain_config, + nemotron_3_nano_peft_8gpu_gb200_bf16_config, nemotron_3_nano_pretrain_8gpu_gb200_bf16_config, + nemotron_3_nano_sft_8gpu_gb200_bf16_config, ) @@ -26,5 +28,7 @@ "nemotron_3_5_nano_pretrain_8k_fsdp_config", "nemotron_3_5_nano_sft_openmathinstruct2_packed_tp1_config", "nemotron_3_nano_gb200_pretrain_config", + "nemotron_3_nano_peft_8gpu_gb200_bf16_config", "nemotron_3_nano_pretrain_8gpu_gb200_bf16_config", + "nemotron_3_nano_sft_8gpu_gb200_bf16_config", ] diff --git a/src/megatron/bridge/recipes/nemotronh/gb200/nemotron_3_nano.py b/src/megatron/bridge/recipes/nemotronh/gb200/nemotron_3_nano.py index 136a7ded2f..b6c8203f1a 100644 --- a/src/megatron/bridge/recipes/nemotronh/gb200/nemotron_3_nano.py +++ b/src/megatron/bridge/recipes/nemotronh/gb200/nemotron_3_nano.py @@ -12,12 +12,17 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""GB200 training recipes for Nemotron 3 Nano.""" +"""GB200 library recipes for Nemotron 3 and 3.5 Nano.""" import torch from megatron.bridge import AutoBridge +from megatron.bridge.peft.base import PEFT from megatron.bridge.recipes.common import _pretrain_common +from megatron.bridge.recipes.nemotronh._nemotron_3_nano import ( + _nemotron_3_nano_peft_reference_config, + _nemotron_3_nano_sft_reference_config, +) from megatron.bridge.recipes.nemotronh.h100.nemotron_3_nano import ( nemotron_3_5_nano_pretrain_config, nemotron_3_5_nano_sft_openmathinstruct2_packed_config, @@ -70,16 +75,16 @@ def nemotron_3_nano_pretrain_8gpu_gb200_bf16_config() -> ConfigContainer: cfg.model.moe_token_dispatcher_type = "flex" cfg.model.moe_flex_dispatcher_backend = "hybridep" - cfg.model.moe_flex_dispatcher_num_sms = None - cfg.model.moe_hybridep_num_sms = 16 + cfg.model.moe_flex_dispatcher_num_sms = 16 + cfg.model.moe_hybridep_num_sms = None cfg.model.moe_shared_expert_overlap = False cfg.model.moe_router_force_load_balancing = False cfg.train.train_iters = 39735 - cfg.train.global_batch_size = 512 + cfg.train.global_batch_size = 3072 cfg.train.micro_batch_size = 2 - cfg.train.manual_gc = True - cfg.train.manual_gc_interval = 100 + cfg.train.manual_gc = False + cfg.train.manual_gc_interval = 0 cfg.model.transformer_impl = "transformer_engine" @@ -120,40 +125,24 @@ def nemotron_3_nano_pretrain_8gpu_gb200_bf16_config() -> ConfigContainer: cfg.comm_overlap = CommOverlapConfig( tp_comm_bootstrap_backend="nccl", - tp_comm_overlap=True, + tp_comm_overlap=False, ) cfg.comm_overlap.delay_wgrad_compute = False cfg.comm_overlap.overlap_moe_expert_parallel_comm = False cfg.checkpoint.save_interval = 200 - cfg.checkpoint.async_save = False cfg.checkpoint.ckpt_assume_constant_structure = True cfg.checkpoint.dist_ckpt_strictness = "log_all" cfg.ddp.overlap_grad_reduce = True cfg.ddp.overlap_param_gather = True - cfg.ddp.check_for_nan_in_grad = True - cfg.ddp.check_for_large_grads = True + cfg.ddp.check_for_nan_in_grad = False cfg.ddp.use_distributed_optimizer = True cfg.ddp.grad_reduce_in_fp32 = False - cfg.rerun_state_machine.check_for_nan_in_loss = True cfg.model.init_method_std = 0.0173 cfg.model.use_fused_weighted_squared_relu = True - cfg.env_vars = { - **COMMON_RECIPE_ENV_VARS, - "CUDA_DEVICE_MAX_CONNECTIONS": 32, - "NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN": 8, - "NUM_OF_TOKENS_PER_CHUNK_COMBINE_API": 128, - "NVLINK_DOMAIN_SIZE": 72, - "USE_MNNVL": 1, - "NVTE_BWD_LAYERNORM_SM_MARGIN": 20, - "NVTE_FWD_LAYERNORM_SM_MARGIN": 20, - "NVTE_NORM_BWD_USE_CUDNN": 1, - "NVTE_NORM_FWD_USE_CUDNN": 1, - } - return cfg @@ -238,6 +227,56 @@ def nemotron_3_5_nano_sft_openmathinstruct2_packed_tp1_config() -> ConfigContain return cfg +def _apply_gb200_finetune_execution_config(cfg: ConfigContainer) -> None: + """Apply safe GB200 packed-finetuning execution settings.""" + 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.expert_model_parallel_size = 8 + + # DeepEP is unsupported on compute capability 10, while packed HybridEP + # lacks matching workload evidence. Use the correctness-first fallback. + cfg.model.moe_token_dispatcher_type = "alltoall" + cfg.model.moe_flex_dispatcher_backend = None + cfg.model.moe_shared_expert_overlap = False + cfg.model.moe_hybridep_num_sms = None + cfg.model.moe_flex_dispatcher_num_sms = None + cfg.model.moe_router_force_load_balancing = False + + +def nemotron_3_nano_sft_8gpu_gb200_bf16_config() -> ConfigContainer: + """Return the Nemotron 3 Nano SFT config for eight GB200 GPUs. + + Packed SFT remains eager because CUDA graphs require static input shapes. + Optimizer, schedule, data, and routing semantics remain unchanged. + + Returns: + GB200 BF16 SFT configuration. + """ + cfg = _nemotron_3_nano_sft_reference_config() + _apply_gb200_finetune_execution_config(cfg) + return cfg + + +def nemotron_3_nano_peft_8gpu_gb200_bf16_config( + peft_scheme: str | PEFT = "lora", +) -> ConfigContainer: + """Return the Nemotron 3 Nano PEFT config for eight GB200 GPUs. + + Args: + peft_scheme: PEFT scheme, or a custom PEFT instance. + + Returns: + GB200 BF16 PEFT configuration. + """ + cfg = _nemotron_3_nano_peft_reference_config(peft_scheme=peft_scheme) + _apply_gb200_finetune_execution_config(cfg) + return cfg + + # NeMo-CI appends ``_pretrain_config`` to MODEL_RECIPE_NAME. This explicit # alias lets the GB200 release case select the hardware recipe without changing # the legacy ``nemotron_3_nano_pretrain_config`` default. @@ -249,5 +288,7 @@ def nemotron_3_5_nano_sft_openmathinstruct2_packed_tp1_config() -> ConfigContain "nemotron_3_5_nano_pretrain_8k_fsdp_config", "nemotron_3_5_nano_sft_openmathinstruct2_packed_tp1_config", "nemotron_3_nano_gb200_pretrain_config", + "nemotron_3_nano_peft_8gpu_gb200_bf16_config", "nemotron_3_nano_pretrain_8gpu_gb200_bf16_config", + "nemotron_3_nano_sft_8gpu_gb200_bf16_config", ] diff --git a/src/megatron/bridge/recipes/nemotronh/h100/nemotron_3_nano.py b/src/megatron/bridge/recipes/nemotronh/h100/nemotron_3_nano.py index 15651cef5e..51bc58cf2f 100644 --- a/src/megatron/bridge/recipes/nemotronh/h100/nemotron_3_nano.py +++ b/src/megatron/bridge/recipes/nemotronh/h100/nemotron_3_nano.py @@ -12,20 +12,16 @@ # See the License for the specific language governing permissions and # limitations under the License. +"""H100 library recipes for Nemotron 3 and 3.5 Nano.""" -from typing import cast - -import torch -from megatron.core.activations import squared_relu - -from megatron.bridge import AutoBridge -from megatron.bridge.models.hybrid.hybrid_provider import HybridModelProvider from megatron.bridge.peft.base import PEFT -from megatron.bridge.peft.lora import LoRA -from megatron.bridge.recipes.common import _peft_common, _pretrain_common, _sft_common -from megatron.bridge.recipes.utils.dataset_utils import default_openmathinstruct2_config, default_peft_config +from megatron.bridge.recipes.nemotronh._nemotron_3_nano import ( + _nemotron_3_nano_peft_reference_config, + _nemotron_3_nano_pretrain_reference_config, + _nemotron_3_nano_sft_reference_config, +) +from megatron.bridge.recipes.utils.dataset_utils import default_openmathinstruct2_config from megatron.bridge.recipes.utils.environment_utils import COMMON_RECIPE_ENV_VARS -from megatron.bridge.training.comm_overlap import CommOverlapConfig from megatron.bridge.training.config import ConfigContainer from megatron.bridge.training.mixed_precision import get_mixed_precision_config from megatron.bridge.utils.cuda_graph import set_cuda_graph_modules @@ -37,232 +33,89 @@ _OPENMATHINSTRUCT2_REVISION = "469216e3f46f4dacf476b382e192485ea51a143e" # pragma: allowlist secret -def _nemotron_3_nano_finetune_model() -> HybridModelProvider: - """Build the Nemotron 3 Nano finetuning provider.""" - model = cast( - HybridModelProvider, - AutoBridge.from_hf_pretrained(_NEMOTRON_3_NANO_MODEL_ID).to_megatron_provider(load_weights=False), - ) - - model.seq_length = 2048 - model.apply_rope_fusion = False - model.attention_backend = "fused" - model.init_method_std = 0.0173 - model.use_fused_weighted_squared_relu = True - model.calculate_per_token_loss = True - model.tensor_model_parallel_size = 1 - model.pipeline_model_parallel_size = 1 - model.pipeline_dtype = torch.bfloat16 - model.virtual_pipeline_model_parallel_size = None - model.context_parallel_size = 1 - model.sequence_parallel = False - model.expert_tensor_parallel_size = 1 - model.expert_model_parallel_size = 8 - return model - - def nemotron_3_nano_pretrain_8gpu_h100_bf16_config() -> ConfigContainer: - """Return a pre-training config for Nemotron 3 Nano (30B-A3B MoE). - - This is a MoE (Mixture of Experts) model with the following default parallelism: - - TP=1, PP=1, ETP=1, EP=8, SP=False - - HybridEP enabled for MoE token dispatch + """Return the Nemotron 3 Nano pretraining config for eight H100 GPUs. + + TP8 retains the perf recipe's PP1/CP1/EP8/ETP1 HybridEP topology, + selective recompute, and native vocab-parallel cross entropy. TP + communication overlap is disabled because its persistent userbuffers + exhaust checkpoint-restore headroom on 80 GB H100s. The compiled native + cross-entropy wrapper is disabled because its temporary workspace does not + fit after FP32 optimizer-state allocation; the underlying native loss is + unchanged. Unused CUDA cache is released after each optimizer step so the + first lazy MoE metric collective can allocate after checkpoint resume. + Validation uses microbatch one without changing its global batch. CUDA + graphs remain disabled to preserve general-training headroom. Returns: - ConfigContainer: Pre-training configuration for Nemotron 3 Nano. + H100 BF16 pretraining configuration. """ - cfg = _pretrain_common() - - # Model Configuration (MoE) - cfg.model = HybridModelProvider( - # Architecture (Nemotron 3 Nano 30B-A3B) - hybrid_layer_pattern="MEMEM*EMEMEM*EMEMEM*EMEMEM*EMEMEM*EMEMEMEM*EMEMEMEME", - num_layers=52, - hidden_size=2688, - mamba_num_heads=64, - kv_channels=128, - mamba_state_dim=128, - ffn_hidden_size=1856, - num_attention_heads=32, - mamba_head_dim=64, - seq_length=8192, - num_query_groups=2, - # MoE - num_moe_experts=128, - moe_ffn_hidden_size=1856, - moe_shared_expert_intermediate_size=3712, - moe_router_topk=6, - moe_router_topk_scaling_factor=2.5, - moe_router_num_groups=1, - moe_router_group_topk=1, - # NemotronH base - mamba_num_groups=8, - make_vocab_size_divisible_by=128, - activation_func=squared_relu, - masked_softmax_fusion=True, - apply_query_key_layer_scaling=False, - persist_layer_norm=True, - attention_softmax_in_fp32=False, - first_last_layers_bf16=True, - is_hybrid_model=True, - moe_aux_loss_coeff=0.0001, - moe_router_score_function="sigmoid", - moe_router_enable_expert_bias=True, - moe_router_load_balancing_type="seq_aux_loss", - moe_router_dtype="fp32", - moe_grouped_gemm=True, - moe_token_dispatcher_type="alltoall", - moe_permute_fusion=True, - moe_shared_expert_overlap=True, - tensor_model_parallel_size=1, - pipeline_model_parallel_size=1, - pipeline_dtype=torch.bfloat16, - virtual_pipeline_model_parallel_size=None, - context_parallel_size=1, - sequence_parallel=False, - expert_tensor_parallel_size=1, - expert_model_parallel_size=8, - ) - # Tokenizer (--tokenizer-model) - cfg.tokenizer.tokenizer_model = _NEMOTRON_3_NANO_MODEL_ID + cfg = _nemotron_3_nano_pretrain_reference_config() - # Dataset Configuration - cfg.dataset.seq_length = 8192 - cfg.dataset.blend = None # Pass the path to the dataset here if not using mock data, along with weight. Ex: (["path/to/data1"], 0.2), [("path/to/data2", 0.8)] - cfg.dataset.num_workers = 8 - cfg.dataset.mmap_bin_files = False - - # Parallelism Settings (MoE-specific) - cfg.model.pipeline_model_parallel_layout = None + cfg.model.tensor_model_parallel_size = 8 + 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.expert_model_parallel_size = 8 - # MoE Token Dispatcher Settings cfg.model.moe_token_dispatcher_type = "flex" cfg.model.moe_flex_dispatcher_backend = "hybridep" - cfg.model.moe_hybridep_num_sms = 16 - - # Training Configuration - cfg.train.train_iters = 39735 - cfg.train.global_batch_size = 1024 - cfg.train.micro_batch_size = 1 - cfg.train.manual_gc = True - cfg.train.manual_gc_interval = 100 - - # Transformer Engine (TE) - cfg.model.transformer_impl = "transformer_engine" - - # CUDA Graph - cfg.model.cuda_graph_impl = "transformer_engine" - set_cuda_graph_modules(cfg.model, ["attn", "mamba"]) - cfg.model.cuda_graph_warmup_steps = 3 - cfg.model.use_te_rng_tracker = True - cfg.rng.te_rng_tracker = True - - # Kernel Selections - cfg.model.attention_backend = "fused" - cfg.model.moe_router_fusion = False - cfg.model.moe_permute_fusion = True - cfg.model.moe_grouped_gemm = True - cfg.model.cross_entropy_loss_fusion = True - cfg.model.cross_entropy_fusion_impl = "native" + cfg.model.moe_hybridep_num_sms = None + cfg.model.moe_flex_dispatcher_num_sms = 16 + cfg.model.moe_shared_expert_overlap = False - # Memory Saving (recompute & offloading) cfg.model.recompute_granularity = "selective" - cfg.model.recompute_modules = ["moe", "layernorm"] cfg.model.recompute_method = None cfg.model.recompute_num_layers = None - cfg.model.fine_grained_activation_offloading = False - cfg.model.offload_modules = None - - # ========================================================================= - # FP8 & MXFP8 (Mixed Precision Settings) - # ========================================================================= - # Note: mixed_precision="bf16_mixed" is set in _pretrain_common as default - # FP8 settings (disabled by default, uncomment to enable) - # cfg.mixed_precision.fp8_recipe = "tensorwise" - # cfg.mixed_precision.fp8 = None - # cfg.mixed_precision.fp8_param_gather = False - # cfg.mixed_precision.reuse_grad_buf_for_mxfp8_param_ag = False - cfg.model.moe_router_padding_for_fp8 = False - - # Optimizer Precision Settings - # Match the measured BF16 performance recipe. The 16-GPU convergence - # workload provides DP=2, so distributed optimizer state remains sharded - # while retaining full FP32 optimizer precision. - cfg.optimizer.use_precision_aware_optimizer = False - cfg.optimizer.main_grads_dtype = torch.float32 - cfg.optimizer.main_params_dtype = torch.float32 - cfg.optimizer.exp_avg_dtype = torch.float32 - cfg.optimizer.exp_avg_sq_dtype = torch.float32 - cfg.mixed_precision = get_mixed_precision_config(cfg.mixed_precision) - cfg.mixed_precision.grad_reduce_in_fp32 = False - - # Optimizer hyperparameters - cfg.optimizer.lr = 1.6e-3 - cfg.optimizer.weight_decay = 0.1 - cfg.optimizer.min_lr = 1.6e-5 - cfg.scheduler.lr_warmup_iters = 333 - - # Communication Overlap - cfg.comm_overlap = CommOverlapConfig( - tp_comm_bootstrap_backend="nccl", - tp_comm_overlap=True, - ) - cfg.comm_overlap.delay_wgrad_compute = False - cfg.comm_overlap.overlap_moe_expert_parallel_comm = False - cfg.model.moe_shared_expert_overlap = False - - # Checkpoint Configuration - # Paths are set in _pretrain_common by default. Override here if needed: - # cfg.checkpoint.load = "path/to/load" - # cfg.checkpoint.save = "path/to/save" - cfg.checkpoint.save_interval = 200 - cfg.checkpoint.async_save = False - cfg.checkpoint.ckpt_assume_constant_structure = True - cfg.checkpoint.dist_ckpt_strictness = "log_all" - - # DDP Configuration - cfg.ddp.overlap_grad_reduce = True - cfg.ddp.overlap_param_gather = True - cfg.ddp.check_for_nan_in_grad = True - cfg.ddp.check_for_large_grads = True - cfg.ddp.use_distributed_optimizer = True - cfg.ddp.grad_reduce_in_fp32 = False - cfg.rerun_state_machine.check_for_nan_in_loss = True - - # MoE Force Load Balancing - cfg.model.moe_router_force_load_balancing = False + cfg.model.recompute_modules = ["moe", "layernorm"] + cfg.model.cross_entropy_loss_fusion = False - cfg.model.init_method_std = 0.0173 - cfg.model.apply_rope_fusion = True - cfg.model.use_fused_weighted_squared_relu = True + cfg.comm_overlap.tp_comm_overlap = False + cfg.train.empty_unused_memory_level = 2 + cfg.validation.eval_micro_batch_size = 1 - # Keep the complete process environment visible on the recipe. cfg.env_vars = { **COMMON_RECIPE_ENV_VARS, - "CUDA_DEVICE_MAX_CONNECTIONS": 32, - "NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN": 8, - "NUM_OF_TOKENS_PER_CHUNK_COMBINE_API": 128, - "NVLINK_DOMAIN_SIZE": 8, - "USE_MNNVL": 0, - "NVTE_BWD_LAYERNORM_SM_MARGIN": 20, - "NVTE_FWD_LAYERNORM_SM_MARGIN": 20, - "NVTE_NORM_BWD_USE_CUDNN": 1, - "NVTE_NORM_FWD_USE_CUDNN": 1, } return cfg def nemotron_3_5_nano_pretrain_config() -> ConfigContainer: """Return the Nemotron 3.5 Nano BF16 pretraining config.""" - cfg = nemotron_3_nano_pretrain_8gpu_h100_bf16_config() + cfg = _nemotron_3_nano_pretrain_reference_config() + + # Preserve the current Nemotron 3.5 H100 execution policy independently + # from the memory-safe eight-GPU Nemotron 3 recipe above. cfg.train.global_batch_size = 512 + cfg.train.micro_batch_size = 1 + cfg.train.manual_gc = True + cfg.train.manual_gc_interval = 100 + cfg.model.tensor_model_parallel_size = 1 + cfg.model.sequence_parallel = False # Split the 8K sequence across two ranks so each MTP head materializes only # half of its vocabulary-loss workspace on an 80-GiB H100. P2P retains the # fused-attention path for this model's grouped-query layout. cfg.model.context_parallel_size = 2 cfg.model.cp_comm_type = "p2p" + cfg.model.moe_flex_dispatcher_backend = "hybridep" + cfg.model.moe_hybridep_num_sms = 16 + cfg.model.moe_flex_dispatcher_num_sms = None + cfg.model.cross_entropy_loss_fusion = True + cfg.model.apply_rope_fusion = True + cfg.model.cuda_graph_impl = "transformer_engine" + cfg.model.recompute_granularity = "selective" cfg.model.recompute_modules = ["moe", "layernorm", "core_attn"] set_cuda_graph_modules(cfg.model, ["mamba"]) + cfg.model.use_te_rng_tracker = True + cfg.rng.te_rng_tracker = True + cfg.mixed_precision = get_mixed_precision_config(cfg.mixed_precision) + cfg.mixed_precision.grad_reduce_in_fp32 = False + cfg.ddp.grad_reduce_in_fp32 = False + cfg.ddp.check_for_large_grads = True + cfg.rerun_state_machine.check_for_nan_in_loss = True + cfg.checkpoint.async_save = False cfg.model.mtp_num_layers = 2 cfg.model.mtp_hybrid_override_pattern = "*E" cfg.model.mtp_use_repeated_layer = True @@ -270,135 +123,53 @@ def nemotron_3_5_nano_pretrain_config() -> ConfigContainer: cfg.model.mtp_loss_scaling_factor = 0.3 cfg.model.hf_model_id = _NEMOTRON_3_5_NANO_MODEL_ID cfg.tokenizer.tokenizer_model = _NEMOTRON_3_5_NANO_MODEL_ID + cfg.env_vars = { + **COMMON_RECIPE_ENV_VARS, + "CUDA_DEVICE_MAX_CONNECTIONS": 32, + "NUM_OF_HYBRID_EP_RANKS_PER_NVLINK_DOMAIN": 8, + "NUM_OF_TOKENS_PER_CHUNK_COMBINE_API": 128, + "NVLINK_DOMAIN_SIZE": 8, + "USE_MNNVL": 0, + "NVTE_BWD_LAYERNORM_SM_MARGIN": 20, + "NVTE_FWD_LAYERNORM_SM_MARGIN": 20, + "NVTE_NORM_BWD_USE_CUDNN": 1, + "NVTE_NORM_FWD_USE_CUDNN": 1, + } return cfg -# ============================================================================= -# SFT Config -# ============================================================================= +def _apply_h100_finetune_execution_config(cfg: ConfigContainer) -> None: + """Apply the evidenced H100 packed-finetuning execution contract.""" + 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.expert_model_parallel_size = 8 + + cfg.model.moe_token_dispatcher_type = "flex" + cfg.model.moe_flex_dispatcher_backend = "deepep" + cfg.model.moe_shared_expert_overlap = False + cfg.model.moe_hybridep_num_sms = None + cfg.model.moe_flex_dispatcher_num_sms = 16 + cfg.model.moe_router_force_load_balancing = False def nemotron_3_nano_sft_8gpu_h100_bf16_config() -> ConfigContainer: - """Return a full SFT config for Nemotron 3 Nano. + """Return the Nemotron 3 Nano SFT config for eight H100 GPUs. - Default parallelism: TP=1, PP=1, EP=8, SP=False + Packed SFT retains the established DeepEP dispatcher and eager execution. + TP4 leaves room for full optimizer state and checkpointing. Returns: - ConfigContainer with all settings pre-configured for Nemotron Nano SFT. + H100 BF16 SFT configuration. """ - cfg = _sft_common() - - cfg.model = _nemotron_3_nano_finetune_model() - - # Parallelism settings - cfg.model.pipeline_model_parallel_layout = None - - # Sequence length - cfg.model.seq_length = 2048 - - # DeePEP settings - set to True to enable DeePEP (enabled by default for Nemotron) - enable_deepep = True - if enable_deepep: - cfg.model.moe_token_dispatcher_type = "flex" - cfg.model.moe_flex_dispatcher_backend = "deepep" - cfg.model.moe_shared_expert_overlap = False - else: - cfg.model.moe_token_dispatcher_type = "alltoall" - cfg.model.moe_flex_dispatcher_backend = None - cfg.model.moe_shared_expert_overlap = True - - cfg.model.moe_hybridep_num_sms = 16 - - # TE (Transformer Engine) - cfg.model.transformer_impl = "transformer_engine" - - # CUDA Graph - cfg.model.cuda_graph_impl = "none" - cfg.model.cuda_graph_scope = "full" - cfg.model.cuda_graph_warmup_steps = 3 - - # Kernel selections - cfg.model.attention_backend = "fused" - cfg.model.moe_router_fusion = False - cfg.model.moe_permute_fusion = True - cfg.model.moe_grouped_gemm = True - cfg.model.cross_entropy_loss_fusion = True - cfg.model.cross_entropy_fusion_impl = "native" - - # Memory saving (recompute & offloading) - cfg.model.recompute_granularity = None - cfg.model.recompute_modules = None - cfg.model.fine_grained_activation_offloading = False - cfg.model.offload_modules = None - - # FP8 & MXFP8 settings - # Note: mixed_precision="bf16_mixed" is set as default - # These are defaults for FP8, enable them if using FP8 - FP8 is not enabled by default - # cfg.mixed_precision.fp8_recipe = "tensorwise" - # cfg.mixed_precision.fp8 = None - # cfg.mixed_precision.fp8_param_gather = False - # cfg.mixed_precision.reuse_grad_buf_for_mxfp8_param_ag = False - cfg.optimizer.use_precision_aware_optimizer = False - cfg.optimizer.main_grads_dtype = torch.float32 - cfg.optimizer.main_params_dtype = torch.float32 - cfg.optimizer.exp_avg_dtype = torch.float32 - cfg.optimizer.exp_avg_sq_dtype = torch.float32 - cfg.model.moe_router_padding_for_fp8 = False - - # MoE Force Load Balancing - cfg.model.moe_router_force_load_balancing = False + cfg = _nemotron_3_nano_sft_reference_config() + _apply_h100_finetune_execution_config(cfg) + cfg.model.tensor_model_parallel_size = 4 + cfg.model.sequence_parallel = True - # Training config overrides - cfg.validation.eval_interval = 500 - - # Dataset config - enable_offline_packing=True by default (from _sft_common), seq_length=2048 - # _sft_common already sets seq_length=2048 and enable_offline_packing=True - # Adjust pad_seq_to_mult for context parallelism - if cfg.model.context_parallel_size > 1: - cfg.dataset.offline_packing_specs.pad_seq_to_mult = cfg.model.context_parallel_size * 2 - - # Optimizer overrides - Nemotron uses specific optimizer settings - cfg.optimizer.adam_beta2 = 0.95 - cfg.optimizer.adam_eps = 1e-8 - cfg.optimizer.weight_decay = 0.1 - cfg.scheduler.start_weight_decay = 0.1 - cfg.scheduler.end_weight_decay = 0.1 - cfg.scheduler.lr_decay_style = "cosine" - - # Tokenizer - cfg.tokenizer.tokenizer_model = _NEMOTRON_3_NANO_MODEL_ID - - # Checkpoint config overrides - cfg.checkpoint.save_interval = 200 - cfg.checkpoint.ckpt_format = "torch_dist" - cfg.checkpoint.dist_ckpt_strictness = "log_all" - cfg.checkpoint.ckpt_assume_constant_structure = True - # Uncomment below if using a pretrained checkpoint and provide path to the directory containing pretrained model for finetuning - # cfg.checkpoint.pretrained_checkpoint = "/path/to/checkpoint" - - # Logger config - cfg.logger.log_interval = 10 - cfg.logger.log_timers_to_tensorboard = False - - # RNG config - Nemotron uses seed 1234 - cfg.rng.seed = 1234 - - # DDP config - cfg.ddp.check_for_nan_in_grad = True - cfg.ddp.grad_reduce_in_fp32 = True - cfg.ddp.overlap_grad_reduce = True - cfg.ddp.overlap_param_gather = True - cfg.ddp.use_distributed_optimizer = True - - # Communication overlap settings(default None, can pass CommOverlapConfig for advanced overlap), uncomment to enable - # cfg.comm_overlap = CommOverlapConfig( - # tp_comm_bootstrap_backend="nccl", - # tp_comm_overlap=True, - # ) - # cfg.comm_overlap.delay_wgrad_compute = False - # cfg.comm_overlap.overlap_moe_expert_parallel_comm = False - - # Keep the complete process environment visible on the recipe. cfg.env_vars = { **COMMON_RECIPE_ENV_VARS, } @@ -407,7 +178,11 @@ def nemotron_3_nano_sft_8gpu_h100_bf16_config() -> ConfigContainer: def nemotron_3_5_nano_sft_config() -> ConfigContainer: """Return a full SFT config for Nemotron 3.5 Nano.""" - cfg = nemotron_3_nano_sft_8gpu_h100_bf16_config() + cfg = _nemotron_3_nano_sft_reference_config() + _apply_h100_finetune_execution_config(cfg) + cfg.env_vars = { + **COMMON_RECIPE_ENV_VARS, + } cfg.model.mtp_num_layers = 2 cfg.model.mtp_hybrid_override_pattern = "*E" cfg.model.mtp_use_repeated_layer = True @@ -476,163 +251,20 @@ def nemotron_3_5_nano_sft_openmathinstruct2_packed_config() -> ConfigContainer: return cfg -# ============================================================================= -# PEFT Config -# ============================================================================= - - -def _nemotron_3_nano_peft_8gpu_h100_bf16_config( +def nemotron_3_nano_peft_8gpu_h100_bf16_config( peft_scheme: str | PEFT = "lora", ) -> ConfigContainer: - """Build a PEFT config for Nemotron 3 Nano. - - Default parallelism: TP=1, PP=1, EP=8, SP=False + """Return the Nemotron 3 Nano PEFT config for eight H100 GPUs. Args: - peft_scheme: PEFT scheme - "lora", "dora", or a custom PEFT instance. + peft_scheme: PEFT scheme, or a custom PEFT instance. Returns: - ConfigContainer with all settings pre-configured for Nemotron Nano PEFT. + H100 BF16 PEFT configuration. """ - cfg = _peft_common() - - cfg.model = _nemotron_3_nano_finetune_model() - - # Parallelism settings - cfg.model.pipeline_model_parallel_layout = None - - # Sequence length - cfg.model.seq_length = 2048 - - # DeePEP settings - set to True to enable DeePEP (enabled by default for Nemotron) - enable_deepep = True - if enable_deepep: - cfg.model.moe_token_dispatcher_type = "flex" - cfg.model.moe_flex_dispatcher_backend = "deepep" - cfg.model.moe_shared_expert_overlap = False - else: - cfg.model.moe_token_dispatcher_type = "alltoall" - cfg.model.moe_flex_dispatcher_backend = None - cfg.model.moe_shared_expert_overlap = True - - cfg.model.moe_hybridep_num_sms = 16 - - # TE (Transformer Engine) - cfg.model.transformer_impl = "transformer_engine" - - # CUDA Graph - cfg.model.cuda_graph_impl = "none" - cfg.model.cuda_graph_scope = "full" - cfg.model.cuda_graph_warmup_steps = 3 - - # Kernel selections - cfg.model.attention_backend = "fused" - cfg.model.moe_router_fusion = False - cfg.model.moe_permute_fusion = True - cfg.model.moe_grouped_gemm = True - cfg.model.cross_entropy_loss_fusion = True - cfg.model.cross_entropy_fusion_impl = "native" - - # Memory saving - cfg.model.recompute_granularity = None - cfg.model.recompute_modules = None - cfg.model.fine_grained_activation_offloading = False - cfg.model.offload_modules = None - - # FP8 & MXFP8 settings - # These are defaults for FP8, enable them if using FP8 - FP8 is not enabled by default - # cfg.mixed_precision.fp8_recipe = "tensorwise" - # cfg.mixed_precision.fp8 = None - # cfg.mixed_precision.fp8_param_gather = False - # cfg.mixed_precision.reuse_grad_buf_for_mxfp8_param_ag = False - cfg.optimizer.use_precision_aware_optimizer = False - cfg.optimizer.main_grads_dtype = torch.float32 - cfg.optimizer.main_params_dtype = torch.float32 - cfg.optimizer.exp_avg_dtype = torch.float32 - cfg.optimizer.exp_avg_sq_dtype = torch.float32 - cfg.model.moe_router_padding_for_fp8 = False - - # MoE Force Load Balancing - cfg.model.moe_router_force_load_balancing = False - - # PEFT config - Nemotron uses Mamba-specific target modules - mamba_target_modules = ["linear_qkv", "linear_proj", "linear_fc1", "linear_fc2", "in_proj", "out_proj"] - if isinstance(peft_scheme, str) and peft_scheme.lower() in ["lora", "dora"]: - cfg.peft = default_peft_config(peft_scheme, target_modules=mamba_target_modules) - elif isinstance(peft_scheme, PEFT): - cfg.peft = peft_scheme - else: - # Default to LoRA with Mamba target modules - cfg.peft = LoRA( - target_modules=mamba_target_modules, - dim=32, - alpha=32, - dropout=0.0, - dropout_position="pre", - lora_A_init_method="xavier", - lora_B_init_method="zero", - ) - - # Training config overrides - cfg.validation.eval_interval = 500 - - # Dataset config - enable_offline_packing=True by default (from _peft_common), seq_length=2048 - # _peft_common already sets seq_length=2048 and enable_offline_packing=True - # Adjust pad_seq_to_mult for context parallelism - if cfg.model.context_parallel_size > 1: - cfg.dataset.offline_packing_specs.pad_seq_to_mult = cfg.model.context_parallel_size * 2 - - # Optimizer overrides - cfg.optimizer.adam_beta2 = 0.95 - cfg.optimizer.adam_eps = 1e-8 - cfg.optimizer.weight_decay = 0.1 - cfg.scheduler.start_weight_decay = 0.1 - cfg.scheduler.end_weight_decay = 0.1 - cfg.scheduler.lr_decay_style = "cosine" - - # Tokenizer - cfg.tokenizer.tokenizer_model = _NEMOTRON_3_NANO_MODEL_ID - - # Checkpoint config overrides - cfg.checkpoint.save_interval = 200 - cfg.checkpoint.ckpt_format = "torch_dist" - cfg.checkpoint.dist_ckpt_strictness = "log_all" - cfg.checkpoint.ckpt_assume_constant_structure = True - # Uncomment below if using a pretrained checkpoint and provide path to the directory containing pretrained model for finetuning - # cfg.checkpoint.pretrained_checkpoint = "/path/to/checkpoint" - - # Logger config - cfg.logger.log_interval = 10 - cfg.logger.log_timers_to_tensorboard = False - - # RNG config - Nemotron uses seed 1234 - cfg.rng.seed = 1234 - - # DDP config - cfg.ddp.check_for_nan_in_grad = True - cfg.ddp.grad_reduce_in_fp32 = True - cfg.ddp.overlap_grad_reduce = True - cfg.ddp.overlap_param_gather = True - cfg.ddp.use_distributed_optimizer = True - - # Communication overlap settings(default None, can pass CommOverlapConfig for advanced overlap), uncomment to enable - # cfg.comm_overlap = CommOverlapConfig( - # tp_comm_bootstrap_backend="nccl", - # tp_comm_overlap=True, - # ) - # cfg.comm_overlap.delay_wgrad_compute = False - # cfg.comm_overlap.overlap_moe_expert_parallel_comm = False - - # Keep the complete process environment visible on the recipe. - cfg.env_vars = { - **COMMON_RECIPE_ENV_VARS, - } - return cfg - + cfg = _nemotron_3_nano_peft_reference_config(peft_scheme=peft_scheme) + _apply_h100_finetune_execution_config(cfg) -def nemotron_3_nano_peft_8gpu_h100_bf16_config(peft_scheme: str | PEFT = "lora") -> ConfigContainer: - """Return a PEFT config for Nemotron 3 Nano.""" - cfg = _nemotron_3_nano_peft_8gpu_h100_bf16_config(peft_scheme) cfg.env_vars = { **COMMON_RECIPE_ENV_VARS, } diff --git a/tests/functional_tests/test_groups/recipes/test_nemotronh_recipes_pretrain.py b/tests/functional_tests/test_groups/recipes/test_nemotronh_recipes_pretrain.py index c07e554928..b04b7abe10 100644 --- a/tests/functional_tests/test_groups/recipes/test_nemotronh_recipes_pretrain.py +++ b/tests/functional_tests/test_groups/recipes/test_nemotronh_recipes_pretrain.py @@ -28,11 +28,12 @@ ( nemotron_3_nano_pretrain_config, "nemotron_3_nano", - {"tensor_model_parallel_size": 2, "pipeline_model_parallel_size": 1, "expert_model_parallel_size": 2}, + {"tensor_model_parallel_size": 1, "pipeline_model_parallel_size": 1, "expert_model_parallel_size": 2}, { "hidden_size": 672, "num_layers": 3, "hybrid_layer_pattern": "M*E", + "sequence_parallel": False, "num_moe_experts": 16, "moe_token_dispatcher_type": "alltoall", "moe_shared_expert_overlap": True, diff --git a/tests/unit_tests/recipes/nemotronh/test_nemotron_3_nano.py b/tests/unit_tests/recipes/nemotronh/test_nemotron_3_nano.py index f63d744e94..ed07355f4d 100644 --- a/tests/unit_tests/recipes/nemotronh/test_nemotron_3_nano.py +++ b/tests/unit_tests/recipes/nemotronh/test_nemotron_3_nano.py @@ -16,10 +16,13 @@ Unit tests for Nemotron 3 and 3.5 Nano recipe configuration builders. Tests cover: +- Pretrain configuration defaults (parameterless API) +- SFT configuration (parameterless API for full supervised fine-tuning) +- PEFT configuration with LoRA and DoRA (``peft_scheme`` parameter) +- MoE-specific settings (HybridEP pretraining, DeepEP finetuning, expert parallelism) - Separate Nemotron 3 and 3.5 pretrain configurations - SFT configuration with explicit Nemotron 3.5 MTP overrides - PEFT configuration with explicit Nemotron 3.5 MTP overrides, LoRA, and DoRA -- MoE-specific settings (DeepEP, expert parallelism) - Parallelism and tokenizer configurations """ @@ -61,9 +64,9 @@ def test_pretrain_config_default_parameters(self): assert isinstance(config.model, HybridModelProvider) # Check model configuration defaults - assert config.model.tensor_model_parallel_size == 1 + assert config.model.tensor_model_parallel_size == 8 assert config.model.pipeline_model_parallel_size == 1 - assert config.model.sequence_parallel is False + assert config.model.sequence_parallel is True # Check expert parallelism defaults assert config.model.expert_tensor_parallel_size == 1 @@ -71,8 +74,8 @@ def test_pretrain_config_default_parameters(self): # Check training configuration assert config.train.train_iters == 39735 - assert config.train.global_batch_size == 1024 - assert config.train.micro_batch_size == 1 + assert config.train.global_batch_size == 3072 + assert config.train.micro_batch_size == 2 # Check dataset configuration assert config.dataset.seq_length == 8192 @@ -83,12 +86,11 @@ def test_pretrain_config_default_parameters(self): # Check comm overlap assert config.comm_overlap is not None - assert config.comm_overlap.tp_comm_overlap is True + assert config.comm_overlap.tp_comm_overlap is False assert config.comm_overlap.tp_comm_bootstrap_backend == "nccl" # Check precision - assert config.mixed_precision.bf16 is True - assert config.mixed_precision.grad_reduce_in_fp32 is False + assert config.mixed_precision == "bf16_mixed" # MTP is opt-in for pretraining. assert config.model.mtp_num_layers == 0 @@ -105,7 +107,8 @@ def test_nemotron_3_5_pretrain_config(self): assert config.model.keep_mtp_spec_in_bf16 is True assert config.model.mtp_loss_scaling_factor == 0.3 assert config.model.calculate_per_token_loss == base_config.model.calculate_per_token_loss - assert config.model.use_te_rng_tracker == base_config.model.use_te_rng_tracker + assert config.model.use_te_rng_tracker is True + assert base_config.model.use_te_rng_tracker is False assert recipe_module._NEMOTRON_3_5_NANO_MODEL_ID == ("nvidia/NVIDIA-Nemotron-3.5-Nano-30B-A3B-BF16") assert config.model.hf_model_id == recipe_module._NEMOTRON_3_5_NANO_MODEL_ID assert config.tokenizer.tokenizer_model == recipe_module._NEMOTRON_3_5_NANO_MODEL_ID @@ -138,7 +141,7 @@ def test_pretrain_config_moe_kernel_settings(self): assert config.model.moe_router_fusion is False assert config.model.moe_permute_fusion is True assert config.model.moe_grouped_gemm is True - assert config.model.cross_entropy_loss_fusion is True + assert config.model.cross_entropy_loss_fusion is False assert config.model.cross_entropy_fusion_impl == "native" def test_pretrain_config_optimizer_settings(self): @@ -157,6 +160,7 @@ def test_pretrain_config_optimizer_settings(self): assert config.optimizer.main_params_dtype == torch.float32 assert config.optimizer.exp_avg_dtype == torch.float32 assert config.optimizer.exp_avg_sq_dtype == torch.float32 + assert config.optimizer.optimizer_cpu_offload is False def test_pretrain_config_checkpoint_settings(self): """Test checkpoint settings for pretrain config.""" @@ -185,9 +189,9 @@ def test_sft_config_default_parameters(self): assert isinstance(config.model, HybridModelProvider) # Check default parallelism for SFT - assert config.model.tensor_model_parallel_size == 1 + assert config.model.tensor_model_parallel_size == 4 assert config.model.pipeline_model_parallel_size == 1 - assert config.model.sequence_parallel is False + assert config.model.sequence_parallel is True # Check expert parallelism assert config.model.expert_tensor_parallel_size == 1 @@ -207,11 +211,12 @@ def test_sft_config_default_parameters(self): assert config.mixed_precision == "bf16_mixed" def test_sft_config_deepep_settings(self): - """Test that SFT config has correct DeepEP/MoE dispatcher settings.""" + """Test that SFT config retains the established DeepEP dispatcher.""" config = nemotron_3_nano_sft_config() # Check MoE dispatcher settings - assert config.model.moe_token_dispatcher_type is not None + assert config.model.moe_token_dispatcher_type == "flex" + assert config.model.moe_flex_dispatcher_backend == "deepep" def test_sft_config_custom_parallelism(self): """Test SFT config with custom parallelism applied after creation.""" @@ -273,18 +278,10 @@ def test_finetuning_uses_nemotron_3_base_model( mtp_num_layers, tokenizer_model, ): - """SFT and PEFT derive the provider from Nemotron 3 before applying MTP.""" - provider = HybridModelProvider() - provider.hf_model_id = recipe_module._NEMOTRON_3_NANO_MODEL_ID - bridge = Mock() - bridge.to_megatron_provider.return_value = provider - - with patch.object(recipe_module.AutoBridge, "from_hf_pretrained", return_value=bridge) as from_hf: - config = recipe_factory() - - from_hf.assert_called_once_with(recipe_module._NEMOTRON_3_NANO_MODEL_ID) - bridge.to_megatron_provider.assert_called_once_with(load_weights=False) - assert config.model is provider + """SFT and PEFT use the shared Nano provider before applying MTP.""" + config = recipe_factory() + + assert isinstance(config.model, HybridModelProvider) assert config.model.mtp_num_layers == mtp_num_layers assert config.model.mtp_hybrid_override_pattern == ("*E" if mtp_num_layers else None) assert config.model.mtp_use_repeated_layer is bool(mtp_num_layers) @@ -293,39 +290,27 @@ def test_finetuning_uses_nemotron_3_base_model( assert config.model.hf_model_id == tokenizer_model assert config.tokenizer.tokenizer_model == tokenizer_model - @pytest.mark.parametrize( - ("recipe_factory", "base_factory_name", "recipe_kwargs", "base_args"), - [ - ( - nemotron_3_5_nano_sft_config, - "nemotron_3_nano_sft_8gpu_h100_bf16_config", - {}, - (), - ), - ( - nemotron_3_5_nano_peft_config, - "nemotron_3_nano_peft_8gpu_h100_bf16_config", - {"peft_scheme": "dora"}, - ("dora",), - ), - ], - ) - def test_nemotron_3_5_finetuning_overrides_base_recipe( - self, - recipe_factory, - base_factory_name, - recipe_kwargs, - base_args, - ): - """Nemotron 3.5 SFT and PEFT mutate their corresponding Nemotron 3 config.""" + def test_nemotron_3_5_sft_retains_generic_topology(self): + """Nemotron 3.5 SFT remains independent of the specialized TP4 recipe.""" + config = nemotron_3_5_nano_sft_config() + + assert config.model.tensor_model_parallel_size == 1 + assert config.model.sequence_parallel is False + + def test_nemotron_3_5_peft_overrides_base_recipe(self): + """Nemotron 3.5 PEFT mutates the corresponding Nemotron 3 config.""" base_config = Mock() base_config.model = HybridModelProvider() base_config.tokenizer = Mock() - with patch.object(recipe_module, base_factory_name, return_value=base_config) as base_factory: - config = recipe_factory(**recipe_kwargs) + with patch.object( + recipe_module, + "nemotron_3_nano_peft_8gpu_h100_bf16_config", + return_value=base_config, + ) as base_factory: + config = nemotron_3_5_nano_peft_config(peft_scheme="dora") - base_factory.assert_called_once_with(*base_args) + base_factory.assert_called_once_with("dora") assert config is base_config assert config.model.mtp_num_layers == 2 assert config.model.mtp_hybrid_override_pattern == "*E" @@ -344,15 +329,8 @@ def test_finetuning_recipes_do_not_expose_model_id(self): def test_nemotron_3_5_openmath_sft_recipe_owns_verified_h100_defaults(self): """The dedicated H100 recipe makes the standard packed SFT command concise.""" - provider = HybridModelProvider() - provider.hf_model_id = recipe_module._NEMOTRON_3_NANO_MODEL_ID - bridge = Mock() - bridge.to_megatron_provider.return_value = provider + config = recipe_module.nemotron_3_5_nano_sft_openmathinstruct2_packed_config() - with patch.object(recipe_module.AutoBridge, "from_hf_pretrained", return_value=bridge) as from_hf: - config = recipe_module.nemotron_3_5_nano_sft_openmathinstruct2_packed_config() - - from_hf.assert_called_once_with(recipe_module._NEMOTRON_3_NANO_MODEL_ID) assert config.model.mtp_num_layers == 2 assert config.model.mtp_hybrid_override_pattern == "*E" assert config.model.mtp_use_repeated_layer is True @@ -468,8 +446,9 @@ def test_peft_config_deepep_settings(self): """Test that PEFT config has correct MoE dispatcher settings.""" config = nemotron_3_nano_peft_config() - # Check MoE dispatcher settings exist - assert config.model.moe_token_dispatcher_type is not None + # Check MoE dispatcher settings + assert config.model.moe_token_dispatcher_type == "flex" + assert config.model.moe_flex_dispatcher_backend == "deepep" def test_peft_config_custom_parallelism(self): """Test PEFT config with custom parallelism applied after creation.""" diff --git a/tests/unit_tests/recipes/nemotronh/test_nemotron_3_nano_hardware.py b/tests/unit_tests/recipes/nemotronh/test_nemotron_3_nano_hardware.py new file mode 100644 index 0000000000..d35a929fa9 --- /dev/null +++ b/tests/unit_tests/recipes/nemotronh/test_nemotron_3_nano_hardware.py @@ -0,0 +1,394 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Hardware parity and convergence-safety tests for Nemotron 3 Nano recipes.""" + +import importlib +from collections.abc import Callable + +import pytest +import torch + +import megatron.bridge.recipes as recipes +from megatron.bridge.models.hybrid.hybrid_provider import HybridModelProvider +from megatron.bridge.perf_recipes.nemotronh.gb200.nemotronh import ( + nemotron_3_nano_pretrain_8gpu_gb200_bf16_config as gb200_perf_config, +) +from megatron.bridge.perf_recipes.nemotronh.h100.nemotronh import ( + nemotron_3_nano_pretrain_16gpu_h100_bf16_config as h100_perf_config, +) +from megatron.bridge.recipes.nemotronh._nemotron_3_nano import _nemotron_3_nano_pretrain_reference_config +from megatron.bridge.recipes.nemotronh.gb200.nemotron_3_nano import ( + nemotron_3_nano_peft_8gpu_gb200_bf16_config, + nemotron_3_nano_pretrain_8gpu_gb200_bf16_config, + nemotron_3_nano_sft_8gpu_gb200_bf16_config, +) +from megatron.bridge.recipes.nemotronh.h100.nemotron_3_nano import ( + nemotron_3_nano_peft_8gpu_h100_bf16_config, + nemotron_3_nano_pretrain_8gpu_h100_bf16_config, + nemotron_3_nano_sft_8gpu_h100_bf16_config, +) +from megatron.bridge.training.config import ConfigContainer +from megatron.bridge.utils.cuda_graph import cuda_graph_module_names +from tests.unit_tests.recipes.recipe_test_utils import patch_recipe_module_global + + +class _OfflineAutoBridge: + """Build the reference provider without reading a Hugging Face config.""" + + @classmethod + def from_hf_pretrained(cls, *args: object, **kwargs: object) -> "_OfflineAutoBridge": + del args, kwargs + return cls() + + def to_megatron_provider(self, *args: object, **kwargs: object) -> HybridModelProvider: + del args, kwargs + model = _nemotron_3_nano_pretrain_reference_config().model + assert isinstance(model, HybridModelProvider) + return model + + +@pytest.fixture(autouse=True) +def _patch_gb200_auto_bridge(monkeypatch: pytest.MonkeyPatch) -> None: + """Keep the GB200 public recipe checks deterministic and offline.""" + monkeypatch.setenv("HF_DATASETS_OFFLINE", "1") + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + monkeypatch.setenv("TRANSFORMERS_OFFLINE", "1") + module = importlib.import_module("megatron.bridge.recipes.nemotronh.gb200.nemotron_3_nano") + patch_recipe_module_global(monkeypatch, module, "AutoBridge", _OfflineAutoBridge) + + +_APPROVED_PRETRAIN_PERF_FIELDS = ( + "model.pipeline_model_parallel_size", + "model.virtual_pipeline_model_parallel_size", + "model.context_parallel_size", + "model.expert_tensor_parallel_size", + "model.expert_model_parallel_size", + "model.moe_token_dispatcher_type", + "model.moe_flex_dispatcher_backend", + "model.moe_shared_expert_overlap", + "model.moe_permute_fusion_into_hybridep", + "model.moe_hybridep_num_blocks_permute", + "model.moe_hybridep_num_blocks_unpermute", + "model.moe_hybridep_num_sms_preprocessing", + "model.transformer_impl", + "model.attention_backend", + "model.moe_router_fusion", + "model.moe_permute_fusion", + "model.moe_grouped_gemm", + "model.masked_softmax_fusion", + "model.use_fused_weighted_squared_relu", + "model.fine_grained_activation_offloading", + "model.offload_modules", + "ddp.overlap_grad_reduce", + "ddp.overlap_param_gather", + "ddp.use_distributed_optimizer", + "comm_overlap.overlap_moe_expert_parallel_comm", + "comm_overlap.delay_wgrad_compute", + "comm_overlap.tp_comm_bootstrap_backend", + "comm_overlap.overlap_p2p_comm", + "comm_overlap.batch_p2p_comm", +) + + +def _get_field(config: ConfigContainer, path: str) -> object: + value: object = config + for part in path.split("."): + value = getattr(value, part) + return value + + +def _flex_dispatcher_num_sms(config: ConfigContainer) -> int | None: + return config.model.moe_flex_dispatcher_num_sms or config.model.moe_hybridep_num_sms + + +def _assert_convergence_sensitive_model_contract(config: ConfigContainer) -> None: + """Assert routing, capacity, objective, initialization, and dropout invariants.""" + assert config.model.moe_router_force_load_balancing is False + assert config.model.moe_router_force_biased is None + assert config.model.moe_router_load_balancing_type == "seq_aux_loss" + assert config.model.moe_aux_loss_coeff == 0.0001 + assert config.model.moe_z_loss_coeff is None + assert config.model.moe_input_jitter_eps is None + assert config.model.moe_router_topk == 6 + assert config.model.moe_router_topk_scaling_factor == 2.5 + assert config.model.moe_router_num_groups == 1 + assert config.model.moe_router_group_topk == 1 + assert config.model.moe_router_score_function == "sigmoid" + assert config.model.moe_router_pre_softmax is False + assert config.model.moe_router_topk_limited_devices is None + assert config.model.moe_router_dtype == "fp32" + assert config.model.moe_router_enable_expert_bias is True + assert config.model.moe_router_bias_update_rate == 1e-3 + assert config.model.moe_enable_routing_replay is False + assert config.model.moe_apply_probs_on_input is False + assert config.model.moe_shared_expert_gate is False + + assert config.model.moe_token_dropping is False + assert config.model.moe_expert_capacity_factor is None + assert config.model.moe_expert_rank_capacity_factor is None + assert config.model.moe_pad_expert_input_to_capacity is False + + assert config.model.init_method_std == 0.0173 + assert config.model.hidden_dropout == 0.0 + assert config.model.attention_dropout == 0.0 + + +@pytest.mark.unit +def test_h100_pretrain_approved_fields_match_perf_reference() -> None: + """Approved execution fields should stay aligned with the H100 perf reference.""" + library_config = nemotron_3_nano_pretrain_8gpu_h100_bf16_config() + perf_config = h100_perf_config() + + for field in _APPROVED_PRETRAIN_PERF_FIELDS: + assert _get_field(library_config, field) == _get_field(perf_config, field), field + assert _flex_dispatcher_num_sms(library_config) == _flex_dispatcher_num_sms(perf_config) == 16 + + +@pytest.mark.unit +def test_gb200_pretrain_approved_fields_match_perf_reference() -> None: + """Approved execution fields should stay aligned with the GB200 perf reference.""" + library_config = nemotron_3_nano_pretrain_8gpu_gb200_bf16_config() + perf_config = gb200_perf_config() + + for field in _APPROVED_PRETRAIN_PERF_FIELDS: + assert _get_field(library_config, field) == _get_field(perf_config, field), field + assert _flex_dispatcher_num_sms(library_config) == _flex_dispatcher_num_sms(perf_config) == 16 + + +@pytest.mark.unit +def test_h100_pretrain_uses_8gpu_memory_execution_config() -> None: + """The 8-GPU library uses the validated memory-safe execution exceptions.""" + library_config = nemotron_3_nano_pretrain_8gpu_h100_bf16_config() + perf_config = h100_perf_config() + + assert perf_config.model.tensor_model_parallel_size == 1 + assert perf_config.model.sequence_parallel is False + assert perf_config.model.recompute_granularity == "selective" + assert perf_config.model.recompute_modules == ["moe", "layernorm"] + + assert library_config.model.tensor_model_parallel_size == 8 + assert library_config.model.sequence_parallel is True + assert library_config.model.recompute_granularity == perf_config.model.recompute_granularity == "selective" + assert library_config.model.recompute_method is perf_config.model.recompute_method is None + assert library_config.model.recompute_num_layers is perf_config.model.recompute_num_layers is None + assert library_config.model.recompute_modules == perf_config.model.recompute_modules == ["moe", "layernorm"] + + assert library_config.model.cross_entropy_loss_fusion is False + assert perf_config.model.cross_entropy_loss_fusion is True + assert library_config.model.cross_entropy_fusion_impl == "native" + assert library_config.train.empty_unused_memory_level == 2 + assert perf_config.train.empty_unused_memory_level == 0 + assert library_config.validation.eval_micro_batch_size == 1 + assert perf_config.validation.eval_micro_batch_size is None + assert library_config.validation.eval_global_batch_size is perf_config.validation.eval_global_batch_size is None + assert library_config.optimizer.optimizer_cpu_offload is perf_config.optimizer.optimizer_cpu_offload is False + assert library_config.comm_overlap.tp_comm_overlap is False + assert perf_config.comm_overlap.tp_comm_overlap is True + + assert perf_config.model.cuda_graph_impl == "transformer_engine" + assert cuda_graph_module_names(perf_config.model) == ["attn", "mamba"] + assert perf_config.model.use_te_rng_tracker is True + assert perf_config.rng.te_rng_tracker is True + assert library_config.model.cuda_graph_impl == "none" + assert cuda_graph_module_names(library_config.model) == [] + assert library_config.model.use_te_rng_tracker is False + assert library_config.rng.te_rng_tracker is False + + +@pytest.mark.unit +def test_gb200_pretrain_retains_validated_execution_contract() -> None: + """The GB200 library should retain the converged 4K graph-enabled contract.""" + library_config = nemotron_3_nano_pretrain_8gpu_gb200_bf16_config() + perf_config = gb200_perf_config() + + assert library_config.model.tensor_model_parallel_size == perf_config.model.tensor_model_parallel_size == 1 + assert library_config.model.sequence_parallel is perf_config.model.sequence_parallel is False + assert library_config.model.expert_tensor_parallel_size == perf_config.model.expert_tensor_parallel_size == 1 + assert library_config.model.expert_model_parallel_size == perf_config.model.expert_model_parallel_size == 8 + assert library_config.model.recompute_granularity is perf_config.model.recompute_granularity is None + assert library_config.model.recompute_method is perf_config.model.recompute_method is None + assert library_config.model.recompute_num_layers is perf_config.model.recompute_num_layers is None + assert library_config.model.recompute_modules == perf_config.model.recompute_modules is None + + assert library_config.model.cuda_graph_impl == perf_config.model.cuda_graph_impl == "transformer_engine" + assert ( + cuda_graph_module_names(library_config.model) + == cuda_graph_module_names(perf_config.model) + == [ + "attn", + "mamba", + "moe_router", + "moe_preprocess", + ] + ) + assert library_config.model.cuda_graph_warmup_steps == perf_config.model.cuda_graph_warmup_steps == 3 + assert library_config.model.use_te_rng_tracker is perf_config.model.use_te_rng_tracker is True + assert library_config.rng.te_rng_tracker is perf_config.rng.te_rng_tracker is True + + assert library_config.model.cross_entropy_loss_fusion is perf_config.model.cross_entropy_loss_fusion is True + assert library_config.model.cross_entropy_fusion_impl == "native" + assert library_config.model.apply_rope_fusion is True + assert library_config.model.seq_length == library_config.dataset.seq_length == 4096 + assert library_config.train.train_iters == 39735 + assert library_config.train.global_batch_size == 3072 + assert library_config.train.micro_batch_size == 2 + + assert library_config.mixed_precision.bf16 is True + assert library_config.mixed_precision.fp8 is None + assert library_config.mixed_precision.grad_reduce_in_fp32 is False + assert library_config.ddp.grad_reduce_in_fp32 is False + assert library_config.ddp.check_for_nan_in_grad is False + assert library_config.rerun_state_machine.check_for_nan_in_loss is False + assert library_config.comm_overlap.tp_comm_overlap is False + assert perf_config.comm_overlap.tp_comm_overlap is True + + +@pytest.mark.unit +def test_h100_pretrain_excludes_benchmark_and_convergence_sensitive_overrides() -> None: + """The H100 library must not inherit benchmark-only or semantic settings.""" + library_config = nemotron_3_nano_pretrain_8gpu_h100_bf16_config() + perf_config = h100_perf_config() + + _assert_convergence_sensitive_model_contract(library_config) + assert perf_config.model.moe_router_force_load_balancing is True + assert library_config.model.moe_router_padding_for_fp8 is False + assert library_config.model.calculate_per_token_loss is False + + assert library_config.train.train_iters == 39735 + assert library_config.train.global_batch_size == 3072 + assert library_config.train.micro_batch_size == 2 + assert perf_config.train.train_iters == 50 + + assert library_config.optimizer.lr == 1.6e-3 + assert library_config.optimizer.min_lr == 1.6e-5 + assert library_config.scheduler.lr_warmup_iters == 333 + assert library_config.optimizer.main_grads_dtype == torch.float32 + assert library_config.optimizer.main_params_dtype == torch.float32 + assert library_config.optimizer.exp_avg_dtype == torch.float32 + assert library_config.optimizer.exp_avg_sq_dtype == torch.float32 + assert library_config.mixed_precision == "bf16_mixed" + assert perf_config.mixed_precision.bf16 is True + assert perf_config.mixed_precision.fp8 is None + assert library_config.ddp.grad_reduce_in_fp32 is True + assert perf_config.ddp.grad_reduce_in_fp32 is False + + assert library_config.model.apply_rope_fusion is False + assert library_config.model.cross_entropy_fusion_impl == "native" + assert perf_config.model.apply_rope_fusion is True + assert perf_config.model.cross_entropy_fusion_impl == "te" + + assert library_config.ddp.check_for_nan_in_grad is True + assert library_config.rerun_state_machine.check_for_nan_in_loss is True + assert perf_config.ddp.check_for_nan_in_grad is False + assert perf_config.rerun_state_machine.check_for_nan_in_loss is False + + +@pytest.mark.unit +def test_gb200_pretrain_excludes_convergence_sensitive_overrides() -> None: + """The GB200 library must preserve learned routing and optimizer semantics.""" + library_config = nemotron_3_nano_pretrain_8gpu_gb200_bf16_config() + perf_config = gb200_perf_config() + + _assert_convergence_sensitive_model_contract(library_config) + assert perf_config.model.moe_router_force_load_balancing is True + assert library_config.model.moe_router_padding_for_fp8 is False + assert library_config.model.calculate_per_token_loss is False + + assert library_config.optimizer.lr == 1.6e-3 + assert library_config.optimizer.min_lr == 1.6e-5 + assert library_config.scheduler.lr_warmup_iters == 333 + assert library_config.optimizer.main_grads_dtype == torch.float32 + assert library_config.optimizer.main_params_dtype == torch.float32 + assert library_config.optimizer.exp_avg_dtype == torch.float32 + assert library_config.optimizer.exp_avg_sq_dtype == torch.float32 + assert library_config.optimizer.optimizer_cpu_offload is False + + +@pytest.mark.unit +@pytest.mark.parametrize( + ("recipe_factory", "tp", "sequence_parallel", "dispatcher", "backend", "dispatcher_sms", "learning_rate"), + [ + (nemotron_3_nano_sft_8gpu_h100_bf16_config, 4, True, "flex", "deepep", 16, 5e-6), + (nemotron_3_nano_peft_8gpu_h100_bf16_config, 1, False, "flex", "deepep", 16, 1e-4), + (nemotron_3_nano_sft_8gpu_gb200_bf16_config, 1, False, "alltoall", None, None, 5e-6), + (nemotron_3_nano_peft_8gpu_gb200_bf16_config, 1, False, "alltoall", None, None, 1e-4), + ], +) +def test_finetune_recipes_retain_safe_execution_defaults( + recipe_factory: Callable[[], ConfigContainer], + tp: int, + sequence_parallel: bool, + dispatcher: str, + backend: str | None, + dispatcher_sms: int | None, + learning_rate: float, +) -> None: + """Packed finetuning uses supported dispatch and keeps graphs/recompute disabled.""" + config = recipe_factory() + + assert config.model.tensor_model_parallel_size == tp + assert config.model.pipeline_model_parallel_size == 1 + assert config.model.virtual_pipeline_model_parallel_size is None + assert config.model.context_parallel_size == 1 + assert config.model.sequence_parallel is sequence_parallel + assert config.model.expert_tensor_parallel_size == 1 + assert config.model.expert_model_parallel_size == 8 + + assert config.model.moe_token_dispatcher_type == dispatcher + assert config.model.moe_flex_dispatcher_backend == backend + assert config.model.moe_flex_dispatcher_num_sms == dispatcher_sms + assert config.model.moe_hybridep_num_sms is None + assert config.model.moe_shared_expert_overlap is False + _assert_convergence_sensitive_model_contract(config) + + assert config.model.cuda_graph_impl == "none" + assert cuda_graph_module_names(config.model) == [] + assert config.model.recompute_granularity is None + assert config.model.recompute_modules is None + assert config.dataset.enable_offline_packing is True + assert config.dataset.seq_length == 2048 + assert config.model.seq_length == 2048 + assert config.mixed_precision == "bf16_mixed" + assert config.model.calculate_per_token_loss is True + + assert config.train.train_iters == 1000 + assert config.train.global_batch_size == 128 + assert config.train.micro_batch_size == 1 + assert config.optimizer.lr == learning_rate + assert config.optimizer.min_lr == 0.0 + assert config.scheduler.lr_warmup_iters == 50 + assert config.optimizer.optimizer_cpu_offload is False + assert config.optimizer.main_grads_dtype == torch.float32 + assert config.optimizer.main_params_dtype == torch.float32 + assert config.optimizer.exp_avg_dtype == torch.float32 + assert config.optimizer.exp_avg_sq_dtype == torch.float32 + + assert config.ddp.overlap_grad_reduce is True + assert config.ddp.overlap_param_gather is True + assert config.ddp.use_distributed_optimizer is True + assert config.comm_overlap is None + + +@pytest.mark.unit +def test_gb200_recipes_are_exported_for_discovery() -> None: + """Top-level recipe discovery should expose every GB200 variant.""" + recipe_names = ( + "nemotron_3_nano_pretrain_8gpu_gb200_bf16_config", + "nemotron_3_nano_sft_8gpu_gb200_bf16_config", + "nemotron_3_nano_peft_8gpu_gb200_bf16_config", + ) + + for recipe_name in recipe_names: + assert callable(getattr(recipes, recipe_name))