Add MoE layer support for NemotronH (Nemotron-3 30B/120B) - #188
Merged
Conversation
Add NemotronHMoELayer with: - NemotronHMoEGate: sigmoid routing with score correction bias (bias affects selection only, not final routing weights) - NemotronHMoEBlock: non-gated FCMLP experts, shared expert, optional latent projection (120B) - Stacked 3D expert tensor splitting in preprocess_weights Config changes: - Add 'E' -> 'moe' to hybrid_override_pattern char_map - Fix type_map: 'moe' was incorrectly mapped to 'mlp' - Add moe_latent_size field to NemotronHConfig - Extract moe_shared_expert_intermediate_size from HF config Cache changes: - Handle 'moe' as stateless layer type (like 'mlp') Test config: - Add MoE variant with 4 experts for L1 graph build tests Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
- Add L4 YAML test cases for Nemotron-3 Nano 30B (MoE, 128 experts) and Nemotron-3 Super 120B (MoE, 512 experts) - Update Nemotron-H Nano 4B YAML: skip_reason → ci_skip_reason, L4 only - Generate golden reference data for all three NemotronH models - Fix HF NemotronH _init_weights dt_bias corruption: the remote code re-initialises dt_bias with torch.rand after from_pretrained loads checkpoint weights, silently corrupting the model. Added _fix_nemotron_h_dt_bias() to torch_reference.py to restore correct values from safetensors files. - Add NemotronH models to _XFAIL_REASONS (HF non-determinism from mamba-ssm CUDA kernels); remove bamba (now passes) Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
- Document NemotronH sigmoid routing, shared experts, latent projection, and graph size impact in moe-models skill - Document com.microsoft.MoE compatibility considerations for sigmoid gates - Add HF _init_weights corruption pitfall (#6) to adding-a-new-model skill with diagnosis pattern and fix reference Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
NemotronH experts use squared ReLU (relu2) activation, which the fused MoE op does not support. Additionally, the sigmoid routing with correction bias and shared expert/latent projection patterns are not compatible. Document this in both the model docstring and MoE skill file. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
Performance Comparison
|
Codecov Report❌ Patch coverage is 📢 Thoughts on this report? Let us know! |
🏗️ Architecture Diff
No architecture changes detected. ✅ Legend: ⚪ No change · 🔵 Minor (attrs/inits) · 🟡 Moderate (nodes added/removed) · 🔴 Major (interface changed) |
Contributor
There was a problem hiding this comment.
Pull request overview
Adds NemotronH MoE (Mixture-of-Experts) support to Mobius’ ONNX graph builder, enabling Nemotron-3 30B/120B hybrid (Mamba2 + MoE + Attention) model construction and associated golden/L4 test artifacts.
Changes:
- Implemented NemotronH MoE gate/block/layer and updated NemotronH weight preprocessing to split stacked expert tensors.
- Updated hybrid cache utilities and build-graph tests to treat
moelayers as stateless (no KV/state I/O). - Added/updated L4 golden test cases + documentation for NemotronH MoE and the HF
dt_biascorruption workaround.
Reviewed changes
Copilot reviewed 15 out of 15 changed files in this pull request and generated 3 comments.
Show a summary per file
| File | Description |
|---|---|
src/mobius/models/nemotron_h.py |
Adds NemotronH MoE modules, integrates moe into layer dispatch, and splits stacked expert weights in preprocess_weights(). |
src/mobius/_configs.py |
Extends NemotronH config parsing to recognize MoE (E / "moe") and adds MoE-specific config fields. |
src/mobius/tasks/_cache_utils.py |
Treats MoE layers as stateless in hybrid cache input/output wiring. |
tests/build_graph_test.py |
Adjusts cache-output assertions for MoE and adds a tiny NemotronH MoE config case. |
tests/_test_configs.py |
Adds NemotronH MoE variant tiny config and removes outdated “MoE not implemented” note. |
tests/e2e_golden_test.py |
Marks NemotronH golden comparisons as xfail (per PR rationale). |
src/mobius/_testing/torch_reference.py |
Adds _fix_nemotron_h_dt_bias() and applies it after HF model load for reference inference. |
testdata/cases/causal-lm/*.yaml |
Adds Nemotron-3 30B/120B cases and adjusts NemotronH nano case to L4-only w/ CI skip reason. |
testdata/golden/causal-lm/*.json |
Adds golden references for the three NemotronH/Nemotron-3 cases. |
.agents/skills/moe-models/SKILL.md |
Documents NemotronH MoE architecture, routing differences, and graph-size implications. |
.agents/skills/adding-a-new-model/SKILL.md |
Adds troubleshooting guidance for HF _init_weights corruption (NemotronH dt_bias). |
justinchuby
force-pushed
the
justinchu/nemotron3-moe-support
branch
from
April 22, 2026 17:15
88b2c29 to
f3834ed
Compare
- Cast NemotronHMoEGate routing computation to float32 for numerical stability (eps=1e-20 underflows in fp16/bf16). Cast routing_weights back to original dtype before returning. - Use huggingface_hub.snapshot_download(local_files_only=True) to resolve exact snapshot path in _fix_nemotron_h_dt_bias instead of unreliable lexicographic glob sort. Add warning when no dt_bias params found. - Add test_nemotron_h_moe_preprocess_weights verifying 3D stacked expert tensor splitting into per-expert 2D weights with correct keys and shapes. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
justinchuby
force-pushed
the
justinchu/nemotron3-moe-support
branch
from
April 22, 2026 17:16
f3834ed to
261e4f1
Compare
- Extend _fix_nemotron_h_init_weights to also restore out_proj.weight params corrupted by HF _init_weights (rescale_prenorm_residual) - Set rescale_prenorm_residual=False before loading NemotronH models in load_torch_model to prevent corruption at load time - Patch NemotronH prepare_inputs_for_generation cache_position bug in generate_golden.py for transformers 5.x compatibility - Enable L5 generation tests for nemotron-h-nano-4b - Regenerate L4 golden with corrected weights (deterministic now) Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
Resolve conflict in tests/e2e_golden_test.py: keep empty _XFAIL_REASONS dict since NemotronH non-determinism was fixed by _fix_nemotron_h_init_weights. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Adds MoE (Mixture of Experts) support for NemotronH architecture, enabling ONNX graph construction for NVIDIA Nemotron-3 30B and 120B models.
Changes
MoE implementation:
NemotronHMoEGate: Sigmoid routing with correction bias and top-k selectionNemotronHMoEBlock: Non-gated FCMLP experts + shared expert + optional latent projectionNemotronHMoELayer: Standalone MoE decoder layer (RMSNorm → MoE → residual)preprocess_weights()for 3D→2D expert weight splittingL4 test cases:
nemotron-h-nano-4b(4B dense),nemotron-3-nano-30b(30B, 128 experts),nemotron-3-super-120b(120B, 512 experts)HF dt_bias corruption fix:
_init_weightsre-initializesdt_biaswithtorch.rand()after checkpoint loading_fix_nemotron_h_dt_bias()to read correct values from safetensorsFused MoE op investigation:
com.microsoft.MoEis incompatible: squared ReLU activation not supported, sigmoid routing incompatibleTesting