refactor(recipes): flatten GB200 recipe construction

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

Author: cuichenx

src/megatron/bridge/recipes/nemotronh/gb200/nemotron_3_nano.py

@@ -14,12 +14,19 @@
 
 """GB200 pretraining recipe for Nemotron 3 Nano."""
 
-from megatron.bridge.recipes.nemotronh.nemotron_3_nano import nemotron_3_nano_pretrain_config
+import torch
+
+from megatron.bridge import AutoBridge
+from megatron.bridge.recipes.common import _pretrain_common
+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
 
 
+_NEMOTRON_3_NANO_MODEL_ID = "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16"
+
+
 def nemotron_3_nano_pretrain_8gpu_gb200_bf16_config() -> ConfigContainer:
     """Return the Nemotron 3 Nano BF16 pretraining config for eight GB200 GPUs.
 
@@ -31,13 +38,21 @@ def nemotron_3_nano_pretrain_8gpu_gb200_bf16_config() -> ConfigContainer:
     Returns:
         GB200 BF16 pretraining configuration.
     """
-    cfg = nemotron_3_nano_pretrain_config()
+    cfg = _pretrain_common()
+
+    cfg.model = AutoBridge.from_hf_pretrained(_NEMOTRON_3_NANO_MODEL_ID).to_megatron_provider(load_weights=False)
+    cfg.tokenizer.tokenizer_model = _NEMOTRON_3_NANO_MODEL_ID
 
     cfg.model.seq_length = 4096
     cfg.dataset.seq_length = 4096
+    cfg.dataset.blend = None
+    cfg.dataset.num_workers = 8
+    cfg.dataset.mmap_bin_files = False
 
     cfg.model.tensor_model_parallel_size = 1
     cfg.model.pipeline_model_parallel_size = 1
+    cfg.model.pipeline_model_parallel_layout = None
+    cfg.model.pipeline_dtype = torch.bfloat16
     cfg.model.virtual_pipeline_model_parallel_size = None
     cfg.model.context_parallel_size = 1
     cfg.model.sequence_parallel = False
@@ -51,28 +66,68 @@ def nemotron_3_nano_pretrain_8gpu_gb200_bf16_config() -> ConfigContainer:
     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 = 3072
+    cfg.train.micro_batch_size = 2
+    cfg.train.manual_gc = False
+    cfg.train.manual_gc_interval = 0
+
+    cfg.model.transformer_impl = "transformer_engine"
+
     # Match the validated GB200 performance recipe's TE-scoped graph set.
     cfg.model.cuda_graph_impl = "transformer_engine"
     set_cuda_graph_modules(cfg.model, ["attn", "mamba", "moe_router", "moe_preprocess"])
     cfg.model.cuda_graph_warmup_steps = 3
     cfg.model.use_te_rng_tracker = True
     cfg.rng.te_rng_tracker = True
 
-    # Retain performance-recipe parity. Nemotron 3 Nano uses no positional
-    # embeddings, so this remains a no-op unless the architecture changes.
+    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.apply_rope_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.rerun_state_machine.check_for_nan_in_loss = False
-    cfg.ddp.check_for_nan_in_grad = 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
 
     # Keep BF16 compute while reducing gradients in BF16 instead of FP32.
     cfg.mixed_precision = get_mixed_precision_config(cfg.mixed_precision)
     cfg.mixed_precision.grad_reduce_in_fp32 = False
