refactor(export): split unified_export_hf into layered modules - #2088
refactor(export): split unified_export_hf into layered modules#2088Fridah-nv wants to merge 10 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. |
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughThe change extracts Hugging Face preparation, quantized weight export, and Diffusers serialization into dedicated modules. Unified exporters, plugins, MoE helpers, tests, documentation, and import-layering checks 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
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (5 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% 76.28% -2.44%
==========================================
Files 522 525 +3
Lines 60129 63407 +3278
==========================================
+ Hits 47335 48369 +1034
- Misses 12794 15038 +2244
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.
…module APIs Review follow-ups on the export split: - `unified_export_hf` no longer aliases `diffusers_utils._HAS_DIFFUSERS`. The guard was load-bearing only while `is_diffusers_object` could be undefined; now that the import is unconditional, that helper already early-returns `False` when diffusers is missing, so the dispatch collapses to a single call and `unified_export_diffusers` holds the sole remaining alias. - Declare `__all__` on the three modules the split added, per the coding standards. - Hoist `import importlib` in `hf_export_prep` to module scope; it is neither optional, heavy, nor circular. - Repoint the last stale `unified_export_hf.py` pointer, in the `tied_modules` test helper, at `hf_export_prep.py`. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
`export/__init__.py` always pulls `unified_export_hf` in first, so a reintroduced cycle resolves under that one "good" order and stays invisible to every existing test. It would surface only for someone importing a submodule while the package init is partially executed. Import each export module first in its own interpreter, which is what actually exercises the DAG the split established. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
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 (1)
tests/unit/torch/export/test_export_import_layering.py (1)
47-51: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winBound the child import.
If an import stalls,
subprocess.runwaits indefinitely. Pass a finite timeout shorter than thetests/unitlimit and includemodulein the timeout failure.🤖 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 `@tests/unit/torch/export/test_export_import_layering.py` around lines 47 - 51, Update the subprocess.run call in the import test to use a finite timeout shorter than the tests/unit limit, and catch the timeout failure so the resulting assertion or error includes the module name. Preserve the existing import command and captured output behavior.Source: Path instructions
🤖 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 `@tests/unit/torch/export/test_export_import_layering.py`:
- Around line 47-48: Update the subprocess invocation in the export import test
to bootstrap the package path without executing
modelopt.torch.export.__init__.py, then import the requested module via
sys.argv[1]. Pass module as a subprocess argument and add a timeout so hung
imports fail promptly, preserving the existing import validation behavior.
---
Nitpick comments:
In `@tests/unit/torch/export/test_export_import_layering.py`:
- Around line 47-51: Update the subprocess.run call in the import test to use a
finite timeout shorter than the tests/unit limit, and catch the timeout failure
so the resulting assertion or error includes the module name. Preserve the
existing import command and captured output behavior.
🪄 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: e70c29e8-ee1d-4064-b030-d4e0f8003818
📒 Files selected for processing (6)
modelopt/torch/export/hf_export_prep.pymodelopt/torch/export/hf_weight_export.pymodelopt/torch/export/unified_export_diffusers.pymodelopt/torch/export/unified_export_hf.pytests/_test_utils/torch/quantization/tied_modules.pytests/unit/torch/export/test_export_import_layering.py
🚧 Files skipped from review as they are similar to previous changes (3)
- modelopt/torch/export/hf_weight_export.py
- modelopt/torch/export/unified_export_diffusers.py
- modelopt/torch/export/hf_export_prep.py
…ng it The subprocess-per-module test did not test what it claimed. Importing `modelopt.torch.export.<submodule>` runs the package initializer first, and `export/__init__.py` imports its submodules in one fixed order, so all nine parametrizations replayed the identical import sequence — the requested module was never the first one loaded. It also cost ~2 min of wall clock for that. Read the graph instead of running it: parse the module-scope intra-package imports out of each source file, assert the result is acyclic, and assert the two leaves import no exporter at all. Only module-scope imports can form an import-time cycle, which is exactly what the parse sees, and the result is order-independent rather than dependent on whichever order `__init__` happens to use. `plugins` is folded in as one node so a cycle routed through a plugin still shows up. Verified against injected regressions: a module-scope `hf_export_prep` -> `unified_export_hf` edge is reported as `hf_export_prep -> unified_export_hf -> hf_export_prep`, and a function-local exporter import in a leaf — invisible to the cycle check, since it is the cycle-dodging shape the split removed rather than an import-time cycle — is caught by the leaf assertion. Runs in 0.3s. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
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
🤖 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 `@tests/unit/torch/export/test_export_import_layering.py`:
- Around line 65-70: The _relative_imports helper currently misses absolute
imports within modelopt.torch.export, allowing dependency cycles to evade the
graph tests. Update _relative_imports to parse ast.Import targets and level-zero
ast.ImportFrom modules beginning with modelopt.torch.export, recording the
relevant top-level imported module names consistently with relative imports; add
regression cases covering both import modelopt.torch.export.unified_export_hf
and from modelopt.torch.export import unified_export_hf.
🪄 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: 9f8a382f-45e7-4c90-b19d-f509a5639016
📒 Files selected for processing (1)
tests/unit/torch/export/test_export_import_layering.py
…elative The parser only recognized relative `ImportFrom` nodes, so `import modelopt.torch.export.x` and `from modelopt.torch.export[.x] import ...` were invisible — an exporter dependency could be added in either spelling and both graph tests would still pass. Not hypothetical: `plugins/vllm_fakequant_megatron.py` already imports absolutely, so real edges were being dropped. Parse `ast.Import` and level-zero `ast.ImportFrom` targets under the export package alongside the relative forms, and cover all five spellings with parser regression cases. Widening the parse surfaced a pre-existing module-scope cycle it had been hiding: `plugins` -> `unified_export_megatron` -> `plugins`. It predates the HF export split, sits entirely in the megatron half, and resolves today only because `export/__init__.py` imports `.plugins` first and what `unified_export_megatron` needs back are `plugins` submodules, which resolve against a partially initialized package. Left alone as out of scope here, but recorded as an explicit `KNOWN_CYCLE_EDGES` entry rather than silently excluded, so a new cycle still fails and the entry is asserted to still exist. Verified: all three absolute spellings of a leaf -> exporter edge are now reported by both the cycle check and the leaf assertion. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
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
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
tests/unit/torch/export/test_export_import_layering.py (2)
76-92: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winTraverse class bodies during import-scope analysis.
At Lines [89-92],
ClassDefcauses the walker to skip the complete class body. An import inside a class body runs when the containing module is imported. The graph can therefore miss an import cycle. Skip only function bodies and continue traversingClassDefnodes. Update the docstring at Lines [76-77].Proposed fix
- if module_scope_only and isinstance( - node, ast.FunctionDef | ast.AsyncFunctionDef | ast.ClassDef - ): + if module_scope_only and isinstance(node, ast.FunctionDef | ast.AsyncFunctionDef): continueAs per path instructions, tests must exercise the behavior they claim to validate.
🤖 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 `@tests/unit/torch/export/test_export_import_layering.py` around lines 76 - 92, Update the module-scope AST traversal to skip only FunctionDef and AsyncFunctionDef nodes, while continuing into ClassDef bodies so imports executed during class definition are analyzed. Revise the surrounding docstring to state that function bodies are skipped but class bodies are traversed, and add or update tests covering imports inside class bodies.Source: Path instructions
147-172: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winFix import-graph false negatives.
_find_cyclehandles self-cycles and disconnected components. The structural checks cover the listed leaves, exporters, and dispatcher dependencies._intra_package_imports(..., module_scope_only=True)skipsClassDefbodies, although class-body imports execute during module import. Include class-body imports in cycle detection while continuing to exclude function-local imports.KNOWN_CYCLE_EDGESare removed before cycle detection. A new cycle that reuses an excluded edge can remain hidden. Validate that the exclusion covers only the documented cycle.- Add a parser regression test for level-2 relative imports because plugin parsing uses
level=2.🤖 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 `@tests/unit/torch/export/test_export_import_layering.py` around lines 147 - 172, Update _intra_package_imports with module_scope_only=True to traverse ClassDef bodies while continuing to skip function-local imports, and add coverage for level-2 relative imports using the parser’s plugin configuration. Revise KNOWN_CYCLE_EDGES handling and its validation so excluded edges cannot hide any cycle beyond the documented one, while preserving _find_cycle behavior for self-cycles and disconnected components.
🧹 Nitpick comments (1)
tests/unit/torch/export/test_export_import_layering.py (1)
175-196: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd regression coverage for plugin-relative imports.
At Lines [175-196], the test always passes
level=1. The graph useslevel=2for files underpluginsat Lines [136-139]. Add cases forfrom ..unified_export_hf import ...andfrom .. import unified_export_hf, and pass the parameterized level to_intra_package_imports.Proposed test adjustment
- "source", + ("source", "level"), [ - pytest.param("from .unified_export_hf import export_hf_checkpoint", id="relative"), + pytest.param( + "from .unified_export_hf import export_hf_checkpoint", 1, id="relative" + ), + pytest.param( + "from ..unified_export_hf import export_hf_checkpoint", 2, id="plugin-relative" + ), + pytest.param("from .. import unified_export_hf", 2, id="plugin-relative-bare"), ], ) -def test_every_import_spelling_is_parsed(source): +def test_every_import_spelling_is_parsed(source, level): - found = _intra_package_imports(ast.parse(source), 1, module_scope_only=True) + found = _intra_package_imports(ast.parse(source), level, module_scope_only=True)As per path instructions, tests must exercise the behavior they claim to validate.
🤖 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 `@tests/unit/torch/export/test_export_import_layering.py` around lines 175 - 196, Extend test_every_import_spelling_is_parsed with parameterized cases for parent-relative imports using “from ..unified_export_hf” and “from .. import unified_export_hf”, and parameterize the expected import level alongside each source. Pass that level instead of the hardcoded 1 to _intra_package_imports so the test covers both regular package imports and plugin-relative level-2 imports.Source: Path instructions
🤖 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 `@tests/unit/torch/export/test_export_import_layering.py`:
- Around line 141-143: Update _import_graph so _find_cycle analyzes the complete
graph instead of deleting KNOWN_CYCLE_EDGES beforehand. Permit only the
documented exact two-node cycle between plugins and unified_export_megatron, and
reject cycles that add intermediate nodes or alternate paths; retain the
existing edge-presence validation where applicable.
---
Outside diff comments:
In `@tests/unit/torch/export/test_export_import_layering.py`:
- Around line 76-92: Update the module-scope AST traversal to skip only
FunctionDef and AsyncFunctionDef nodes, while continuing into ClassDef bodies so
imports executed during class definition are analyzed. Revise the surrounding
docstring to state that function bodies are skipped but class bodies are
traversed, and add or update tests covering imports inside class bodies.
- Around line 147-172: Update _intra_package_imports with module_scope_only=True
to traverse ClassDef bodies while continuing to skip function-local imports, and
add coverage for level-2 relative imports using the parser’s plugin
configuration. Revise KNOWN_CYCLE_EDGES handling and its validation so excluded
edges cannot hide any cycle beyond the documented one, while preserving
_find_cycle behavior for self-cycles and disconnected components.
---
Nitpick comments:
In `@tests/unit/torch/export/test_export_import_layering.py`:
- Around line 175-196: Extend test_every_import_spelling_is_parsed with
parameterized cases for parent-relative imports using “from ..unified_export_hf”
and “from .. import unified_export_hf”, and parameterize the expected import
level alongside each source. Pass that level instead of the hardcoded 1 to
_intra_package_imports so the test covers both regular package imports and
plugin-relative level-2 imports.
🪄 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: ad323521-8539-450e-b3cc-b279ccfa11a1
📒 Files selected for processing (1)
tests/unit/torch/export/test_export_import_layering.py
| if drop_known_cycles: | ||
| for src, dst in KNOWN_CYCLE_EDGES: | ||
| graph.get(src, set()).discard(dst) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | 🏗️ Heavy lift
Keep the known-cycle exception narrow.
At Lines [141-143], _import_graph deletes plugins -> unified_export_megatron before _find_cycle runs. This hides any new cycle that uses that edge. For example, plugins -> unified_export_megatron -> new_exporter -> plugins passes after the deletion. The check at Lines [205-211] only confirms that the edge still exists.
Run cycle analysis on the full graph. Allow only the exact documented two-node cycle. Alternatively, remove both known cycle edges and assert that no alternate path from unified_export_megatron to plugins exists.
Based on the provided import-graph implementation, the exception must cover only the pre-existing cycle.
🤖 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 `@tests/unit/torch/export/test_export_import_layering.py` around lines 141 -
143, Update _import_graph so _find_cycle analyzes the complete graph instead of
deleting KNOWN_CYCLE_EDGES beforehand. Permit only the documented exact two-node
cycle between plugins and unified_export_megatron, and reject cycles that add
intermediate nodes or alternate paths; retain the existing edge-presence
validation where applicable.
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