refactor(export): split unified_export_hf into layered modules - #2088
refactor(export): split unified_export_hf into layered modules#2088Fridah-nv wants to merge 6 commits into
Conversation
The transformers and diffusers export paths shared a file but almost no code: they meet only at the dispatch in export_hf_checkpoint. Move the diffusers half -- _export_diffusers_checkpoint, _postprocess_safetensors, _fuse_qkv_linears_diffusion and four helpers, 498 lines -- to unified_export_diffusers.py. unified_export_hf.py goes 1685 -> 1187. The diffusers-only imports (generate_diffusion_dummy_forward_fn, get_diffusion_components, merge_diffusion_checkpoint and the rest) leave with it; only is_diffusers_object, is_qkv_projection and get_qkv_group_key stay, for the dispatch check and the shared QKV fusion. The dispatch imports _export_diffusers_checkpoint lazily for now, because the diffusers module still imports the module-walking helpers back from here. The following commits move those out and the lazy import goes away. Two test files imported _postprocess_safetensors from the old location and are updated rather than shimmed; test_export_diffusers.py's monkeypatches move to the new module, since a `from X import Y` binding is not affected by patching Y on X. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
Every unified HF exporter runs the same preparation before packing a single weight: resolve the dtype, prepare MoE input quantizers, resmooth and fuse shared-input modules, adjust the quant config, and patch transformers while artifacts are written. That code sat in unified_export_hf.py, so the exporters had to import it back from the module that dispatches to them -- which is the only reason the lazy imports exist. Move those 13 symbols (364 lines) to hf_export_prep.py. It imports nothing else from the export package, so it sits at the bottom of the graph and the three exporters can depend on it without a cycle. unified_export_hf.py goes 1187 -> 823. The QKV fusion helpers travel with _fuse_shared_input_modules, so only is_diffusers_object remains of the diffusers imports here. External importers are repointed rather than shimmed: plugins/vllm_fakequant_hf.py for collect_shared_input_modules, and tests/gpu/.../test_fsdp2_export.py for requantize_resmooth_fused_llm_layers. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
_export_quantized_weight is the leaf of the pipeline -- it packs one module's weight and registers the scale buffers beside it -- but it lived in the same file as the exporters that call it, so moe_utils.py and hf_export_handlers.py had to reach it through function-local imports to dodge the cycle. Move it, _compressed_per_block_scale, _dispatch_export_handler and _process_quantized_modules (349 lines) to hf_weight_export.py. Like hf_export_prep, it imports nothing else from the export package. unified_export_hf.py goes 823 -> 474. Thirteen files are repointed rather than shimmed: moe_utils.py, hf_export_handlers.py, the two other exporters, and nine test modules. The patch targets in test_fused_experts.py move too, since patching a name on the old module no longer reaches the callers' bindings. The lazy imports in moe_utils.py and hf_export_handlers.py still point at the new module; hoisting them is the next commit. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
With preparation and weight packing in their own modules, the export package is
a DAG:
hf_export_prep, hf_weight_export -> (nothing in the package)
unified_export_hf_streaming -> prep
unified_export_diffusers -> prep, weight
unified_export_hf -> the three exporters, prep, weight
so the four function-local imports that existed only to dodge a cycle become
ordinary module-scope ones:
- moe_utils.py and hf_export_handlers.py reach _export_quantized_weight
directly. These predate this work -- they were dodging the cycle through
unified_export_hf.
- export_hf_checkpoint imports both the diffusers and streaming exporters at
module scope. The streaming one was added in #2008 with a comment saying it
could go once the shared helpers moved; this is that.
Verified by importing each of the eight modules first, and the package.
One test consequence, since hoisting changes name binding: the spies in
test_fused_experts.py patched _export_quantized_weight on the module that
defines it, which worked while moe_utils imported it lazily. Now that
moe_utils holds a module-scope reference, the patch has to target
moe_utils._export_quantized_weight instead.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
|
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. |
|
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 (8)
💤 Files with no reviewable changes (1)
🚧 Files skipped from review as they are similar to previous changes (3)
📝 WalkthroughWalkthroughThe change extracts Hugging Face preparation, quantized weight export, and Diffusers serialization into dedicated modules. Unified exporters, plugins, MoE helpers, and tests now use the new module boundaries. ChangesExport pipeline modularization
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant Caller
participant unified_export_hf
participant hf_export_prep
participant hf_weight_export
participant ExportRegistry
Caller->>unified_export_hf: start checkpoint export
unified_export_hf->>hf_export_prep: prepare model and quantizers
hf_export_prep-->>unified_export_hf: prepared model
unified_export_hf->>hf_weight_export: process quantized modules
hf_weight_export->>ExportRegistry: dispatch export handlers
ExportRegistry-->>hf_weight_export: exported weights and scales
hf_weight_export-->>unified_export_hf: processed checkpoint
unified_export_hf-->>Caller: saved checkpoint
Possibly related PRs
Suggested reviewers: Important Pre-merge checks failedPlease resolve all errors before merging. Addressing warnings is optional. ❌ Failed checks (1 warning, 1 inconclusive)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 1📝 Generate docstrings 💡
🧪 Generate unit tests (beta)
Comment |
|
/claude review |
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Warning
CodeRabbit couldn't request changes on this pull request because it doesn't have sufficient GitHub permissions.
Please grant CodeRabbit Pull requests: Read and write permission and re-run the review.
Actionable comments posted: 1
🧹 Nitpick comments (4)
modelopt/torch/export/hf_export_prep.py (2)
411-423: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winMove
import importlibto module scope.
importlibis a lightweight standard-library module. It is not an optional dependency and it creates no circular import. The coding guidelines require module-scope imports unless one of those justifications applies.♻️ Proposed fix
+import importlib import re import warningsdef _try_patch_module(mod_path: str) -> tuple[Any, Any] | None: """Try to patch revert_weight_conversion in a single module.""" - import importlib - try:Based on coding guidelines: "Keep imports at the top of Python source and test files; use local imports only for justified circular dependencies, optional dependencies, or unusually heavy imports, with a brief explanatory comment."
🤖 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/hf_export_prep.py` around lines 411 - 423, Move the importlib import from inside _try_patch_module to module scope with the other standard-library imports, then keep _try_patch_module’s importlib.import_module usage unchanged.Source: Coding guidelines
26-30: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd
__all__to the new module.This new module exports
collect_shared_input_modulesandrequantize_resmooth_fused_llm_layersto other packages. The coding guidelines require each module to declare its public API.♻️ Proposed addition
from .registry import ExportContext, PrepareMoEInputsRegistry +__all__ = ["collect_shared_input_modules", "requantize_resmooth_fused_llm_layers"] + try:Based on coding guidelines: "Define each module's public API with
__all__ = [...]."🤖 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/hf_export_prep.py` around lines 26 - 30, Add a module-level __all__ declaration in hf_export_prep.py listing the public functions collect_shared_input_modules and requantize_resmooth_fused_llm_layers, so the module explicitly defines its exported API.Source: Coding guidelines
modelopt/torch/export/unified_export_diffusers.py (1)
334-346: 🗄️ Data Integrity & Integration | 🔵 Trivial | ⚡ Quick win
_remove_promoted_quantizer_tensorsdeletes buffers it did not create.
_promote_quantizer_tensors_to_moduleregisters buffers only on modules whereis_quantlinear(sub_module)is true._remove_promoted_quantizer_tensorsdeletes the three buffer names from every submodule, without that filter and without tracking what was promoted.Two consequences follow. A non-quantlinear submodule that legitimately owns a buffer named
pre_quant_scale,svdquant_lora_a, orsvdquant_lora_bloses it after export. A quantlinear that already ownedpre_quant_scalehas it overwritten at line 324 and then deleted, so the original value is lost. Both contradict the docstring claim that the live module is unchanged after export.Track the promoted
(module, buffer_name)pairs and remove only those.♻️ Proposed refactor
-def _promote_quantizer_tensors_to_module(component: nn.Module) -> None: +def _promote_quantizer_tensors_to_module(component: nn.Module) -> None: @@ + promoted: list[tuple[nn.Module, str]] = [] for _, sub_module in component.named_modules(): if not is_quantlinear(sub_module): continue @@ if pre_quant_scale is not None: sub_module.register_buffer("pre_quant_scale", pre_quant_scale.detach().clone()) + promoted.append((sub_module, "pre_quant_scale")) @@ if lora_a is not None and lora_b is not None: sub_module.register_buffer("svdquant_lora_a", lora_a.detach().clone()) sub_module.register_buffer("svdquant_lora_b", lora_b.detach().clone()) + promoted.append((sub_module, "svdquant_lora_a")) + promoted.append((sub_module, "svdquant_lora_b")) + component._modelopt_promoted_export_buffers = promoted- for _, sub_module in component.named_modules(): - for buffer_name in ("svdquant_lora_a", "svdquant_lora_b", "pre_quant_scale"): - if buffer_name in getattr(sub_module, "_buffers", {}): - del sub_module._buffers[buffer_name] + promoted = getattr(component, "_modelopt_promoted_export_buffers", []) + for sub_module, buffer_name in promoted: + sub_module._buffers.pop(buffer_name, None) + if hasattr(component, "_modelopt_promoted_export_buffers"): + del component._modelopt_promoted_export_buffers🤖 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/unified_export_diffusers.py` around lines 334 - 346, Update _promote_quantizer_tensors_to_module and _remove_promoted_quantizer_tensors to track each (module, buffer_name) pair actually registered or overwritten during promotion, and remove only those tracked buffers during cleanup. Preserve any pre-existing buffers, including on quantlinear modules, and avoid deleting same-named buffers from non-quantlinear submodules while maintaining repeated-export module reuse.modelopt/torch/export/hf_export_handlers.py (1)
45-48: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winRemove the obsolete lazy-import comment.
The comment still describes a lazy import that this PR removed.
_export_quantized_weightnow resolves from the module-level import at line 25. The stale text tells the next reader that a cycle still forces a function-local import, which is the opposite of the dependency structure this PR establishes.♻️ Proposed cleanup
def _export_weight( module: nn.Module, ctx: ExportContext, weight_name: str = "weight", ) -> None: - # Imported lazily to avoid a cycle: unified_export_hf imports this module to - # install the built-in handlers while retaining this legacy helper's import path. - _export_quantized_weight(module, ctx.dtype, weight_name, _tied_cache=ctx.tied_cache)🤖 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/hf_export_handlers.py` around lines 45 - 48, Remove the obsolete lazy-import comment immediately above the _export_quantized_weight call; the function now uses the module-level import, so leave the call and its arguments unchanged.
🤖 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.
Inline comments:
In `@modelopt/torch/export/hf_export_prep.py`:
- Around line 251-254: Update the condition guarding the MoE quantization path
near is_moe(module) to handle a None quantization_format before performing the
substring check, and replace the fragile identity comparison against
QUANTIZATION_NONE with value inequality consistent with the existing usage.
Preserve the current AWQ and NVFP4_SVDQUANT selection behavior for non-None
formats.
---
Nitpick comments:
In `@modelopt/torch/export/hf_export_handlers.py`:
- Around line 45-48: Remove the obsolete lazy-import comment immediately above
the _export_quantized_weight call; the function now uses the module-level
import, so leave the call and its arguments unchanged.
In `@modelopt/torch/export/hf_export_prep.py`:
- Around line 411-423: Move the importlib import from inside _try_patch_module
to module scope with the other standard-library imports, then keep
_try_patch_module’s importlib.import_module usage unchanged.
- Around line 26-30: Add a module-level __all__ declaration in hf_export_prep.py
listing the public functions collect_shared_input_modules and
requantize_resmooth_fused_llm_layers, so the module explicitly defines its
exported API.
In `@modelopt/torch/export/unified_export_diffusers.py`:
- Around line 334-346: Update _promote_quantizer_tensors_to_module and
_remove_promoted_quantizer_tensors to track each (module, buffer_name) pair
actually registered or overwritten during promotion, and remove only those
tracked buffers during cleanup. Preserve any pre-existing buffers, including on
quantlinear modules, and avoid deleting same-named buffers from non-quantlinear
submodules while maintaining repeated-export module reuse.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 9046bcc1-c082-4c2a-b75b-ad9b86d70998
📒 Files selected for processing (20)
modelopt/torch/export/hf_export_handlers.pymodelopt/torch/export/hf_export_prep.pymodelopt/torch/export/hf_weight_export.pymodelopt/torch/export/moe_utils.pymodelopt/torch/export/plugins/vllm_fakequant_hf.pymodelopt/torch/export/unified_export_diffusers.pymodelopt/torch/export/unified_export_hf.pymodelopt/torch/export/unified_export_hf_streaming.pytests/gpu/torch/export/test_export_embedding.pytests/gpu/torch/export/test_export_weight_gpu.pytests/gpu/torch/export/test_fsdp2_export.pytests/gpu/torch/quantization/test_gptq.pytests/gpu_trtllm/torch/export/test_export_compressed_nvfp4.pytests/unit/torch/export/test_export_diffusers.pytests/unit/torch/export/test_export_registry.pytests/unit/torch/export/test_export_weight.pytests/unit/torch/export/test_nvfp4_utils.pytests/unit/torch/export/test_offload_export.pytests/unit/torch/export/test_unified_export_hf.pytests/unit/torch/quantization/plugins/test_fused_experts.py
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #2088 +/- ##
==========================================
- Coverage 78.72% 78.05% -0.67%
==========================================
Files 522 525 +3
Lines 60129 60546 +417
==========================================
- Hits 47335 47260 -75
- Misses 12794 13286 +492
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:
|
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
|
Claude review Scope: Full review. Trigger comment carried no scoping instructions. 20 files changed (1493+/1361-); I reviewed all 8 modelopt/ files and all 12 test files. Verification method Because this PR claims moved-verbatim, the highest-value check was proving that rather than re-reading the logic. I hashed every relocated block against
No function from the original file is missing, and no new function appeared. Also confirmed via grep that zero references to the old addresses remain anywhere in modelopt/, tests/, examples/, or docs/ — the 13-file repoint is complete. The layering claim holds The stated DAG checks out. Two things I specifically checked and found not to be problems:
Findings CRITICAL: 0 / IMPORTANT: 0 / SUGGESTION: 3 All three are comment/docstring residue from the split, not logic:
One more, outside the diff: On the flagged compatibility caveat The PR body raises that out-of-tree deep imports of Risk assessment Low. A refactor whose every moved byte is verifiably unchanged, whose call sites are all repointed with no shims left to drift, and which deletes four real import cycles rather than moving them. The one behavioral change (name binding on the hoisted import) is understood, documented in the PR body, and correctly handled in the two affected tests. The only follow-up worth doing is carrying the two orphaned comments to where their code went. No blocking issues. |
…split Review findings on the module split. HAS_DIFFUSERS was the real one. Replacing the `import diffusers` probe with an import from .diffusers_utils changed behavior: that module catches its own diffusers ImportError and still imports cleanly, so the except branch could never fire and the flag was unconditionally True. Verified: without diffusers it read True here while unified_export_diffusers read False -- two identically named flags disagreeing. Both now read diffusers_utils._HAS_DIFFUSERS, so there is one probe. unified_export_diffusers keeps a use-site `import diffusers` for the one place it needs __version__. Also from the split: - hf_export_prep wrapped the QKV helpers in an `except ImportError` that could not fire, whose None fallback would have turned a missing dependency into `TypeError: 'NoneType' object is not callable` at the call site. Removed. - Both new module docstrings claimed to depend on nothing else in the export package; each imports several leaf modules. They now state the real invariant: leaf helpers only, never an exporter. - hf_export_handlers kept the comment explaining a lazy import that commit 04c5904 deleted, describing a cycle that no longer exists. - registry.py, model_utils.py and the ptq skill reference still pointed at unified_export_hf.py for code that moved. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
…it documents Commit b3037f8 moved _revert_weight_conversion_noop and _patch_revert_weight_conversion to hf_export_prep.py but left the TODO explaining them behind in unified_export_hf.py, where it dangled between _export_transformers_checkpoint and export_speculative_decoding -- two functions it has nothing to do with. That note is the only record of the transformers 5.12.0 0-d-scalar bug and the condition for dropping the workaround, so detached it left the patch helpers with no rationale and pointed anyone revisiting them at the wrong file. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
Edwardf0t1
left a comment
Edwardf0t1
left a comment
There was a problem hiding this comment.
Reviewed as a refactor-correctness question rather than a re-read of the logic, since the PR claims pure movement.
Verification
- AST-level identity: parsed every top-level
def/classinmain:unified_export_hf.pyand in the union of the five resulting modules. All 40 symbols present, none duplicated across modules, and every body identical except the two intentional deltas (_export_diffusers_checkpoint's use-siteimport diffusers,export_hf_checkpoint's dropped lazy import). - Comments too — AST unparse drops them, so I diffed comment lines separately: exactly 2 lost across the whole split, both from the deleted lazy import. The transformers-5.12.0 TODO and every other inline note survived.
- No module-level mutable state moved: zero
globalstatements, no module-scope constants beyond__all__and the two import guards. - DAG holds: imported each of the 9 export modules first in a fresh interpreter, plus the package and
plugins/vllm_fakequant_hf— clean under every order. Also confirmed the handler-registration side effect survives:hf_weight_exportno longer transitively importshf_export_handlers, butexport/__init__.pyruns before any submodule, soExportModuleRegistryis populated (5 entries) even when onlyhf_weight_exportis imported. - Tests:
tests/unit/torch(excl.puzzletron, which fails to collect locally on unrelated missing deps) — 2256 passed, 0 failed. Import targets of all five changed GPU test files resolve against the new layout; thetest_fused_experts.pyretargets are correct,moe_utils._export_quantized_weightis the used-site binding the spies need.
LGTM. Four non-blocking notes inline.
On the shim question in the description: I'd skip them. requantize_resmooth_fused_llm_layers and collect_shared_input_modules are absent from __all__ and from docs/source; the only documented entry points (export_hf_checkpoint, plus _export_transformers_checkpoint as used by examples/llm_qat/export.py) all stayed put. Shims would reintroduce exactly the two-addresses-per-symbol problem the PR removes.
Minor: the line-count table in the description drifted after ee7f050/c927f89 — actual is 363/571/457/416, not 373/574/455/415.
|
|
||
| # _HAS_DIFFUSERS is diffusers_utils' own probe; re-deriving it here would drift, since | ||
| # that module imports cleanly whether or not diffusers is installed. | ||
| from .diffusers_utils import _HAS_DIFFUSERS as HAS_DIFFUSERS |
There was a problem hiding this comment.
Now that this import is unconditional, the flag it carries is dead weight. diffusers_utils.is_diffusers_object already early-returns False when _HAS_DIFFUSERS is false, so the only use site —
is_diffusers_obj = False
if HAS_DIFFUSERS:
is_diffusers_obj = is_diffusers_object(model)— can collapse to is_diffusers_obj = is_diffusers_object(model), and this alias import (plus its two-line comment) can go entirely.
The guard was load-bearing on main, where is_diffusers_object could be undefined if import diffusers failed; it isn't anymore. Dropping it also removes one of the two duplicate HAS_DIFFUSERS aliases that ee7f050 set out to de-duplicate — unified_export_diffusers would be the only module left holding one, and it needs its copy for the pipeline check.
| _dispatch_export_handler(name, sub_module, ctx) | ||
|
|
||
|
|
||
| def _export_transformers_checkpoint( |
There was a problem hiding this comment.
The layering stops one module short of the goal, and this function is why: the resident exporter stayed in the dispatch module, so plugins/hf_spec_export.py:275 and :445 still need function-local from ..unified_export_hf import _export_transformers_checkpoint to dodge the unified_export_hf → .plugins → hf_spec_export cycle. Those are the last cycle-dodging lazy imports in the package, and this PR's own framing is that such imports are the symptom worth removing.
The cut looks clean: .plugins is only touched by export_speculative_decoding (has_spec_opt, SpeculativeDecodingExporter, sanitize_hf_config_for_deployment), not by _export_transformers_checkpoint. So moving the resident exporter into its own module — mirroring the streaming and diffusers splits — would let both call sites hoist and leave this file as pure dispatch. examples/llm_qat/export.py:25 would need repointing too.
Follow-up, not this PR.
| Like :mod:`hf_export_prep`, this imports only leaf helpers (model_config, quant_utils, | ||
| registry) and never an exporter, so the exporters and the MoE/handler plugins can import | ||
| it directly instead of lazily. |
There was a problem hiding this comment.
This invariant is the whole point of the PR, and nothing pins it. CI can't catch a regression here: export/__init__.py always imports unified_export_hf first, so a future cycle stays hidden behind the "good" import order and only surfaces for someone importing a submodule in a context where the package init is already partially executed.
A subprocess-per-module import test would pin it cheaply — roughly what I ran by hand to check this PR:
MODULES = ["unified_export_hf", "unified_export_diffusers", "hf_export_prep",
"hf_weight_export", "unified_export_hf_streaming", "hf_export_handlers",
"moe_utils", "registry", "model_utils"]
@pytest.mark.parametrize("mod", MODULES)
def test_module_imports_first(mod):
subprocess.run([sys.executable, "-c", f"import modelopt.torch.export.{mod}"], check=True)Optional, but it's the difference between a documented invariant and an enforced one.
| ``` | ||
|
|
||
| **Known VLM export issue**: The export step (`requantize_resmooth_fused_llm_layers` in `unified_export_hf.py`) may try to run a dummy forward pass on the full VLM instead of the language model backbone. This currently only handles Nemotron VLMs. If hit, patch the export to use `is_multimodal_model()` for the VLM check instead of model-specific string matching. | ||
| **Known VLM export issue**: The export step (`requantize_resmooth_fused_llm_layers` in `hf_export_prep.py`) may try to run a dummy forward pass on the full VLM instead of the language model backbone. This currently only handles Nemotron VLMs. If hit, patch the export to use `is_multimodal_model()` for the VLM check instead of model-specific string matching. |
There was a problem hiding this comment.
One sibling of this pointer got missed. tests/_test_utils/torch/quantization/tied_modules.py:103 still says:
the
model_typegate inside_reorder_canonical_first(mirrors the existing whisper / nemotron-vl dispatch inunified_export_hf.py)
That dispatch is now hf_export_prep.py:266 — same class of stale reference ee7f050 fixed in registry.py and model_utils.py, just in a file this PR didn't otherwise touch.
2 participants
What does this PR do?
Type of change: Refactor (no functional change)
unified_export_hf.pyhad grown to 1685 lines and was the largest file inmodelopt/torch/export/. More importantly, it mixed four unrelated jobs — dispatch, the resident exporter, model-level preparation, and per-module weight packing — which forced the other exporters to import their shared helpers back from the module that dispatches to them. That cycle is why several function-local imports exist today.Splitting by layer rather than by size makes the package a DAG:
Four commits, each independently green so they can be reviewed one at a time:
602879b280unified_export_diffusers.py(498 lines)b3037f8910hf_export_prep.py(364 lines)04a7382b10hf_weight_export.py(349 lines)04c5904396Resulting layout:
unified_export_hf.pyunified_export_diffusers.pyhf_export_prep.pyunified_export_hf_streaming.pyhf_weight_export.pyThe payoff is the last commit. Three of the four removed lazy imports predate this work:
moe_utils.pyandhf_export_handlers.pyreached_export_quantized_weightthrough function-local imports purely to dodge the cycle, and #2008 added a third for the streaming dispatch with a comment saying it could go once the shared helpers moved. This is that.Usage
No API change.
export_hf_checkpointis unaffected and still dispatches to the right exporter:Testing
Run after each commit, not just at the end:
tests/unit— 3130 passed, 15 skippedtests/gpu/torch/export/+tests/gpu/torch/quantization/test_gptq.py— 123 passed, 2 skipped (both pre-existing:sm90requirement,INT4_AWQ_CFGon Qwen3 MoE)tests/gpu/torch/export/test_export_diffusers.pywas excluded from the local GPU run because it exceeds our relay's time limit; its unit-test counterpart passes, and CI covers it.Two review notes, both consequences of the mechanics rather than incidental:
moe_utils.py,hf_export_handlers.py,plugins/vllm_fakequant_hf.py, the other two exporters, and 9 test modules. Imports are updated rather than shimmed, so no symbol ends up with two addresses.moe_utilsnow holds a module-scope reference, so patching_export_quantized_weightwhere it is defined no longer intercepts it. The spies intest_fused_experts.pymove to patching where it is used (moe_utils._export_quantized_weight). Same reasontest_export_diffusers.py's monkeypatches move in commit 1.Before your PR is "Ready for review"
modelopt.torch.export.__all__is unchanged (export_hf_checkpoint,export_speculative_decoding), and both stay inunified_export_hf. Worth flagging one caveat: deep imports of two non-underscore internals,requantize_resmooth_fused_llm_layersandcollect_shared_input_modules, now resolve fromhf_export_prep. They were never in__all__, and every in-repo caller is updated, but out-of-tree code importing them directly fromunified_export_hfwould need a one-line change. Happy to add re-export shims if reviewers would rather not break that.CONTRIBUTING.md: N/A — no new dependencies; all code is moved verbatim within the repo.Additional Information
Follow-up to #2008. The extraction was suggested there by @Edwardf0t1, who scoped the diffusers block and
_export_quantized_weightas separate work; this PR does both plus the preparation layer, because splitting all three is what actually removes the cycles rather than relocating them.Deliberately left alone:
layer_utils.py(1991 lines) andquant_utils.py(1664), which are now the two largest files in the package. Both are worth a look, but neither is entangled with the exporter layering this PR is fixing.Summary by CodeRabbit
New Features
Bug Fixes
Refactor