+
+    cfg.comm_overlap = CommOverlapConfig(
+        tp_comm_bootstrap_backend="nccl",
+        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.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 = False
+    cfg.ddp.use_distributed_optimizer = True
     cfg.ddp.grad_reduce_in_fp32 = False
 
-    # TP communication overlap requires TP > 1 and sequence parallelism.
-    if cfg.comm_overlap is not None:
-        cfg.comm_overlap.tp_comm_overlap = False
+    cfg.model.init_method_std = 0.0173
+    cfg.model.use_fused_weighted_squared_relu = True
 
     return cfg
 

src/megatron/bridge/recipes/qwen/gb200/qwen35.py

@@ -20,8 +20,6 @@
 from transformers import AutoConfig
 
 from megatron.bridge import AutoBridge
-from megatron.bridge.models.gpt_provider import GPTModelProvider
-from megatron.bridge.models.hybrid.hybrid_provider import HybridModelProvider
 from megatron.bridge.recipes.common import _pretrain_common
 from megatron.bridge.training.comm_overlap import CommOverlapConfig
 from megatron.bridge.training.config import ConfigContainer
@@ -32,20 +30,15 @@
 _QWEN35_35B_A3B_BASE = "Qwen/Qwen3.5-35B-A3B-Base"
 
 
-def _qwen35_text_provider(model_id: str, architecture: str) -> GPTModelProvider | HybridModelProvider:
-    """Build a language-model provider from a unified Qwen3.5 config."""
-    text_config = AutoConfig.from_pretrained(model_id).text_config
-    # The nested text config intentionally omits ``architectures``. AutoBridge
-    # needs it to select the registered causal-LM bridge instead of the VLM.
-    text_config.architectures = [architecture]
-    return AutoBridge.from_hf_config(text_config).to_megatron_provider(load_weights=False)
-
-
 def qwen35_text_9b_pretrain_8gpu_gb200_bf16_config() -> ConfigContainer:
     """Return a text-only Qwen3.5-9B pretraining config for eight GB200 GPUs."""
     cfg = _pretrain_common()
 
-    cfg.model = _qwen35_text_provider(_QWEN35_9B_BASE, "Qwen3_5ForCausalLM")
+    text_config = AutoConfig.from_pretrained(_QWEN35_9B_BASE).text_config
+    # The nested text config intentionally omits ``architectures``. AutoBridge
+    # needs it to select the registered causal-LM bridge instead of the VLM.
+    text_config.architectures = ["Qwen3_5ForCausalLM"]
+    cfg.model = AutoBridge.from_hf_config(text_config).to_megatron_provider(load_weights=False)
     cfg.tokenizer.tokenizer_model = _QWEN35_9B_BASE
     cfg.dataset.seq_length = 4096
     cfg.dataset.blend = None
@@ -114,7 +107,11 @@ def qwen35_text_35b_a3b_pretrain_8gpu_gb200_bf16_config() -> ConfigContainer:
     """Return a text-only Qwen3.5-35B-A3B pretraining config for eight GB200 GPUs."""
     cfg = _pretrain_common()
 
-    cfg.model = _qwen35_text_provider(_QWEN35_35B_A3B_BASE, "Qwen3_5MoeForCausalLM")
+    text_config = AutoConfig.from_pretrained(_QWEN35_35B_A3B_BASE).text_config
+    # The nested text config intentionally omits ``architectures``. AutoBridge
+    # needs it to select the registered causal-LM bridge instead of the VLM.
+    text_config.architectures = ["Qwen3_5MoeForCausalLM"]
+    cfg.model = AutoBridge.from_hf_config(text_config).to_megatron_provider(load_weights=False)
     cfg.tokenizer.tokenizer_model = _QWEN35_35B_A3B_BASE
     cfg.dataset.seq_length = 4096
     cfg.dataset.blend = None

tests/unit_tests/recipes/test_nemotronh_recipes.py

@@ -59,8 +59,9 @@ def to_megatron_provider(self, *args, **kwargs):
 
 @pytest.fixture(autouse=True)
 def _patch_hf_backed_recipe_providers(monkeypatch: pytest.MonkeyPatch) -> None:
-    """Keep Super and Ultra recipe construction deterministic and offline."""
+    """Keep AutoBridge-backed recipe construction deterministic and offline."""
     for module_name in (
+        "megatron.bridge.recipes.nemotronh.gb200.nemotron_3_nano",
         "megatron.bridge.recipes.nemotronh.nemotron_3_super",
         "megatron.bridge.recipes.nemotronh.nemotron_3_ultra",
     ):