[Refactor] Model-family spec registry for export: MoE expert naming by h-guo18 · Pull Request #1828 · NVIDIA/Model-Optimizer

h-guo18 · 2026-06-25T20:06:14Z

What does this PR do?

Type of change: Refactor / code health

Adds a declarative per-model-family registry, modelopt/torch/export/modeling/, and
migrates the first data category: MoE expert linear naming. Each family is one
small spec file (families/qwen.py, deepseek.py, mixtral.py, dbrx.py,
gptoss.py, nemotron.py, gemma.py); importing the package registers all specs.

Two HF-path consumers now read the registry instead of branching on class names:

layer_utils.get_expert_linear_names — resolves spec.expert_linear_names first;
the legacy if/elif chain is preserved verbatim as fallback, so unmatched families
are byte-for-byte unaffected.
layer_utils.get_experts_list — reads spec.expert_linear_names gated by
spec.has_iterable_experts, which reproduces the legacy "supported 4 families,
else NotImplementedError" behavior.

Scope note: per MODEL_SPECIFIC_REFACTOR.md, this series targets the unified HF
path only. The TRT-LLM checkpoint path keeps its branches (to be moved out as-is
on a separate track); the Megatron path already follows a per-family registry
pattern and is untouched.

Relationship to #1939 (why a second registry?)

The two registries are deliberate, orthogonal layers — not duplication:

Registry-based module dispatch for unified_export_hf.py #1939's ExportModuleRegistry / PrepareMoEInputsRegistry dispatch on module
structure: which handler processes this module?
modeling/ resolves family data: what are this family's values?

Example: when exporting a QuantMixtralSparseMoeBlock, #1939 matches it to the
shared iterable-experts handler (the same handler serves Qwen/DeepSeek/Gemma4);
inside, get_expert_linear_names resolves the mixtral spec and returns
("w1", "w2", "w3"). One handler serves many families, and one family spec will
serve many handlers — a many-to-many relation indexed from both sides. Folding one
registry into the other would either turn declarative data into per-family
functions, or turn behavior into spec fields. Sharing the matcher machinery between
the two is planned for P4 (see the doc).

Usage

# Adding a MoE model family is a small declarative file; no engine edits.
# modelopt/torch/export/modeling/families/<family>.py
register(
    ModelSpec(
        name="qwen_moe",
        moe_block_names=("Qwen3MoeSparseMoeBlock",),
        expert_linear_names=("gate_proj", "down_proj", "up_proj"),
        has_iterable_experts=True,
    )
)

Testing

New tests/unit/torch/export/test_modeling_specs.py: registry matching
(substring / quantized-class-name / unmatched→None), fallback preservation, and
get_experts_list grouping + unsupported-family rejection.
Existing tests/unit/torch/export/test_layer_utils.py expert-naming tests pass
unchanged through the new spec path (equivalence check).
pytest tests/unit/torch/export/ — 116 passed (test_export_diffusers.py
skipped: pre-existing CUDA/glibc collection error unrelated to this change).
Docs build clean with --fail-on-warning (fixes the previous build-docs failure).

Before your PR is "Ready for review"

Is this change backward compatible?: ✅ (fallback-first; behavior-preserving)
If you copied code from any other sources or added a new PIP dependency: N/A
Did you write any new necessary tests?: ✅ (test_modeling_specs.py)
Did you update Changelog?: N/A (internal refactor, no public API change)

Additional Information

Follow-ups (see MODEL_SPECIFIC_REFACTOR.md): P2 PQS_FUSE_MODULE_MAPPING →
spec.pqs_fuse_rules; P3 is_moe name list → spec.moe_block_names; P4 handlers
consume specs directly + unify matcher machinery with #1939's registry.

copy-pr-bot · 2026-06-25T20:06:18Z

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

coderabbitai · 2026-06-25T20:06:23Z

Important

Review skipped

Draft detected.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: c82bb9cf-8702-4bd1-a47b-6cd03fdfa45e

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

🔍 Trigger review

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests
Commit unit tests in branch haoguo/export-modelinglib

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

codecov · 2026-06-25T20:16:47Z

Codecov Report

❌ Patch coverage is 89.69072% with 20 lines in your changes missing coverage. Please review.
✅ Project coverage is 74.43%. Comparing base (f479e78) to head (fae15c9).
⚠️ Report is 5 commits behind head on main.

Files with missing lines	Patch %	Lines
modelopt/torch/export/layer_utils.py	85.36%	6 Missing ⚠️
modelopt/torch/export/unified_export_hf.py	28.57%	5 Missing ⚠️
modelopt/torch/export/plugins/vllm_fakequant_hf.py	25.00%	3 Missing ⚠️
modelopt/torch/export/quant_utils.py	40.00%	3 Missing ⚠️
modelopt/torch/export/hf_export_handlers.py	0.00%	2 Missing ⚠️
modelopt/modeling/registry.py	98.11%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1828      +/-   ##
==========================================
- Coverage   77.87%   74.43%   -3.44%     
==========================================
  Files         522      544      +22     
  Lines       58452    61841    +3389     
