Support MOE Export for Nemotron H

[jenchen13]

What does this PR do?

Type of change: New feature

Overview: Support Mamba-MOE export for Nemotron H

Usage

# Add a code snippet demonstrating how to use this

Testing

Will test MLM import/export using MLM scripts

• test import from HF
• test export to HF - verify state dicts look the same

Before your PR is "Ready for review"

• Make sure you read and follow Contributor guidelines and your commits are signed.
• Is this change backward compatible?: Yes/No
• Did you write any new necessary tests?: Yes/No
• Did you add or update any necessary documentation?: Yes/No
• Did you update Changelog?: Yes/No

Additional Information

Summary by CodeRabbit

• New Features
  • Improved support for Mixture-of-Experts (MoE) models in import/export workflows.
  • Added handling for expert routing and both local and shared expert projection paths so MoE behavior is preserved during model transfer.
  • Enhanced MLP import/export metadata so router and expert components are correctly detected and exported.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[coderabbitai]

Walkthrough

Added four imports (COL_ETP, ROW_ETP, QKVMerging, QKVSlicing) to the plugin's public surface and extended Nemotron import/export mapping dictionaries to include Mixture-of-Experts (MoE) keys: router, local_experts.linear_fc1, local_experts.linear_fc2, shared_experts.linear_fc1, shared_experts.linear_fc2.

Changes

Cohort / File(s)	Summary
Public imports modelopt/torch/export/plugins/mcore_nemotron.py	Added imports exported from modelopt/torch/export/plugins/mcore_custom: COL_ETP, ROW_ETP, QKVMerging, QKVSlicing.
MoE mapping updates modelopt/torch/export/plugins/mcore_nemotron.py	Extended nemotron_h_causal_lm_import and nemotron_h_causal_lm_export dictionaries with MoE entries: "router", "local_experts.linear_fc1", "local_experts.linear_fc2", "shared_experts.linear_fc1", "shared_experts.linear_fc2".

Sequence Diagram(s)

sequenceDiagram
    autonumber
    participant Caller as Caller
    participant Plugin as mcore_nemotron.py
    participant McoreCustom as mcore_custom
    participant Mappings as Import/Export Dicts

    Note over Plugin,McoreCustom: expose additional symbols and extend mappings

    Caller->>Plugin: request nemotron mappings & exports
    Plugin->>McoreCustom: import COL_ETP, ROW_ETP, QKVMerging, QKVSlicing
    Plugin->>Mappings: build base nemotron import/export dicts
    Mappings->>Mappings: insert "router"
    Mappings->>Mappings: insert "local_experts.linear_fc1"/"linear_fc2"
    Mappings->>Mappings: insert "shared_experts.linear_fc1"/"linear_fc2"
    Plugin->>Caller: return updated mappings and exported symbols

Loading

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~25 minutes

Poem

🐇 I nibble at mappings, row and col,

Routers hum and experts scroll.
New imports hop into place,
Up and down the tensors race,
A cheerful rabbit bounds with grace.

Pre-merge checks and finishing touches

✅ Passed checks (3 passed)

Check name	Status	Explanation
Title Check	✅ Passed	The pull request title "Support MOE Export for Nemotron H" is fully aligned with the main change in the changeset. The PR's primary objective is to add support for Mamba-MOE export for Nemotron H, and the actual changes consist of adding MoE-related keys to the export mapping and importing additional components in the mcore_nemotron.py file. The title is concise, specific, and uses clear technical terminology (MOE for Mixture of Experts) that a team member scanning the history would immediately understand as referring to this feature addition.
Docstring Coverage	✅ Passed	No functions found in the changes. Docstring coverage check skipped.
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.

✨ Finishing touches

• 📝 Generate docstrings

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Post copyable unit tests in a comment
• Commit unit tests in branch jennifchen/nmh-moe-export

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 73.39%. Comparing base (4476f21) to head (7d3245d).

Additional details and impacted files

@@           Coverage Diff           @@
##             main     #447   +/-   ##
=======================================
  Coverage   73.39%   73.39%           
=======================================
  Files         180      180           
  Lines       17976    17976           
=======================================
  Hits        13193    13193           
  Misses       4783     4783           

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[coderabbitai]

Actionable comments posted: 3

📜 Review details

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 6ef9954 and 752c70a.

📒 Files selected for processing (1)

• modelopt/torch/export/plugins/mcore_nemotron.py (4 hunks)

🧰 Additional context used

🧬 Code graph analysis (1)

modelopt/torch/export/plugins/mcore_nemotron.py (1)

modelopt/torch/export/plugins/mcore_custom.py (1)

• NameRemapping (82-91)

🪛 GitHub Actions: Code Quality

modelopt/torch/export/plugins/mcore_nemotron.py

[error] 1-1: pre-commit checks failed. Ruff check reported issues and formatting changes were applied by hooks; the pre-commit run exited with code 1.

⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (8)

• GitHub Check: partial-install (torch)
• GitHub Check: partial-install (onnx)
• GitHub Check: multi-transformers (min)
• GitHub Check: multi-torch (26)
• GitHub Check: multi-py (11)
• GitHub Check: multi-torch (27)
• GitHub Check: windows
• GitHub Check: multi-py (10)

🔇 Additional comments (3)

modelopt/torch/export/plugins/mcore_nemotron.py (3)

20-22: LGTM! New imports are necessary for MoE support.

The COL_ETP and ROW_ETP imports are correctly added and used in the local_experts mappings below.

---

101-106: LGTM! MoE export mappings are consistent.

The MoE mappings for export are well-structured and consistent with the existing export dictionary patterns, using the correct "backbone.layers" prefix throughout.

---

1-1: All pre-commit and ruff checks now pass—no action required.

The trailing whitespace on line 76 and other ruff formatting issues have been resolved. Verification confirms zero remaining errors and no trailing whitespace in the file.