Registry-based module dispatch for unified_export_hf.py#1939
Conversation
Signed-off-by: Shiyang Chen <shiychen@nvidia.com>
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (1)
📝 WalkthroughWalkthroughThis PR adds registry-based handler dispatch for unified Hugging Face export, introduces export-scoped context state, moves quantized and MoE handling into built-in handlers, and adds tests for registry semantics and export integration. ChangesUnified HF export registry
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant Flow as _process_quantized_modules
participant Context as ExportContext
participant Registry as ExportModuleRegistry
participant Handler as hf_export_handlers
Flow->>Context: create context with model, dtype, and caches
loop for each sub_module
Flow->>Registry: match sub_module
Registry-->>Flow: matched handler or None
Flow->>Handler: invoke with name, module, and context
end
Flow->>Registry: match experts container
Registry-->>Flow: preparation handler or None
Flow->>Handler: set expert input amax values
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
cjluo-nv
left a comment
cjluo-nv
left a comment
There was a problem hiding this comment.
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.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #1939 +/- ##
==========================================
+ Coverage 76.45% 77.31% +0.85%
==========================================
Files 520 522 +2
Lines 58322 58396 +74
==========================================
+ Hits 44589 45147 +558
+ Misses 13733 13249 -484
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
| return hasattr(module, f"{first_proj_attr}_weight_quantizers") | ||
|
|
||
|
|
||
| @ExportModuleRegistry.register("QuantMoELinear", predicate=lambda m: hasattr(m, "experts")) |
There was a problem hiding this comment.
how about moving this registration to separate files?
There was a problem hiding this comment.
Yes. @shengliangxu mentioned this as well during our infra meeting. I guess we may eventually move them to modellib. But does it make sense to keep it here for now? @h-guo18
There was a problem hiding this comment.
I moved them out to hf_export_handlers.py now.
cjluo-nv
left a comment
cjluo-nv
left a comment
There was a problem hiding this comment.
Bot review — DM the bot to share feedback.
Re: operator question "is this a pure no-op refactor?" — It is mostly behavior-preserving, but I would not call it a pure no-op. It's a faithful mechanical translation of two if/elif dispatchers into an ordered first-match-wins ExportModuleRegistry, and I traced the main paths carefully — but there are a few genuine behavioral divergences (all appear benign today, but they rest on unverified assumptions and the affected paths are GPU/large-model-only, so a human should confirm on a real MoE export).
Behavioral parity trace:
_MoELinearExporter(QuantMoELinear+experts) reproduces the old amax-fill +continue(first-match stops dispatch). ✓_QuantLinearExporter/_QuantEmbeddingExporterpreserve the format-NONE guards and AssertionError wrapping. ✓is_quantlinear("QuantMoELinear")is correctly False (no substring), so no double-export. ✓- Prepass
NotImplementedErrorfor unknown non-iterable experts is retained viamatch() is None. ✓
Not-quite-no-op divergences (why not "pure"):
- New
_IterableExpertsExportercatch-all in_process_quantized_modules. The old whole-model walk had no iterable catch-all; now anynn.ModuleList/nn.Sequentialmatches it. Itsexport()is a no-op so it's harmless, but it is new dispatch behavior that didn't exist before. - Registry ordering assumption for DBRX. Old
_process_quantized_moduleshad no DBRX handling and ran the fused-experts branch (_has_fused_experts_quantizers) unconditionally; the new registry puts_DbrxExpertsExporter(no-opexport) before_FusedExpertsExporter. Parity holds only because DBRX usesw1_linear/w2_linear/v1_linear(not<first_proj>_weight_quantizers), so_has_fused_experts_quantizersis False for it — true today, but the equivalence now silently depends on that. - Prepass key change from quant class names (
QuantGptOssExperts/QuantLlama4TextExperts) to base names (GptOssExperts/Llama4TextExperts) matched via MRO — equivalent in practice, different mechanism.
Testing: CPU suite (test_export_registry.py) covers MRO/name/predicate matching, precedence, prepend, re-registration, and a tiny FP8 ToyModel export — validates wiring, not the full-size DBRX/Llama4/GptOss/fused/QuantMoELinear exports (the codebase's known GPU-only gap). So the "pure no-op" claim is untested on exactly the paths where the ordering divergences above would matter.
Previous-review items:
- 💬 The "why a 2nd registry vs extending
QuantModuleRegistry/_DMRegistryCls" rationale is in theregistry.pydocstring (quantization swaps classes in place; export must emit weights without changing class + needs predicate/structural matching). Reasonable, but still only in the source docstring — the PR body just says "refactor." Worth a maintainer confirming the second registry over extending the existing one. - The inline comment "move this registration to separate files?" (unified_export_hf.py:788) is not addressed — the exporter classes still live inline in
unified_export_hf.py. Minor/optional.
Licensing: new registry.py/test use the standard NVIDIA Apache-2.0 header (year 2026, matches canonical) — no licensing concern. No prompt-injection in untrusted fields.
Net: strong, well-tested refactor and very likely behavior-preserving, but not provably a pure no-op — the iterable catch-all and the DBRX-vs-fused registration ordering are new behaviors whose equivalence hinges on assumptions untested by the CPU suite. Recommend a human sign-off on a real MoE (DBRX/Llama4/GptOss) export before merging.
This way we can have different order and not catch all condition for either Signed-off-by: Shiyang Chen <shiychen@nvidia.com>
cjluo-nv
left a comment
cjluo-nv
left a comment
There was a problem hiding this comment.
Bot review (gpt-5.6-sol) — DM the bot to share feedback.
The registry split resolves the prior dispatch concerns: preparation and export now have independent precedence, the iterable fallback exists only in preparation, DBRX can prepare via its type-specific handler while independently taking the fused export path, and the handler registrations have been moved out of unified_export_hf.py. The focused dispatch tests cover these ordering properties, and I found no new correctness bug. The new files also use the canonical 2026 NVIDIA header.
Remaining reasons for owner sign-off:
- The CPU tests validate matching and a toy linear export, but still do not execute the actual DBRX/Llama4/GptOss/fused-experts/QuantMoELinear handlers. Given the operator's confirmation that current behavior is identical, this is not a code-level blocker, but the GPU/large-model paths remain the main unverified risk.
- This introduces a second registry mechanism alongside
QuantModuleRegistry/_DMRegistryCls. The source docstring now gives a reasonable rationale (export handlers do not replace classes, need structural matching, and preparation/export require independent ordering), but the PR body still does not record that design choice as required for a new registry subsystem. Please capture that short rationale in the PR description or have an owner explicitly sign off on it.
There was a problem hiding this comment.
🧹 Nitpick comments (1)
modelopt/torch/export/registry.py (1)
92-124: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winReplace the
assertguard with an explicit exception.
assert keys or predicate is not None(Line 111) is the only guard preventing a handler from being registered with no keys and no predicate. If that happens,match()'sif keys and .../if predicate is not None ...checks both short-circuit toFalse, so the handler unconditionally matches every module.assertstatements are stripped underpython -O, silently removing this protection. Bandit's B101 flags this pattern in non-test files, and the coding guideline disallows# nosecbypasses, so raising a real exception is required rather than relying onassert.🔒️ Proposed fix
- assert keys or predicate is not None, "register() requires at least one key or a predicate" + if not keys and predicate is None: + raise ValueError("register() requires at least one key or a predicate")Based on coding guidelines: "Bandit security checks must pass without exceptions.
# noseccomments are not allowed as a bypass for security checks."🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@modelopt/torch/export/registry.py` around lines 92 - 124, Replace the assert guard in ExportModuleRegistry.register with an explicit exception when neither keys nor predicate is provided, preserving the existing validation condition and message semantics. Do not use an assert or add a nosec bypass.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@modelopt/torch/export/registry.py`:
- Around line 92-124: Replace the assert guard in ExportModuleRegistry.register
with an explicit exception when neither keys nor predicate is provided,
preserving the existing validation condition and message semantics. Do not use
an assert or add a nosec bypass.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 4abd008a-e87b-47c8-bc2f-a279ce39fd69
📒 Files selected for processing (4)
modelopt/torch/export/hf_export_handlers.pymodelopt/torch/export/registry.pymodelopt/torch/export/unified_export_hf.pytests/unit/torch/export/test_export_registry.py
🚧 Files skipped from review as they are similar to previous changes (2)
- tests/unit/torch/export/test_export_registry.py
- modelopt/torch/export/unified_export_hf.py
2 participants
What does this PR do?
Type of change: ? refactor
Use registry-based module dispatch for unified_export_hf.py
export handlers do not replace classes, need structural matching, and preparation/export require independent ordering
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.).CONTRIBUTING.md:N/AAdditional Information
Summary by CodeRabbit
Summary