==========================================
+ Hits        45522    46034     +512     
- Misses      12930    15807    +2877

Flag	Coverage Δ
unit	`55.13% <89.69%> (-0.27%)`	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 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.

h-guo18 · 2026-06-25T23:18:29Z

+
+
+@dataclass
+class ModelSpec:


trtllm export path is deprecated.

h-guo18 · 2026-06-25T23:19:27Z

    RelativeAttentionTableConfig,
    RgLruConfig,
 )
 from .model_config_utils import pad_weights


HF export only.

Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>

github-actions · 2026-07-14T22:53:32Z

PR Preview Action v1.8.1
🚀 View preview at https://NVIDIA.github.io/Model-Optimizer/pr-preview/pr-1828/
Built to branch `gh-pages` at 2026-07-16 20:46 UTC. Preview will be ready when the GitHub Pages deployment is complete.

Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>

sychen52 · 2026-07-16T18:36:43Z

@@ -0,0 +1,120 @@
+# Export: Model-Specific Logic Refactor — Action Plan


Is this file meant to be checked in?

No, it will be removed

This PR is still messy. Need some cleaning before ready for review

Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>

h-guo18 changed the title ~~[Refactor] Extract model specific logics in export~~ [Refactor] Extract model specific logics in export lib Jun 25, 2026

h-guo18 commented Jun 25, 2026

View reviewed changes

squash: pr1828

e2a32bb

Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>

h-guo18 force-pushed the haoguo/export-modelinglib branch from 8ddd26c to e2a32bb Compare July 14, 2026 22:49

h-guo18 changed the title ~~[Refactor] Extract model specific logics in export lib~~ [Refactor] Model-family spec registry for export: MoE expert naming Jul 14, 2026

h-guo18 added 9 commits July 14, 2026 23:45

move more model logic

c7f474b

Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>

top level dir

8d81535

Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>

rename

2c60c7e

Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>

categorized spec

ddbf88f

Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>

better matching

c5c68b2

Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>

dev

0fea97c

Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>

dev

ed2eb7a

Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>

dev

10d4682

Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>

dev

1bb5139

Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>

sychen52 reviewed Jul 16, 2026

View reviewed changes

h-guo18 added 3 commits July 16, 2026 20:41

dev

fae15c9

Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>

dev

23b7ca8

Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>

dev

58629a6

Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[Refactor] Model-family spec registry for export: MoE expert naming#1828

[Refactor] Model-family spec registry for export: MoE expert naming#1828
h-guo18 wants to merge 13 commits into
mainfrom
haoguo/export-modelinglib

h-guo18 commented Jun 25, 2026 •

edited

Loading

Uh oh!

copy-pr-bot Bot commented Jun 25, 2026

Uh oh!

coderabbitai Bot commented Jun 25, 2026 •

edited

Loading

Review skipped

Uh oh!

codecov Bot commented Jun 25, 2026 •

edited

Loading

Uh oh!

h-guo18 Jun 25, 2026

Uh oh!

h-guo18 Jun 25, 2026

Uh oh!

github-actions Bot commented Jul 14, 2026 •

edited

Loading

Built to branch `gh-pages` at 2026-07-16 20:46 UTC.
Preview will be ready when the GitHub Pages deployment is complete.

Uh oh!

sychen52 Jul 16, 2026

Uh oh!

h-guo18 Jul 16, 2026

Uh oh!

h-guo18 Jul 16, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

		@@ -0,0 +1,120 @@
		# Export: Model-Specific Logic Refactor — Action Plan

Uh oh!

Conversation

h-guo18 commented Jun 25, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do?

Relationship to #1939 (why a second registry?)

Usage

Testing

Before your PR is "Ready for review"

Additional Information

Uh oh!

copy-pr-bot Bot commented Jun 25, 2026

Uh oh!

coderabbitai Bot commented Jun 25, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Review skipped

Uh oh!

codecov Bot commented Jun 25, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

h-guo18 Jun 25, 2026

Choose a reason for hiding this comment

Uh oh!

h-guo18 Jun 25, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions Bot commented Jul 14, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Built to branch gh-pages at 2026-07-16 20:46 UTC. Preview will be ready when the GitHub Pages deployment is complete.

Uh oh!

sychen52 Jul 16, 2026

Choose a reason for hiding this comment

Uh oh!

h-guo18 Jul 16, 2026

Choose a reason for hiding this comment

Uh oh!

h-guo18 Jul 16, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

h-guo18 commented Jun 25, 2026 •

edited

Loading

coderabbitai Bot commented Jun 25, 2026 •

edited

Loading

codecov Bot commented Jun 25, 2026 •

edited

Loading

github-actions Bot commented Jul 14, 2026 •

edited

Loading

Built to branch `gh-pages` at 2026-07-16 20:46 UTC.
Preview will be ready when the GitHub Pages deployment is complete.