Support NemotronH hybrid Mamba+Attention+MoE training (e.g. Nemotron Nano-30B-A3B)

scripts/models/nemotron-nano-30B-A3B.sh

@@ -0,0 +1,37 @@
+#!/bin/bash
+# Nemotron Nano-30B-A3B (NemotronH hybrid Mamba2+Attn+MoE) — Megatron MODEL_ARGS.
+# Authoritative values dumped from mb-nano's nemotron_3_nano_finetune_config provider
+# (the exact config our SFT trained with). 52 layers, pattern MEMEM*E..., 128 experts.
+MODEL_ARGS=(
+   --num-layers 52
+   --hidden-size 2688
+   --ffn-hidden-size 1856
+   --num-attention-heads 32
+   --group-query-attention
+   --num-query-groups 2
+   --kv-channels 128
+   --normalization RMSNorm
+   --position-embedding-type none
+   --disable-bias-linear
+   --squared-relu
+   --untie-embeddings-and-output-weights
+   --make-vocab-size-divisible-by 128
+   # --- hybrid Mamba/Attn/MLP allocation ---
+   --hybrid-override-pattern "MEMEM*EMEMEM*EMEMEM*EMEMEM*EMEMEM*EMEMEMEM*EMEMEMEME"
+
+   --mamba-num-heads 64
+   --mamba-head-dim 64
+   --mamba-state-dim 128
+   --mamba-num-groups 8
+   # --- MoE ---
+   --num-experts 128
+   --moe-ffn-hidden-size 1856
+   --moe-router-topk 6
+   --moe-shared-expert-intermediate-size 3712
+   --moe-grouped-gemm
+   --moe-router-load-balancing-type seq_aux_loss
+   --moe-token-dispatcher-type alltoall
+   --transformer-impl local
+   --no-persist-layer-norm
+   --no-gradient-accumulation-fusion
+)

slime/backends/megatron_utils/model.py

@@ -61,6 +61,19 @@ def _should_update_microbatch_pbar(model) -> bool:
     return mpu.is_pipeline_last_stage(ignore_virtual=True)
 
 
+def _drop_unsupported_loss_mask_kwarg(model, forward_kwargs):
+    # mcore MambaModel.forward (hybrid NemotronH) has no loss_mask kwarg
+    # (GPTModel does). Drop it when unsupported; loss masking happens in
+    # slime's own loss function, not the model.
+    import inspect
+
+    m = model
+    while hasattr(m, "module"):
+        m = m.module
+    if "loss_mask" not in inspect.signature(m.forward).parameters:
+        forward_kwargs.pop("loss_mask", None)
+
+
 def _wrap_forward_step_with_microbatch_pbar(forward_step_func, pbar):
     if pbar is None:
         return forward_step_func
@@ -428,6 +441,7 @@ def forward_step(
             "packed_seq_params": packed_seq_params,
             "loss_mask": batch["full_loss_masks"],
         }
+        _drop_unsupported_loss_mask_kwarg(model, forward_kwargs)
         if batch["multimodal_train_inputs"] is not None:
             forward_kwargs.update(batch["multimodal_train_inputs"])
         output_tensor = model(**forward_kwargs)
@@ -623,6 +637,7 @@ def forward_step(data_iterator: DataIterator, model: GPTModel, return_schedule_p
                 "packed_seq_params": batch["packed_seq_params"],
                 "loss_mask": batch["full_loss_masks"],
             }
+            _drop_unsupported_loss_mask_kwarg(model, forward_kwargs)
 
             if batch["multimodal_train_inputs"] is not None:
                 forward_kwargs.update(batch["multimodal_train_inputs"])

slime/backends/megatron_utils/model_provider.py

@@ -140,6 +140,36 @@ def model_provider(pre_process: bool = True, post_process: bool = True, vp_stage
         # Experimental loading arguments from yaml
         config: TransformerConfig = core_transformer_config_from_args(args)
 
+        # --- Hybrid Mamba-Attention-MoE (e.g. NVIDIA NemotronH / Nano) ---
+        # When a hybrid layer pattern is given, the decoder is a MambaStack, not a
+        # pure transformer, so build a MambaModel (mamba_stack_spec) instead of a
+        # GPTModel. This makes the param names (decoder.layers.N.mixer.*, GroupedMLP
+        # experts.experts, final_norm) match a Megatron-core NemotronH checkpoint.
+        if getattr(args, "hybrid_override_pattern", None):
+            from megatron.core.models.mamba import MambaModel
+            from megatron.core.models.mamba.mamba_layer_specs import mamba_stack_spec
+
+            model = MambaModel(
+                config=config,
+                mamba_stack_spec=mamba_stack_spec,
+                vocab_size=args.padded_vocab_size,
+                max_sequence_length=args.max_position_embeddings,
+                pre_process=pre_process,
+                post_process=post_process,
+                hybrid_attention_ratio=getattr(args, "hybrid_attention_ratio", 0.0),
+                hybrid_mlp_ratio=getattr(args, "hybrid_mlp_ratio", 0.0),
+                hybrid_override_pattern=args.hybrid_override_pattern,
+                fp16_lm_cross_entropy=args.fp16_lm_cross_entropy,
+                parallel_output=True,
+                share_embeddings_and_output_weights=not args.untie_embeddings_and_output_weights,
+                position_embedding_type=args.position_embedding_type,
+                rotary_percent=args.rotary_percent,
+                rotary_base=args.rotary_base,
+            )
+            if post_process and role == "critic":
+                model.output_layer = LinearForLastLayer(input_size=config.hidden_size, output_size=1, config=config)
+            return model
+
         if args.spec is not None:
             transformer_layer_spec = import_module(args.spec)
             # Allow the spec to be a function so that user can use customized Megatron easier.