diff --git a/scripts/models/nemotron-nano-30B-A3B.sh b/scripts/models/nemotron-nano-30B-A3B.sh new file mode 100644 index 0000000000..1f61337ee5 --- /dev/null +++ b/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 +) diff --git a/slime/backends/megatron_utils/model.py b/slime/backends/megatron_utils/model.py index 1ad6cd7957..23c290cfb2 100644 --- a/slime/backends/megatron_utils/model.py +++ b/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"]) diff --git a/slime/backends/megatron_utils/model_provider.py b/slime/backends/megatron_utils/model_provider.py index 090b692f18..afa271f846 100644 --- a/slime/backends/megatron_utils/model_provider.py +++ b/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.