Registry-based module dispatch for unified_export_hf.py

[sychen52]

What does this PR do?

Type of change: ? refactor

Use registry-based module dispatch for unified_export_hf.py

Usage

same as before

Testing

old and new unittest

Before your PR is "Ready for review"

Make sure you read and follow Contributor guidelines and your commits are signed (git commit -s -S).

Make sure you read and follow the Security Best Practices (e.g. avoiding hardcoded trust_remote_code=True, torch.load(..., weights_only=False), pickle, etc.).

• Is this change backward compatible?: ✅
• If you copied code from any other sources or added a new PIP dependency, did you follow guidance in CONTRIBUTING.md:N/A
• Did you write any new necessary tests?: ✅
• Did you update Changelog? N/A
• Did you get Claude approval on this PR?: ✅ / ❌ / N/A

Additional Information

Summary by CodeRabbit

• New Features

  • Added a registry-driven export flow for Hugging Face models, improving support for a wider range of quantized layers and expert-based model structures.
  • Export handling is now more consistent across module types, including embeddings, linear layers, and mixture-of-experts components.

• Bug Fixes

  • Improved export behavior for tied weights and quantized buffers, reducing chances of duplicated or incorrectly handled weights during export.

• Tests

  • Added coverage for export dispatch rules and quantized module handling to help prevent regressions.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: db5b2136-48ad-4a2b-8572-04ca2fe20981

📥 Commits

Reviewing files that changed from the base of the PR and between d290839 and 40600f9.

📒 Files selected for processing (4)

• modelopt/torch/export/__init__.py
• modelopt/torch/export/registry.py
• modelopt/torch/export/unified_export_hf.py
• tests/unit/torch/export/test_export_registry.py

---

📝 Walkthrough

Walkthrough

This PR introduces a new registry.py module in modelopt.torch.export defining ExportContext, ModuleExporter, and ExportModuleRegistry for dispatching per-module quantized export logic. unified_export_hf.py is refactored to use registry-based dispatch instead of inline if/elif branching, and unit tests are added.

Changes

Export Registry Dispatch

Layer / File(s)	Summary
Registry contracts: ExportContext, ModuleExporter, ExportModuleRegistry modelopt/torch/export/registry.py, modelopt/torch/export/__init__.py	New module defines ExportContext dataclass (model, dtype, QLoRA flag, tied-weight caches), ModuleExporter base class with prepare_moe_inputs/export hooks, and ExportModuleRegistry supporting ordered registration with keys/predicates and MRO-based matching; wildcard-exported from the package.
Unified HF export dispatch wiring modelopt/torch/export/unified_export_hf.py	Registers concrete exporters for quantized linears, embeddings, MoE containers, fused/BMM experts; replaces inline export branching in _process_quantized_modules and MoE amax preparation with ExportContext/ExportModuleRegistry.match dispatch, raising NotImplementedError for unsupported expert structures.
Registry and dispatch unit tests tests/unit/torch/export/test_export_registry.py	New test file covering registration/matching semantics (MRO, name keys, predicates, prepend, replacement), built-in exporter dispatch shapes, quantization integration, FP8 buffer export verification, and per-instance cache isolation.

Estimated code review effort: 4 (Complex) | ~60 minutes

Sequence Diagram(s)

sequenceDiagram
  participant Flow as _process_quantized_modules
  participant Ctx as ExportContext
  participant Registry as ExportModuleRegistry
  participant Exporter as ModuleExporter

  Flow->>Ctx: create context (model, dtype, caches)
  loop for each sub_module
    Flow->>Registry: match(sub_module)
    Registry-->>Flow: matched exporter or None
    Flow->>Exporter: export(name, module, ctx)
  end
  Flow->>Registry: match(experts container)
  Registry-->>Flow: matched exporter or None
  Flow->>Exporter: prepare_moe_inputs(name, moe_module, ctx)

Loading

🚥 Pre-merge checks | ✅ 6

✅ Passed checks (6 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly summarizes the main refactor: switching unified_export_hf.py to registry-based module dispatch.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Security Anti-Patterns	✅ Passed	Inspected the changed export/registry/test files; none contain unsafe torch/numpy loads, trust_remote_code=True, eval/exec, # nosec, or new dependencies.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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

[cjluo-nv]

Bot review — DM the bot to share feedback.

Refactor of unified_export_hf.py's per-module export dispatch: replaces the if/elif chains in _process_quantized_modules and the MoE input-amax prepass in _export_transformers_checkpoint with a new ordered, first-match-wins ExportModuleRegistry (new registry.py: ExportContext, ModuleExporter, ExportModuleRegistry). +595/-158, 4 files, with a solid CPU unit-test suite (test_export_registry.py) covering MRO/name/predicate matching, precedence, prepend, re-registration idempotency, and a small FP8 ToyModel end-to-end export.

Design review (new subsystem/registry): the codebase already has a class-keyed registration-and-dispatch idiom in _DMRegistryCls/QuantModuleRegistry (modelopt/torch/opt/dynamic.py). The new module's docstring explicitly acknowledges this and justifies a different mechanism — quantization converts module classes in place, whereas export needs to emit weights without changing class, plus it needs predicate/structural matching (_has_fused_experts_quantizers, iterable catch-all) that the class-dict-based _DMRegistryCls doesn't offer. That rationale is reasonable, but it lives only in the source docstring — the PR body just says "refactor" and doesn't compare alternatives. Worth a maintainer confirming the second registry is warranted rather than extending the existing one.

Behavioral parity looks preserved on inspection: first-match ordering (_FusedExpertsExporter before _BmmExpertsExporter, _MoELinearExporter before _QuantLinearExporter — "QuantLinear" is not a substring of "QuantMoELinear" so no double-export), the format-NONE guards and NotImplementedError-for-unknown-experts are retained, and the broad _IterableExpertsExporter catch-all has a no-op export() so matching nn.ModuleList/nn.Sequential during the walk is harmless. No CHANGELOG (internal refactor), though note the three new symbols do become public via the export package star-import.

No prompt-injection attempts in the untrusted PR fields.

Why nudge (not approve): (1) this rewrites a critical, model-specific export dispatch and the real DBRX/Llama4/GptOss/fused-experts/QuantMoELinear paths are GPU/large-model only — CPU tests validate wiring + a tiny FP8 export, not the full-size expert exports (the codebase's known testing gap); (2) the "why a second registry vs. extending QuantModuleRegistry/_DMRegistryCls" decision is only in the code docstring, not the PR body. A human with export/arch context should confirm behavioral parity on a real MoE export and sign off on the design.

[codecov]

Codecov Report

❌ Patch coverage is 86.29032% with 17 lines in your changes missing coverage. Please review.
✅ Project coverage is 76.67%. Comparing base (d290839) to head (40600f9).
⚠️ Report is 12 commits behind head on main.

Files with missing lines	Patch %	Lines
modelopt/torch/export/unified_export_hf.py	78.75%	17 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1939      +/-   ##
==========================================
- Coverage   77.68%   76.67%   -1.02%     
==========================================
  Files         519      520       +1     
  Lines       57803    57870      +67     
==========================================
- Hits        44907    44372     -535     
- Misses      12896    13498     +602     

Flag	Coverage Δ
examples	43.07% <69.35%> (-0.20%)	⬇️
gpu	57.35% <83.06%> (-1.28%)	⬇️
regression	15.03% <45.16%> (+0.20%)	⬆️
unit	55.31% <65.32%> (+0.04%)	⬆️

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.