Skip to content

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

Open
HelloWorldLTY wants to merge 1 commit into
THUDM:mainfrom
HelloWorldLTY:nemotron-h-mamba-training
Open

Support NemotronH hybrid Mamba+Attention+MoE training (e.g. Nemotron Nano-30B-A3B)#2211
HelloWorldLTY wants to merge 1 commit into
THUDM:mainfrom
HelloWorldLTY:nemotron-h-mamba-training

Conversation

@HelloWorldLTY

Copy link
Copy Markdown

Motivation

slime currently always builds a GPTModel (pure transformer). Hybrid Mamba-Attention-MoE models such as NVIDIA NemotronH (e.g. Nemotron Nano-30B-A3B: Mamba2 + attention + MoE, pattern MEMEM*E...) therefore get param names that don't match a Megatron-core NemotronH checkpoint (input_layernorm instead of mixer.*), and the training forward crashes because mcore MambaModel.forward() does not accept the loss_mask kwarg that GPTModel does.

This PR adds NemotronH (hybrid mamba arch) training support.

Modifications

  • slime/backends/megatron_utils/model_provider.py: when --hybrid-override-pattern is set, build megatron.core.models.mamba.MambaModel with mamba_stack_spec instead of GPTModel. Param names (decoder.layers.N.mixer.*, GroupedMLP experts.experts, final_norm) then match a Megatron-core NemotronH checkpoint. The critic role still swaps in LinearForLastLayer.
  • slime/backends/megatron_utils/model.py: drop the loss_mask kwarg via a signature check before calling the model. Loss masking already happens in slime's own loss function, so behavior is unchanged for GPT-family models.
  • scripts/models/nemotron-nano-30B-A3B.sh: Megatron MODEL_ARGS for Nemotron Nano-30B-A3B (52 layers, hybrid pattern MEMEM*E..., 128 experts top-k 6, mamba 64 heads / state 128 / 8 groups, squared-ReLU).

Usage notes

  • Weight sync to the rollout engine should use --megatron-to-hf-mode bridge (megatron-bridge ships a NemotronH bridge; the raw converters don't cover nemotron_h). SGLang serves nemotron_h natively.
  • When initializing from a Megatron SFT checkpoint in bridge mode, pass --finetune --no-load-optim --no-load-rng explicitly (optimizer state saved under a different TP layout cannot be resharded).
  • The training environment needs mamba-ssm and causal-conv1d.

Validation

Single-node 8-GPU GRPO on Nemotron-3 Nano-30B-A3B (from an SFT checkpoint; colocate mode, actor TP2/PP1/EP8 + sequence parallel, 2× TP4 SGLang engines, CPU-offloaded Adam, selective moe moe_act recompute): two full iterations of rollout → reward → GRPO update → weight sync completed cleanly. grad_norm 0.34/0.14, train-vs-rollout logprob abs diff ≈ 0.017 (weight sync consistent with the training model).

…Nano-30B-A3B)

slime always builds a GPTModel (pure transformer), so hybrid Mamba-attention-MoE
models like NVIDIA NemotronH produce param names (input_layernorm, no mixer.*)
that do not match a Megatron-core NemotronH checkpoint, and the training forward
crashes on the loss_mask kwarg that mcore MambaModel.forward does not accept.

- model_provider: when --hybrid-override-pattern is set, build
  megatron.core.models.mamba.MambaModel with mamba_stack_spec, yielding
  decoder.layers.N.mixer.* + GroupedMLP experts + final_norm that match a
  Megatron-core NemotronH checkpoint (critic role keeps LinearForLastLayer).
- megatron_utils/model.py: drop the loss_mask kwarg by signature check before
  calling the model; loss masking already happens in slime's own loss function,
  so GPT-family behavior is unchanged.
- scripts/models/nemotron-nano-30B-A3B.sh: Megatron MODEL_ARGS for Nano-30B-A3B
  (52 layers, pattern MEMEM*E..., 128 experts topk6, mamba 64 heads).

Validated end-to-end with single-node 8-GPU GRPO on Nemotron-3 Nano-30B-A3B
(SFT checkpoint, colocate mode, actor TP2/EP8, 2x TP4 SGLang engines,
--megatron-to-hf-mode bridge for weight sync): 2 full iterations of
rollout -> reward -> update -> weight sync; train-vs-rollout logprob abs
diff ~0.017.

Co-authored-by: Wang Yuchen <yw.yy953e@alibaba-inc.com>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01YPhWWrF2GjHDDehtKaMoh7
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant