Add id/ tied_weight_names based export shared weight deduplication logic and remove legacy data_ptr approach - #2092
Add id/ tied_weight_names based export shared weight deduplication logic and remove legacy data_ptr approach #2092juhi10071998 wants to merge 15 commits into
Conversation
|
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:
📝 WalkthroughWalkthroughHF checkpoint export now resolves declared tied weights by model names. A shared ChangesTied-weight export
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant unified_export_hf
participant TiedGroupResolver
participant sync_tied_input_amax
participant _process_quantized_modules
participant postprocess_state_dict
unified_export_hf->>TiedGroupResolver: resolve declared tied parameter names
unified_export_hf->>sync_tied_input_amax: synchronize tied input amax values
unified_export_hf->>_process_quantized_modules: process quantized modules with shared resolver
unified_export_hf->>postprocess_state_dict: deduplicate state dict with shared resolver
postprocess_state_dict-->>unified_export_hf: return canonicalized checkpoint
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #2092 +/- ##
==========================================
- Coverage 78.74% 77.13% -1.62%
==========================================
Files 522 522
Lines 60368 61506 +1138
==========================================
- Hits 47538 47440 -98
- Misses 12830 14066 +1236
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:
|
a6a6f60 to
b190aa3
Compare
![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: 6
🧹 Nitpick comments (6)
tests/unit/torch/export/test_export_registry.py (1)
306-312: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd a test for resolver reuse.
This test covers the branch where
ExportContextbuilds a resolver. It does not cover the branch where the caller supplies one. That branch carries the newresolver=plumbing from_process_quantized_modules, so a regression that overwrote the field would pass this test and silently rebuild the map.💚 Proposed additional test
def test_export_context_reuses_caller_provided_resolver(): model = nn.Linear(2, 2) resolver = TiedGroupResolver(model) ctx = ExportContext(model=model, dtype=torch.float16, resolver=resolver) # A caller-supplied resolver is kept, so the export builds the alias map once. assert ctx.resolver is resolverAdd the import at the top of the file:
from modelopt.torch.export.model_utils import TiedGroupResolver🤖 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_registry.py` around lines 306 - 312, Add a unit test alongside test_export_context_builds_per_instance_resolver that constructs a TiedGroupResolver, passes it through ExportContext(..., resolver=resolver), and asserts ctx.resolver is the same instance. Import TiedGroupResolver from model_utils so the caller-supplied resolver path is covered without altering the existing per-instance resolver test.Source: Path instructions
modelopt/torch/export/unified_export_hf.py (1)
1028-1032: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚖️ Poor tradeoff
_reorder_canonical_firstis now redundant for declared ties.The comment states the reorder feeds "postprocess_state_dict's first-wins data_ptr dedup". The name-based pass now selects the canonical side from the declaration, so iteration order no longer decides which key survives for a declared tie.
The address pass still runs first-wins, but
_reorder_canonical_firstderives its patterns from_tied_weights_keys, and every tie declared there is already handled by name. The reorder therefore has no remaining effect on the keys it can match.Consider removing the call and the helper, or at minimum correct the comment so a later reader does not treat the address pass as authoritative. The helper is gated to DiffusionGemma, so confirm that export before removing it.
Run the following script to find every user of the reorder helper before removing it:
#!/bin/bash # Description: Locate all references to the canonical-first reorder helpers. set -euo pipefail rg -nP -C4 '_reorder_canonical_first|_collect_canonical_tied_patterns' --type=py🤖 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_hf.py` around lines 1028 - 1032, Remove the now-redundant _reorder_canonical_first call and its _collect_canonical_tied_patterns helper after confirming the repository search shows no other required users and the DiffusionGemma export remains covered by name-based tie handling. Also remove any imports or comments that only support this reorder, while preserving the existing quantized_state_dict processing flow.tests/unit/torch/export/test_unified_export_hf.py (1)
274-314: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd a test for a key marked by both dedup passes.
postprocess_state_dictusesdict.fromkeys(keys_to_delete)atmodelopt/torch/export/quant_utils.pyLine 1181 to tolerate a key that both the name pass and an earlier pass mark. No test drives that path, so a regression that reintroduces a doublepopon a missing key would pass.A declared alias key that also matches the LoRA filter covers it in one case.
💚 Proposed test for the double-marked key path
def test_postprocess_double_marked_alias_is_deleted_once(): """A key marked by both the LoRA filter and the name-based tie pass deletes cleanly.""" enc, dec = make_tied_linear_pair() parent = wrap_in_parent_with_tied_keys(enc, dec, decoder_canonical=True) resolver = TiedGroupResolver(parent) sd = { "encoder.weight": torch.randn(4, 4), # declared alias "decoder.weight": torch.randn(4, 4), # canonical } # is_modelopt_qlora routes "lora" keys into keys_to_delete as well. out = postprocess_state_dict( sd, maxbound=448, quantization=None, is_modelopt_qlora=False, resolver=resolver ) assert "decoder.weight" in out assert "encoder.weight" not in out🤖 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_unified_export_hf.py` around lines 274 - 314, Add a unit test near test_postprocess_name_based_drops_tied_expert_subtree_by_name that creates a declared encoder/decoder tied pair, includes the alias and canonical weight keys, and runs postprocess_state_dict with the resolver and LoRA filtering enabled as required to mark the alias in both deletion passes. Assert the canonical decoder key remains and the encoder alias is removed without error, covering single deletion of a double-marked key.Source: Path instructions
tests/unit/torch/quantization/plugins/test_fused_experts.py (1)
654-659: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winThe added
_tied_weights_keysdeclaration does not affect either test.Both tests call
_export_fused_experts(container, dtype)directly. That function no longer takes a resolver and no longer dedups, so nothing in this file reads_tied_weights_keys. Thetie=Trueandtie=Falsecases now differ only by whether the 3-D source Parameters are shared, and thetorch.equalassertion follows from that sharing alone.The class is named
TestExportFusedExpertsTiedDedup, which implies dedup coverage. Dedup coverage lives intests/unit/torch/export/test_unified_export_hf.py.Make the declaration load-bearing with a resolver assertion, or remove it so the fixture states only what it exercises.
💚 Proposed assertion that makes the declaration load-bearing
def test_tied_fused_experts_pack_independently_to_equal_values(self): """Tied FusedExperts pack independently (distinct storage), byte-identical values. There is no per-module dedup cache: each container splits and packs the shared source itself, so per-expert buffers have DIFFERENT data_ptrs but EQUAL bytes -- the duplicate keys are then dropped by name in postprocess_state_dict. """ parent = _build_two_moe_blocks(tie=True) + # The declaration in _build_two_moe_blocks is what postprocess_state_dict later + # uses to drop the encoder subtree, so assert the resolver reads it. + resolver = TiedGroupResolver(parent) + assert ( + resolver.container_group_key("encoder.experts", "gate_up_proj") + == "decoder.experts" + ) expert_type = type(parent.encoder.experts)Add the import at the top of the file:
from modelopt.torch.export.model_utils import TiedGroupResolver🤖 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/quantization/plugins/test_fused_experts.py` around lines 654 - 659, Remove the unused parent._tied_weights_keys declaration from the tie=True branch, since _export_fused_experts does not consume it and these tests only exercise shared 3-D Parameters. Keep tie_fused_experts_3d_params and the existing assertions unchanged.Source: Coding guidelines
modelopt/torch/export/quant_utils.py (2)
1055-1055: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueAnnotate the
resolverparameter in both public functions. Both functions acceptresolver=Nonewith no type annotation, while every other parameter in the same signatures is annotated. The shared cause is thatquant_utils.pyhas noTYPE_CHECKINGimport forTiedGroupResolver, which a runtime import would make circular.registry.pyalready solves this with aTYPE_CHECKINGblock plus a quoted annotation.
modelopt/torch/export/quant_utils.py#L1055-L1055: annotate asresolver: "TiedGroupResolver | None" = Noneinpostprocess_state_dict.modelopt/torch/export/quant_utils.py#L1630-L1630: annotate asresolver: "TiedGroupResolver | None" = Noneinsync_tied_input_amax.Add the guarded import once at module scope:
from typing import TYPE_CHECKING if TYPE_CHECKING: from .model_utils import TiedGroupResolver🤖 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/quant_utils.py` at line 1055, Update modelopt/torch/export/quant_utils.py at lines 1055-1055 and 1630-1630 to annotate the resolver parameters in postprocess_state_dict and sync_tied_input_amax as "TiedGroupResolver | None", preserving the default None. Add a single module-scope TYPE_CHECKING guard importing TiedGroupResolver from .model_utils to avoid a runtime circular import.
1653-1658: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winMove the
defaultdictimport to module scope and justify the localTiedGroupResolverimport.
from collections import defaultdictis a stdlib import inside the function. It is not a circular dependency, an optional dependency, or an unusually heavy import, so it belongs at module scope.The
TiedGroupResolverimport is a justified circular-import workaround, but it carries no explanatory comment.♻️ Proposed import cleanup
- from collections import defaultdict - - from .model_utils import TiedGroupResolver + # Imported here to avoid a circular import: model_utils imports from this module. + from .model_utils import TiedGroupResolverAdd the stdlib import at the top of the file:
from collections import defaultdictAs per 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/quant_utils.py` around lines 1653 - 1658, Move the defaultdict import from the function to module scope, and retain the local TiedGroupResolver import only as a documented circular-dependency workaround by adding a brief explanatory comment next to it. Keep the resolver initialization behavior unchanged.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.
Inline comments:
In `@CHANGELOG.rst`:
- Line 24: Update the CHANGELOG entry to remove claims that a fused-MoE cache
survives or provides a resident-path optimization. State that fused-MoE
in-memory tied caching and per-module dedup caching were removed, while
preserving the description of name-based tied-weight resolution and
authoritative dedup.
In `@modelopt/torch/export/model_utils.py`:
- Around line 407-440: Restrict alias-prefix deduplication in alias_prefix_pairs
and postprocess_state_dict to declared aliased parameters and their exported
companion keys, rather than every descendant under the alias module prefix.
Preserve untied sibling parameters by ensuring canonical_state_dict_key only
rewrites eligible exported keys. Add a regression test covering an alias module
with an untied sibling and verify the sibling remains in the state dict.
In `@modelopt/torch/export/quant_utils.py`:
- Around line 1132-1142: Update the name-alias pass around
resolver.canonical_state_dict_key so it does not mark key when its canonical_key
is already present in keys_to_delete; maintain a set for marked keys and skip
such counterparts, ensuring bidirectional aliases retain at least one state-dict
entry.
In `@modelopt/torch/export/registry.py`:
- Around line 64-70: Create a single TiedGroupResolver before export preparation
and reuse it throughout the export flow, including streaming. Pass that resolver
into every ExportContext used by _prepare_moe_inputs() and
_process_quantized_modules(), and reuse it for final post-processing instead of
constructing additional resolvers; keep ExportContext.__post_init__ as the
fallback only when no resolver is supplied.
In `@modelopt/torch/export/unified_export_hf.py`:
- Around line 991-995: Update the comments in
modelopt/torch/export/unified_export_hf.py at lines 991-995 and 914-916: remove
the per-module MoE cache from the resolver consumer list near
TiedGroupResolver(model), and replace the obsolete per-call MoE dedup-cache
description with wording that ExportContext builds a resolver when the caller
does not provide one.
In `@tests/unit/torch/export/test_offload_export.py`:
- Around line 302-318: Update the docstring of
test_tied_weights_exported_independently_without_cache to describe only the
independent packing and byte-identical weight behavior it actually verifies;
remove the claim that it guards or covers the offload path.
---
Nitpick comments:
In `@modelopt/torch/export/quant_utils.py`:
- Line 1055: Update modelopt/torch/export/quant_utils.py at lines 1055-1055 and
1630-1630 to annotate the resolver parameters in postprocess_state_dict and
sync_tied_input_amax as "TiedGroupResolver | None", preserving the default None.
Add a single module-scope TYPE_CHECKING guard importing TiedGroupResolver from
.model_utils to avoid a runtime circular import.
- Around line 1653-1658: Move the defaultdict import from the function to module
scope, and retain the local TiedGroupResolver import only as a documented
circular-dependency workaround by adding a brief explanatory comment next to it.
Keep the resolver initialization behavior unchanged.
In `@modelopt/torch/export/unified_export_hf.py`:
- Around line 1028-1032: Remove the now-redundant _reorder_canonical_first call
and its _collect_canonical_tied_patterns helper after confirming the repository
search shows no other required users and the DiffusionGemma export remains
covered by name-based tie handling. Also remove any imports or comments that
only support this reorder, while preserving the existing quantized_state_dict
processing flow.
In `@tests/unit/torch/export/test_export_registry.py`:
- Around line 306-312: Add a unit test alongside
test_export_context_builds_per_instance_resolver that constructs a
TiedGroupResolver, passes it through ExportContext(..., resolver=resolver), and
asserts ctx.resolver is the same instance. Import TiedGroupResolver from
model_utils so the caller-supplied resolver path is covered without altering the
existing per-instance resolver test.
In `@tests/unit/torch/export/test_unified_export_hf.py`:
- Around line 274-314: Add a unit test near
test_postprocess_name_based_drops_tied_expert_subtree_by_name that creates a
declared encoder/decoder tied pair, includes the alias and canonical weight
keys, and runs postprocess_state_dict with the resolver and LoRA filtering
enabled as required to mark the alias in both deletion passes. Assert the
canonical decoder key remains and the encoder alias is removed without error,
covering single deletion of a double-marked key.
In `@tests/unit/torch/quantization/plugins/test_fused_experts.py`:
- Around line 654-659: Remove the unused parent._tied_weights_keys declaration
from the tie=True branch, since _export_fused_experts does not consume it and
these tests only exercise shared 3-D Parameters. Keep
tie_fused_experts_3d_params and the existing assertions unchanged.
🪄 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: 2243e3ff-b346-4c7e-990e-640cf314822e
📒 Files selected for processing (11)
CHANGELOG.rstmodelopt/torch/export/hf_export_handlers.pymodelopt/torch/export/model_utils.pymodelopt/torch/export/moe_utils.pymodelopt/torch/export/quant_utils.pymodelopt/torch/export/registry.pymodelopt/torch/export/unified_export_hf.pytests/unit/torch/export/test_export_registry.pytests/unit/torch/export/test_offload_export.pytests/unit/torch/export/test_unified_export_hf.pytests/unit/torch/quantization/plugins/test_fused_experts.py
aa4511d to
1f5f360
Compare
Why
|
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
modelopt/torch/export/quant_utils.py (2)
1132-1154: 🗄️ Data Integrity & Integration | 🟠 Major | 🏗️ Heavy liftRestrict name-based removal to entries derived from declared tied parameters.
A declaration for
encoder.weight -> decoder.weightcreates anencoder -> decoderprefix mapping. The current pass also removesencoder.biaswhendecoder.biasexists, even when only the weights are tied. The exported checkpoint then loses an untied parameter.Only rewrite packed weights and scale entries that derive from the declared parameter. Keep unrelated parameters and buffers. Add a regression with tied weights and distinct biases.
🤖 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/quant_utils.py` around lines 1132 - 1154, Restrict the name-based deduplication loop around resolver.canonical_state_dict_key to keys derived from the explicitly declared tied parameter, rather than every key sharing its alias prefix. Ensure tied packed-weight and scale entries are still deduplicated, while unrelated parameters such as distinct encoder.bias and decoder.bias remain in post_state_dict. Add a regression covering tied weights with separate biases.
1177-1189: 🗄️ Data Integrity & Integration | 🟠 Major | 🏗️ Heavy liftUse complete view metadata before storage-based removal.
The current identity does not include dtype, shape, or stride. A tensor and its transpose can have the same
data_ptr()and byte count but serialize to different logical values. This loop removes one required state-dict key.Deduplicate only tensors with equivalent serialization metadata. Preserve or clone non-equivalent shared views. Add a regression for equal-span transposed views.
🤖 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/quant_utils.py` around lines 1177 - 1189, Update the tensor identity in the shared-storage deduplication loop to include dtype, shape, and stride alongside device, data_ptr, and byte count. Only remove a key when the complete serialization metadata matches; preserve or clone non-equivalent shared views, including equal-span transposed tensors, and add a regression covering that case.
🤖 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.
Outside diff comments:
In `@modelopt/torch/export/quant_utils.py`:
- Around line 1132-1154: Restrict the name-based deduplication loop around
resolver.canonical_state_dict_key to keys derived from the explicitly declared
tied parameter, rather than every key sharing its alias prefix. Ensure tied
packed-weight and scale entries are still deduplicated, while unrelated
parameters such as distinct encoder.bias and decoder.bias remain in
post_state_dict. Add a regression covering tied weights with separate biases.
- Around line 1177-1189: Update the tensor identity in the shared-storage
deduplication loop to include dtype, shape, and stride alongside device,
data_ptr, and byte count. Only remove a key when the complete serialization
metadata matches; preserve or clone non-equivalent shared views, including
equal-span transposed tensors, and add a regression covering that case.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 8282bfac-b764-4551-9beb-7d3be78deb0a
📒 Files selected for processing (5)
modelopt/torch/export/model_utils.pymodelopt/torch/export/quant_utils.pymodelopt/torch/export/unified_export_hf.pytests/unit/torch/export/test_offload_export.pytests/unit/torch/export/test_unified_export_hf.py
🚧 Files skipped from review as they are similar to previous changes (3)
- modelopt/torch/export/model_utils.py
- tests/unit/torch/export/test_offload_export.py
- modelopt/torch/export/unified_export_hf.py
cjluo-nv
left a comment
cjluo-nv
left a comment
There was a problem hiding this comment.
Bot review (bedrock-claude-opus-5) — DM the bot to share feedback.
Replaces the data_ptr()-keyed tied-weight dedup in HF export with a name-based TiedGroupResolver built from _tied_weights_keys / tie_word_embeddings, deletes the dense and fused-MoE pack-time dedup caches, and makes postprocess_state_dict the single dedup authority. The motivation (address identity is wrong under FSDP gather/offload and can be recycled) is well argued and the CPU unit tests are a genuine improvement over the old address-based tests. However I think there are correctness risks that should be settled before merge:
- Declared-but-not-actually-tied is not verified (highest risk).
_build_tied_alias_maprecords an alias purely because a name matches a dict-style_tied_weights_keyspattern; it never checks that the alias and canonical Parameters are the same object. Recent transformers declares the dict form on the class independently ofconfig.tie_word_embeddings, so a model with tying disabled could havelm_head.weightdropped from the exported checkpoint. The streaming exporter in this very repo already guards this exact case ("_tied_weights_keyscan list keys whose weights are not actually shared … which would incorrectly droplm_head.weight"); the new path has no equivalent guard and no test. - Prefix rewriting is broader than the declaration.
alias_prefix_pairs()+canonical_state_dict_key()turn a per-parameter tie into a whole-subtree prefix substitution, so untied siblings under an alias module prefix are dropped whenever a same-named canonical key exists. Given the bug being fixed is "a false dedup dropped an independent expert weight", a cheap shape/dtype (or value) check before dropping would be a good backstop. - Partial-quantization ties become inconsistent — the deleted
test_export_quantized_weight_skips_alias_when_one_tied_side_is_unquantizedscenario now yields orphanweight_scale/input_scalekeys on the alias side with itsweightdropped; the drop is per-key rather than per-group and that coverage wasn't replaced. - Design question not addressed in the PR body: pre-pack
nn.Parameterobject identity (groupnamed_parameters(remove_duplicate=False)byid) is stable under FSDP resharding/offload and immune to allocator reuse, and would give exact tie groups without_canonical_via_pattern_pair, backreference templates, or prefix rewriting — with declarations used only to pick the canonical name. Please explain why regex name resolution was preferred. - Smaller items:
ExportContext.resolveris no longer read by any handler yet is built eagerly (streaming path builds two unused resolvers over a large model);sync_tied_input_amaxsilently loses amax merging for undeclared shares that the address backstop still drops;resolverparams are untyped; the PR body claims the fused-MoE cache is "kept purely as a resident-path optimization" while the diff (and CHANGELOG) removes it entirely, and the Testing section says e2e is "in progress" while Additional Information reports it complete. Also the Slack description cites NVBug 6530966 while the body cites 6525352 — worth reconciling. The CHANGELOG entry is a 15-line paragraph where neighbours are one-liners.
No licensing concerns (no license files or third-party code touched), size is reasonable (+644/-311), and no prompt-injection content was observed.
| if not isinstance(tied, dict) or not tied: | ||
| continue | ||
| prefix = f"{mod_name}." if mod_name else "" | ||
| plen = len(prefix) |
There was a problem hiding this comment.
Bot comment.
param_names is only used to match names — nothing here verifies that the alias and canonical parameters are actually the same object, so a declared-but-unapplied tie becomes a real dedup entry and postprocess_state_dict will drop the alias key.
That is a live scenario: recent transformers declares the dict form of _tied_weights_keys (e.g. {"lm_head.weight": "model.embed_tokens.weight"}) on the class regardless of config.tie_word_embeddings, and unified_export_hf_streaming.py already documents exactly this hazard:
Only apply when
tie_word_embeddings=True:_tied_weights_keyscan list keys whose weights are not actually shared … which would incorrectly droplm_head.weight.
Since the resolver is built before packing, you can settle it exactly: keep the {name: param} mapping from named_parameters(remove_duplicate=False) and only record alias -> canonical when params[alias] is params[canonical] (object identity survives FSDP resharding and offload, unlike data_ptr). Please add a unit test with a dict-style declaration over untied parameters asserting neither key is dropped.
There was a problem hiding this comment.
Fixed in 7d5155a9f: _build_tied_alias_map now confirms each declared dict-style tie is actually applied — alias and canonical must resolve to the same live nn.Parameter. That check is done pre-pack while every parameter is resident, so no freed-then-reused address can forge it. A class-level _tied_weights_keys with tie_word_embeddings=False therefore no longer drops an independent lm_head.weight. New test test_build_tied_alias_map_skips_declared_but_unapplied_tie.
| return head | ||
|
|
||
|
|
||
| def _canonical_via_pattern_pair(alias_pat: str, canonical_pat: str, name: str) -> str | None: |
There was a problem hiding this comment.
Bot comment.
This heuristic (longest common suffix of two pattern strings, plus a "literal head" check that permits .) can cut a regex mid-token and silently produce a canonical name that the canonical pattern would never match. Consider validating the result before accepting it — e.g. re.fullmatch(canonical_pat, canonical_rel) (or re.search) and falling back to alias_re.sub(...) otherwise. Right now the only coverage is the happy-path DiffusionGemma shape; a declaration whose trailing structures are not actually parallel would resolve to a wrong canonical name and drop a real weight.
There was a problem hiding this comment.
Fixed in ad17a5a95: _canonical_via_pattern_pair now validates its result with re.fullmatch(canonical_pat, candidate) and returns None (caller falls back to alias_re.sub) when the character-based common-suffix split lands mid-token, so a name the pattern would never produce is never emitted as canonical.
| ``encoder…experts.gate_up_proj``); stripping the trailing parameter component | ||
| yields the owning-module prefix. State-dict dedup rewrites any exported key | ||
| under an alias prefix — packed weight, ``weight_scale`` / ``weight_scale_2`` / | ||
| ``input_scale``, and per-expert splits like ``…experts.3.gate_proj.weight`` — |
There was a problem hiding this comment.
Bot comment.
Widening a per-parameter declaration into a module-prefix rewrite is broader than what the model declared: every exported key under the alias prefix is rewritten, so untied siblings get dropped whenever a same-named canonical key exists (e.g. a tie declared only on …experts.gate_up_proj/down_proj also collapses any other key in that container; lm_head.weight -> …embed_tokens.weight makes every lm_head.* scale collide with …embed_tokens.*).
Given the bug this PR fixes is precisely "dedup falsely dropped an independent weight", please either restrict the rewrite to keys derived from the declared parameter (the weight plus its known scale suffixes and the per-expert split pattern) or add a cheap shape/dtype equality check on the two tensors before appending to keys_to_delete.
There was a problem hiding this comment.
Addressed by the atomic drop in 7d5155a9f: postprocess_state_dict now drops an alias module-prefix's keys all-or-nothing, and only when every key has a canonical counterpart. An untied sibling under an alias prefix (no canonical twin) makes the group incomplete, so nothing is dropped. The declared-but-unapplied object-identity gate (see the model_utils.py:320 thread) further keeps a falsely-declared prefix out of the map entirely.
| if resolver is not None: | ||
| alias_prefixes = resolver.alias_prefix_pairs() | ||
| if alias_prefixes: | ||
| for key in post_state_dict: |
There was a problem hiding this comment.
Bot comment.
The drop is decided per key, so a tie whose two sides have different quantization states is left inconsistent: if the alias side is quantized and the canonical is not, encoder.weight is dropped (canonical weight exists) while encoder.weight_scale / encoder.input_scale survive (no canonical counterpart) — the checkpoint then carries scales for a weight that is gone. This is the scenario the now-deleted test_export_quantized_weight_skips_alias_when_one_tied_side_is_unquantized covered; there is no replacement test. Suggest treating the alias key group atomically (only drop when the canonical side has the matching key set, otherwise keep both and warn) and re-adding coverage.
There was a problem hiding this comment.
Fixed in 7d5155a9f: the drop is now atomic per alias module prefix — if any key in the group lacks a canonical counterpart (e.g. the alias is quantized but the canonical is not), none are dropped, so weight_scale/input_scale are never orphaned. New test test_postprocess_keeps_both_sides_when_tied_quant_state_differs replaces the deleted one.
| maxbound: float, | ||
| quantization: str | None, | ||
| is_modelopt_qlora: bool = False, | ||
| resolver=None, |
There was a problem hiding this comment.
Bot comment.
resolver=None is untyped in this public helper (same in sync_tied_input_amax). Please annotate resolver: "TiedGroupResolver | None" = None with a TYPE_CHECKING import, matching how registry.py types the same field.
There was a problem hiding this comment.
Fixed: annotated resolver: "TiedGroupResolver | None" = None on both postprocess_state_dict and sync_tied_input_amax (7d5155a9f), and in ad17a5a95 the TiedGroupResolver import was hoisted to module level (model_utils imports only stdlib + torch.nn, so there is no cycle).
| key = ("moe", first_proj.data_ptr(), m.down_proj.data_ptr()) | ||
| by_dp[key].append(m) | ||
| gk = resolver.container_group_key(name, first_proj_attr) | ||
| if gk is not None: |
There was a problem hiding this comment.
Bot comment.
Behavioural regression for undeclared ties: the previous grouping keyed on weight.data_ptr() and so merged amaxes for genuinely shared Parameters even when the model never declared the tie. Now only declared ties are grouped — but the address backstop in postprocess_state_dict still drops one side of an undeclared share, so the retained side's input_scale no longer covers the dropped side's activation range (silent clipping at inference, exactly the failure mode this helper exists to prevent). Grouping by nn.Parameter object identity here (id(m.weight)) would restore that coverage without reintroducing data_ptr fragility.
There was a problem hiding this comment.
Fixed in 7d5155a9f: sync_tied_input_amax gains an id(weight) fallback so an undeclared physical share — which the address backstop still collapses — has its input amaxes max-merged too, so the surviving input_scale covers both sides (no silent clipping). Grouping by object id (pre-pack, all resident) can't collide the way data_ptr can. New test test_sync_tied_input_amax_merges_undeclared_shared_weight.
| self.tied_cache = None | ||
| self.moe_tied_cache = None | ||
| # Import here to avoid a circular import at module load time. | ||
| from .model_utils import TiedGroupResolver |
There was a problem hiding this comment.
Bot comment.
After this change no handler reads ctx.resolver (both hf_export_handlers.py call sites now pass nothing), so this field is unused state that is nevertheless built eagerly. _export_transformers_checkpoint_streaming constructs two contexts (its own plus the one inside _prepare_moe_inputs), each building a full alias map that nobody reads — and _build_tied_alias_map is O(#declaring modules × #patterns × #params) regex searches, which is not free on a MiniMax-scale model. Either drop resolver from ExportContext or make it lazily computed.
There was a problem hiding this comment.
Fixed in 7d5155a9f: dropped the ExportContext.resolver field and its __post_init__ build. No handler read it; the driver builds one resolver and passes it directly to sync_tied_input_amax / postprocess_state_dict (its only consumers). This removes the redundant per-context alias-map build, including the streaming path's two contexts.
| @@ -77,8 +48,6 @@ def _delete_fused_moe_source_attrs(module: nn.Module) -> None: | |||
| def _export_fused_experts( | |||
There was a problem hiding this comment.
Bot comment.
Removing the fused-MoE cache means each tied experts container now re-unpacks and re-packs all experts, and both packed copies stay resident until postprocess_state_dict drops one — a compute and peak-memory increase on the resident path for tied-expert models (DiffusionGemma). The PR body states the fused-MoE cache is "kept purely as a resident-path compute/memory optimization … name-keyed", which contradicts this diff and the CHANGELOG; please update the body and confirm the DiffusionGemma-26B export peak memory did not regress. Note the correctness goal doesn't require deleting the cache — it can be re-keyed on the resolver's container group key.
There was a problem hiding this comment.
The fused-MoE cache is removed (not kept) — the PR body has been corrected to match the diff/CHANGELOG. Both the dense and fused-MoE caches are gone; the name-based drop is the single authority, and the has_non_resident_weights guard is removed with them. Peak save-memory does rise for tied-MoE (each side is re-packed rather than aliased) and the informational total_parameters index counts tied experts per side (25.8B vs 14.4B on DiffusionGemma-26B) — on-disk weights are byte-identical to the known-good baseline. Re-verified on HSG: DiffusionGemma 47067 tensors / 0 leaked encoder-expert keys, MiniMax 191211 / 15872 experts / 0 missing.
Design note: why tied-weight dedup uses both declared names and object identityA few review comments (thanks @cjluo-nv) touch the same question — why introduce object-identity checks alongside the name-based resolver, isn't one signal enough? Writing the rationale down here since it motivates the resolver, the Dedup asks two independent questions, and no single signal answers both. To safely drop alias
Each single-signal approach has a blind spot:
So the two are combined, each covering the other's blind spot:
Concretely, the resulting rules:
In short: physical sameness (id) and re-tiability (name) together are the actual definition of "safe to dedup" — dropping either check re-introduces one of the failure modes above. |
Edwardf0t1
left a comment
Edwardf0t1
left a comment
There was a problem hiding this comment.
Follow-up review pass. Skipping everything already raised in the existing threads (declared-but-not-tied / tie_word_embeddings, prefix rewriting broader than the declaration, per-key drop leaving orphaned scales, the _canonical_via_pattern_pair mid-token cut, the removed fused-MoE cache, the bidirectional map, and the stale MoE-cache comments) — the 5 inline comments below are points not covered by any existing thread.
| if key in already_marked: | ||
| continue | ||
| if isinstance(value, torch.Tensor) and value.data_ptr() != 0: | ||
| tensor_id = (value.device, value.data_ptr(), value.numel() * value.element_size()) |
There was a problem hiding this comment.
Bot comment.
Adding the size component makes this backstop disagree with the writer it exists to protect. safetensors groups by storage identity, not by tensor extent — _find_shared_tensors keys on (v.device, storage_ptr(v), storage_size(v)) — so a base tensor and a shorter view of it land in the same group and save_file raises. This key uses value.data_ptr() + numel * element_size, so that pair gets two distinct ids, both survive postprocess_state_dict, and the export dies at write time with RuntimeError: Some tensors share memory after all the packing work is done.
Reproduced on the exact shape the new test pins:
base = torch.arange(4, dtype=torch.float32)
sd = {"short": base[:2], "long": base}
# tensor_id here: {'short': (ptr, 8), 'long': (ptr, 16)} -> not collapsed
_find_shared_tensors(sd) # [{'long', 'short'}] -> save_file RAISESSo test_postprocess_state_dict_preserves_tensors_with_different_byte_ranges asserts precisely the state that makes the subsequent save_file fail. The false-positive it was added to fix is real (two independent tensors whose addresses coincide), but the discriminator should match safetensors' notion of sharing — e.g. key on (device, storage_ptr, storage_size) via safetensors.torch.storage_ptr/storage_size, or on value.untyped_storage().data_ptr() — rather than on tensor length.
There was a problem hiding this comment.
Fixed in ad17a5a95: the backstop now keys on safetensors' own shared-storage identity via safetensors.torch.storage_ptr / storage_size — (device, storage_ptr, storage_size). A base tensor and a view of it share storage, so they now collapse here instead of both surviving and crashing save_file. Since two distinct live tensors cannot share a data_ptr without sharing storage, storage identity never false-collapses independent weights. The ..._preserves_tensors_with_different_byte_ranges test is replaced by ..._collapses_view_and_base_sharing_storage.
| if gk is None: | ||
| return None | ||
| suffix = f".{first_proj_attr}" | ||
| return gk.removesuffix(suffix) |
There was a problem hiding this comment.
Bot comment.
removesuffix is a silent no-op when the resolved canonical does not end with .{first_proj_attr} — e.g. a declaration that maps …experts.gate_up_proj onto a differently-named canonical projection, or a re.sub fallback that produced a mangled canonical. The method then returns a full parameter name as a container key, which will not equal the key computed from the container's other projection, so sync_tied_input_amax quietly puts the two tied containers in different groups, skips the merge (len(modules) < 2), and the retained side's input_scale under-covers the dropped side — the exact clipping this function exists to prevent, with no warning.
Suggest returning None (or warning) when not gk.endswith(suffix), so a misresolved canonical degrades to "untied" instead of to a wrong group.
There was a problem hiding this comment.
Fixed in ad17a5a95: container_group_key returns None when the resolved canonical does not end in .{first_proj_attr} (the case where removesuffix was a silent no-op), so a misresolved canonical degrades to 'untied' instead of forming a wrong group and skipping the amax merge.
| @@ -1053,7 +1038,11 @@ def _export_transformers_checkpoint( | |||
| quantized_state_dict = _reorder_canonical_first(quantized_state_dict, model) | |||
There was a problem hiding this comment.
Bot comment.
This call is obsolete under the new scheme and the comment above it is now inaccurate. _reorder_canonical_first's own docstring states its purpose — "Lets the downstream first-wins data_ptr dedup keep canonical names" — but the name-based pass now always keeps the canonical side regardless of iteration order, and the data_ptr pass is a backstop for undeclared shares only (which no reordering can help). For every declared tie the reorder is a no-op on the outcome.
What remains is ~80 lines in model_utils.py (_collect_canonical_tied_patterns + _reorder_canonical_first, a regex-token heuristic behind a DiffusionGemma class-name allowlist) that parse the same _tied_weights_keys structure as _build_tied_alias_map — two parsers of one declaration format that can drift apart, plus a model-specific special case layered on the new shared infrastructure. Since the PR's stated goal is to make postprocess_state_dict the single dedup authority, dropping this path (and its helper + import) would finish that consolidation; if there is a DiffusionGemma case the resolver still cannot express, that is worth a comment here, because right now it reads as leftover scaffolding.
There was a problem hiding this comment.
Done in ad17a5a95: removed _reorder_canonical_first and _collect_canonical_tied_patterns, the call site, the import, and their tests. Confirmed it is a no-op under the name-based drop (the canonical is kept regardless of iteration order; declared aliases are pre-marked so the address backstop skips them; undeclared shares the reorder never sees), and it was a second parser of _tied_weights_keys — removing it finishes the single-dedup-authority consolidation. DiffusionGemma dedup is still 0-leaked without it (HSG re-run).
| self.tied_cache = None | ||
| self.moe_tied_cache = None | ||
| # Import here to avoid a circular import at module load time. | ||
| from .model_utils import TiedGroupResolver |
There was a problem hiding this comment.
Bot comment.
The stated justification does not hold: model_utils.py imports only re, warnings and torch.nn — nothing from this package — so there is no cycle to break and this can be a module-level import (the file already has a TYPE_CHECKING import of the same name at line 37). quant_utils.py:1667 adds a second function-level import of TiedGroupResolver with no stated reason at all.
CONTRIBUTING.md asks for imports at module top, with a function-level import only for a concrete reason such as an unavoidable cycle. Beyond style, the inaccurate comment is the kind of thing that gets copied, and it defers any import error to export time.
There was a problem hiding this comment.
Fixed: the registry __post_init__ and its inaccurate 'avoid circular import' comment are gone (7d5155a9f removed the resolver from ExportContext). The quant_utils function-level import is hoisted to module level in ad17a5a95 — model_utils imports only re, warnings, and torch.nn, so there is no cycle.
|
|
||
| # Reuse a caller-provided resolver when given (built once per export), else build. | ||
| if self.resolver is None: | ||
| self.resolver = TiedGroupResolver(self.model) |
There was a problem hiding this comment.
Bot comment.
Separate from the "nobody reads ctx.resolver" thread above: the offload/streaming path — the motivating case for moving off data_ptr() — is not actually covered by this change, so the two paths now produce different key sets for the same model.
_export_transformers_checkpoint_streaming builds an ExportContext (paying a full named_parameters(remove_duplicate=False) walk to construct a resolver it never reads), never calls postprocess_state_dict, and keeps its own tie-drop logic:
raw_tied_keys = set(getattr(model, "_tied_weights_keys", None) or []) if config.tie_word_embeddings else set()On transformers>=5 the dict form makes set(...) the pattern keys (e.g. lm_head\.weight), which never equal a state-dict key, so the drop silently no-ops and the offload export emits both sides of every declared tie while the resident export emits one. Either thread the resolver into the streaming key filter, or say explicitly in the docstring that the streaming path keeps its own (tie_word_embeddings-gated, exact-name) rule — but the current state builds the resolver there for nothing.
While in here: has_non_resident_weights in modelopt/torch/quantization/utils/core_utils.py:667 now has zero callers repo-wide (this PR removed its only import) and can go.
There was a problem hiding this comment.
Partly fixed, partly documented. ExportContext no longer builds a resolver (7d5155a9f), so the streaming path no longer pays a named_parameters walk for one it never read. The streaming path does keep its own tie_word_embeddings-gated, exact-name drop and does not dedup dict-style / MoE ties — this is now stated explicitly in the streaming exporter's docstring (ad17a5a95), with resolver-threading left as a tracked follow-up (so such models should use the resident path meanwhile). The now-dead has_non_resident_weights is removed in ad17a5a95.
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: 2
🧹 Nitpick comments (4)
modelopt/torch/export/quant_utils.py (2)
1177-1180: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winUse
logger.infofor the expected deduplication outcome.Line 1177 logs at warning level for the normal success path. Every declared tie produces a warning on every export. A per-expert MoE declaration produces one warning per group. The three log calls in this block use the same level, so a reader cannot separate the two skip conditions (which are anomalies) from the successful drop (which is the intended behavior).
Lower the success message to
logger.infoand keeplogger.warningfor the two skip branches at Lines 1161 and 1169.♻️ Proposed change
for k, _ in members: keys_to_delete.append(k) - logger.warning( + logger.info( f"Tied weight (declared): dropping {len(members)} alias key(s) under " f"'{a_base}'; canonical '{alias_prefixes[a_base]}' kept." )🤖 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/quant_utils.py` around lines 1177 - 1180, Change the tied-weight deduplication log in the successful drop path to use logger.info instead of logger.warning. Keep logger.warning unchanged for the two skip branches near the same block, so anomaly conditions remain warnings while expected alias removal is informational.
1203-1211: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winMove the
safetensorsimport to module scope.
safetensorsis a direct dependency, andstorage_ptrandstorage_sizeare publicsafetensors.torchutilities. The dependency is not version-pinned.🤖 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/quant_utils.py` around lines 1203 - 1211, Move the safetensors.torch import for storage_ptr and storage_size from the local scope to module scope in quant_utils.py, while preserving the existing tensor deduplication logic around already_marked, seen_tensors, and post_state_dict.Source: Coding guidelines
modelopt/torch/export/unified_export_hf.py (1)
956-961: 🚀 Performance & Scalability | 🔵 Trivial | 💤 Low valueConsider building the resolver after the offload rejection.
TiedGroupResolver(model)walksnamed_parameters(remove_duplicate=False)and runs a regex match per declared pattern per parameter. Line 969 then raisesNotImplementedErrorfor offloaded models. For that case the whole alias-map build is discarded. Moving the construction below thehas_accelerate_offload(model)check avoids the work on the error path.Note that the resolver must still be built before
_process_quantized_modules, because the object-identity gate requires unpacked Parameters.🤖 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_hf.py` around lines 956 - 961, Move TiedGroupResolver(model) construction below the has_accelerate_offload(model) rejection so offloaded models fail before building the unused alias map. Keep resolver creation before _process_quantized_modules to preserve the object-identity gate’s access to unpacked Parameters, and retain its existing use by amax synchronization and postprocess_state_dict deduplication.modelopt/torch/export/model_utils.py (1)
400-427: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueConsider one helper that returns both the alias prefix and the canonical key.
matched_alias_prefixrepeats the prefix walk ofcanonical_state_dict_key. The two walks also apply different accept conditions:canonical_state_dict_keystops at the longest matching prefix and returnsNonewhen the rewrite equals the key, whilematched_alias_prefixcontinues to shorter prefixes in that case. The caller inmodelopt/torch/export/quant_utils.pythen records a member withcanonical_key=None, which keeps the whole group. That fallback is safe, so this is not a defect.A single method returning
(alias_prefix, canonical_key)removes the duplication and the divergence.🤖 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/model_utils.py` around lines 400 - 427, Consolidate canonical_state_dict_key and matched_alias_prefix into one helper that performs the prefix walk once and returns both the matched alias prefix and canonical key, using consistent longest-prefix and unchanged-rewrite handling. Update callers to consume the tuple and preserve the existing safe None fallback behavior when no rewrite applies.
🤖 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/quant_utils.py`:
- Around line 1697-1704: Restrict or remove the object-identity-based
dense_shared fallback in sync_tied_input_amax, since _process_quantized_modules
replaces undeclared shared weights and postprocess_state_dict retains both
entries. Only merge amax values for weights that remain aliased after
processing, or add matching deduplication for packed undeclared shares; update
the related test and comments to reflect the preserved separate input scales.
In `@modelopt/torch/export/unified_export_hf_streaming.py`:
- Around line 206-213: Correct the streaming export docstring near the
tied-weight handling to state that only list-style _tied_weights_keys
declarations are supported and dict-style declarations are not deduplicated,
including the common embedding-tie case. Keep the implementation unchanged
unless explicitly choosing to add regex-pattern handling for dict declarations.
---
Nitpick comments:
In `@modelopt/torch/export/model_utils.py`:
- Around line 400-427: Consolidate canonical_state_dict_key and
matched_alias_prefix into one helper that performs the prefix walk once and
returns both the matched alias prefix and canonical key, using consistent
longest-prefix and unchanged-rewrite handling. Update callers to consume the
tuple and preserve the existing safe None fallback behavior when no rewrite
applies.
In `@modelopt/torch/export/quant_utils.py`:
- Around line 1177-1180: Change the tied-weight deduplication log in the
successful drop path to use logger.info instead of logger.warning. Keep
logger.warning unchanged for the two skip branches near the same block, so
anomaly conditions remain warnings while expected alias removal is
informational.
- Around line 1203-1211: Move the safetensors.torch import for storage_ptr and
storage_size from the local scope to module scope in quant_utils.py, while
preserving the existing tensor deduplication logic around already_marked,
seen_tensors, and post_state_dict.
In `@modelopt/torch/export/unified_export_hf.py`:
- Around line 956-961: Move TiedGroupResolver(model) construction below the
has_accelerate_offload(model) rejection so offloaded models fail before building
the unused alias map. Keep resolver creation before _process_quantized_modules
to preserve the object-identity gate’s access to unpacked Parameters, and retain
its existing use by amax synchronization and postprocess_state_dict
deduplication.
🪄 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: adf93f39-4626-4a32-9eea-70b31f66aa18
📒 Files selected for processing (9)
modelopt/torch/export/model_utils.pymodelopt/torch/export/quant_utils.pymodelopt/torch/export/registry.pymodelopt/torch/export/unified_export_hf.pymodelopt/torch/export/unified_export_hf_streaming.pymodelopt/torch/quantization/utils/core_utils.pytests/_test_utils/torch/quantization/tied_modules.pytests/unit/torch/export/test_export_registry.pytests/unit/torch/export/test_unified_export_hf.py
💤 Files with no reviewable changes (1)
- modelopt/torch/quantization/utils/core_utils.py
🚧 Files skipped from review as they are similar to previous changes (2)
- tests/unit/torch/export/test_export_registry.py
- tests/unit/torch/export/test_unified_export_hf.py
|
Directionally the change is right to me. Before we had a Now, we build a tied weights map based on name, and dedupe with that in the This fixes a false-positive issue in dedupe in the pointer becomes invalid during the export flow (quant weight packing, moe expansion, etc.) The logic in the PR is becoming a bit hard to understand with consecutive agent patches. It seems the agent is very worried about inconsistencies between name-based tie config and the actual weights. I'd propose to have an agent re-write the PR from scratch, using |
|
Here is a re-write from agent based on my suggested prompt above: https://github.com/NVIDIA/Model-Optimizer/pull/2151/changes |
…ect)
Replace the data_ptr-based tied-weight dedup in the unified HF export with a
name-based scheme driven by the model's own _tied_weights_keys /
tie_word_embeddings declarations. Address identity misfires in several ways: a
freed address recycled by the allocator can falsely alias two unrelated
weights, and the FSDP full-state-dict gather (and offload) materializes tied
weights at distinct addresses so a genuine tie is missed and both copies are
written. Names are stable across packing, FSDP resharding, and offload -- this
implements the TODO already noted on ExportContext.
- Add TiedGroupResolver (model_utils): builds {alias -> canonical} from
dict-style _tied_weights_keys (with per-layer regex backreferences and the
container-level fused-experts tie) plus tie_word_embeddings. Enumerates
named_parameters(remove_duplicate=False) so a genuinely shared Parameter is
seen under both names even when the canonical side is registered first.
- postprocess_state_dict is the authoritative dedup: drop each declared alias
key whose canonical is present (address-independent -> correct under the FSDP
gather / offload). A (device, data_ptr, size) pass is kept as a backstop for
undeclared/coincidental shares; deletion is idempotent.
- One resolver is built per export and threaded through sync_tied_input_amax
(name-based grouping), the ExportContext MoE cache, and postprocess.
- sync_tied_input_amax still runs before packing: the tied group collapses to
one retained weight whose single input_scale must cover every side's range.
- Remove the per-module dense tied cache: both sides pack identically (sync
equalizes scales) and the duplicate is dropped by name; the fused-MoE cache
is retained as a resident-path compute/memory optimization, keyed by the
name-based container group key and disabled under has_non_resident_weights
(FSDP2 / offload), as before.
On-disk output is unchanged for every declared tie and untied model; only the
memory-identity heuristic is replaced. The streaming offload path already
deduped by _tied_weights_keys and is unaffected.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Signed-off-by: Juhi Mittal <juhim@nvidia.com>
…mma)
DiffusionGemma declares ties as {alias_regex: canonical_regex} where the value
is a second regex structurally identical to the alias except for a leading
literal head (e.g. "encoder.language_model.layers\.(?:[^.]+\.)*gate_up_proj" ->
"decoder.layers\.(?:[^.]+\.)*gate_up_proj"), NOT a re.sub backreference template.
_build_tied_alias_map treated the value as a re.sub replacement, emitting the raw
pattern string as the "canonical" name; that name never exists in the state dict,
so postprocess dropped nothing and the encoder experts were written to disk
(under-dedup: ~1536 extra keys on a DiffusionGemma nvfp4_experts_only export).
Add _canonical_via_pattern_pair: when the alias and canonical declarations are
parallel patterns, derive the differing literal head via longest-common-suffix
and swap it, copying the shared trailing structure from the concrete name. The
loop tries this first and falls back to re.sub for genuine backreference
templates (and plain-name canonicals). Regression test covers the exact
DiffusionGemma format, including the post-export per-expert split key rewriting
to the decoder canonical.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Signed-off-by: Juhi Mittal <juhim@nvidia.com>
…ority) Prefer fewer moving parts and one dedup authority over an in-memory optimization. The moe_tied_cache aliased a tied experts container's already-packed per-expert buffers to skip re-packing the second side. It was never load-bearing for correctness -- postprocess_state_dict's name-based drop yields the exact same on-disk checkpoint with or without it -- and it relied on _alias_per_expert_subtree_from_prior, a delicate hand-rolled routine that rebuilds each expert's weight / weight_scale / weight_scale_2 / input_scale aliases and could silently mis-alias a buffer. Removing it makes both dense and fused-MoE tied weights follow one auditable path: pack each side independently to byte-identical tensors, then drop the duplicate keys by name in postprocess. Fewer variables, a smaller surface for silent corruption, and no residency guard to reason about. On-disk output is unchanged: postprocess drops the same declared-alias keys, so the exported checkpoint is byte-identical (verified on DiffusionGemma -- same 47067 keys and total_size as the with-cache run). Accepted tradeoff (does not affect the stored weights or loading): without the in-memory aliasing a tied experts container is re-packed rather than shared, so during save its experts exist as separate tensors until postprocess drops the keys. This raises peak save memory for tied-MoE and inflates the informational `total_parameters` index field (tied experts counted per-side -- e.g. 25.8B vs 14.4B on DiffusionGemma-26B). The simpler single-authority path is worth that cost. - moe_utils: drop _moe_tied_cache/_tied_group_key params, the fast-path alias+return, the cache register, and _alias_per_expert_subtree_from_prior. - ExportContext: drop the moe_tied_cache field and the has_non_resident_weights guard (nothing left to disable; the name-based postprocess drop is FSDP/offload-safe). container_group_key stays on the resolver -- sync_tied_input_amax still groups amax by it. - Tests: tied fused experts now assert independent storage + byte-identical values (dropped later by name); offload/registry cache tests removed or repointed. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Juhi Mittal <juhim@nvidia.com>
…ionable warning Clarify (comment) that the (device,data_ptr,size) backstop only ever fires for unquantized/unpacked shared weights, which keep the single original shared Parameter and thus two keys on one storage. Quantized tied weights cannot reach it: each side packs into its own fresh Parameter (distinct storage, byte-identical), so they are collapsed by the name-based pass, which is the sole authority for quantized ties. The backstop remains because safetensors save_file raises on any two keys sharing storage, so a residual undeclared share must be collapsed here or the export fails at write time. Make the warning actionable: a quantized weight reaching the backstop means its tie was not declared in _tied_weights_keys / tie_word_embeddings and was missed by the name-based dedup. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Juhi Mittal <juhim@nvidia.com>
The changelog entry was written in the first commit, before the per-module caches were removed. Update it to the shipped design: one TiedGroupResolver drives the input-amax sync and the postprocess name-drop; both the dense and fused-MoE data_ptr caches are removed; the (device,data_ptr,size) pass is only a backstop for undeclared unquantized same-storage shares. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Juhi Mittal <juhim@nvidia.com>
…collision, cleanups) - postprocess_state_dict: guard the name-drop against a (pathological) bidirectional alias map (A<->B) that would otherwise mark both sides and leave the loader with a missing tensor. Only drop an alias whose canonical is a terminal canonical (not itself an alias); warn and keep both otherwise. Regression test added. - TiedGroupResolver.alias_prefix_pairs: warn (instead of silently overwriting) when one alias module prefix maps to two different canonical prefixes, so declared per-key rewrites are not misrouted. - Build the name-based resolver once per export and thread it into _prepare_moe_inputs (was built twice: prepare context + driver). - Fix stale comments still referencing the removed per-module MoE dedup cache (_process_quantized_modules and the driver). - Narrow test_tied_weights_exported_independently_without_cache docstring: it checks packing behavior, not the offload/streaming path. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Juhi Mittal <juhim@nvidia.com>
…drop, id sync) Follow-up to the name-based tied-weight dedup, addressing cjluo-nv review: - _build_tied_alias_map now confirms a declared dict-style tie is actually applied (alias and canonical resolve to the SAME live Parameter) before trusting it. A class-level _tied_weights_keys dict is declared regardless of config (e.g. lm_head<->embed_tokens with tie_word_embeddings=False), so a name-only match would drop an independent weight. Object identity is safe here: the map is built pre-packing while all params are resident, so no freed-then-reused address can forge a match (unlike data_ptr). - postprocess_state_dict drops an alias's keys atomically per module prefix: only when EVERY key has a canonical counterpart present, else keep all. Fixes orphaned weight_scale/input_scale when tied sides have mismatched quant state, and stops an untied sibling under an alias prefix from being dropped. - sync_tied_input_amax gains an id(weight) fallback so undeclared physical shares (which the address backstop still collapses) get their input amaxes max-merged too, avoiding silent activation clipping on the surviving side. - _canonical_via_pattern_pair validates its rewrite with re.fullmatch and falls back to re.sub when the common-suffix split lands mid-token. - Drop the dead ExportContext.resolver field: no handler read it; the driver owns the one resolver and feeds it to sync_tied_input_amax / postprocess_state_dict. Removes an avoidable per-context alias-map build. - Type-annotate resolver params via TYPE_CHECKING import. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Juhi Mittal <juhim@nvidia.com>
…ode cleanup) - postprocess_state_dict address backstop keys on safetensors' own shared-storage identity (device, storage_ptr, storage_size) instead of tensor extent. A base tensor and a shorter view share storage and safetensors save_file rejects them, but the numel*element_size key split them into two ids so both survived and the export crashed at write time. Two distinct live tensors cannot share a data_ptr without sharing storage, so storage identity never false-collapses independents. - TiedGroupResolver.container_group_key returns None when the resolved canonical does not end in the projection suffix (removesuffix would otherwise be a silent no-op, mis-grouping tied containers and skipping the amax merge). - Remove _reorder_canonical_first + _collect_canonical_tied_patterns and their call site: a no-op under the name-based drop (which keeps the canonical regardless of iteration order) and a second parser of _tied_weights_keys that could drift from _build_tied_alias_map. Finishes the single-dedup-authority consolidation. - Hoist TiedGroupResolver import in quant_utils to module level (model_utils has no local imports, so no cycle) and drop the function-level import. - Remove has_non_resident_weights (zero callers after the cache removal). - Document that the streaming exporter keeps its own tie_word_embeddings-gated, exact-name tie-drop and does not dedup dict-style / MoE ties (follow-up tracked). 112 unit tests pass. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Juhi Mittal <juhim@nvidia.com>
Simplify the tied-weight path (inspired by the id-grouping in PR NVIDIA#2151) while keeping both safety properties the identity-only approach lacks: declared-only drop and the shared-storage backstop. - _build_tied_alias_map now detects ties by shared object identity: group params by id(parameter) pre-pack (resident), and use _tied_weights_keys / tie_word_embeddings ONLY to label which member of a shared group is canonical. The canonical-side regex is never parsed. Delete _canonical_via_pattern_pair. id is observed once at build time and recorded as NAMES; names survive packing / FSDP gather / offload, so the drop (in postprocess) never needs the packed tensors to still be the same object -- which they aren't. - Several guards become impossible and are removed: the declared-but-unapplied is-gate and the re.fullmatch check (with the regex derivation), the container_group_key removesuffix fail-safe, and the postprocess bidirectional guard (chains can't form when ties are id-groups). - Keep declared-only drop (undeclared shares are not ours to drop) and the storage-identity backstop; postprocess / sync logic otherwise unchanged. - Rename TiedGroupResolver -> TiedWeightMap: a thin immutable view over the id-derived {alias: canonical} map, no longer a regex resolver. - Condense the docstrings/comments added across the review rounds. Net ~120 fewer lines. 112 unit tests pass; DiffusionGemma (47067 tensors, 0 leaked encoder-expert keys) and MiniMax (191211 tensors, 15872/15872 experts) exports verified unchanged on HSG. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Juhi Mittal <juhim@nvidia.com>
2a56654 to
77c2874
Compare
juhi10071998
changed the title
juhi10071998
changed the title
Worked example: what
|
The parameter and driver variable were still named `resolver`, but the type is `TiedWeightMap` (no longer a regex resolver). Rename to `tied_map` in postprocess_state_dict / sync_tied_input_amax, the driver, and the tests, plus two stale comments. No behavior change. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Juhi Mittal <juhim@nvidia.com>
cjluo-nv
left a comment
cjluo-nv
left a comment
There was a problem hiding this comment.
Bot review (bedrock-claude-opus-5) — DM the bot to share feedback.
Re-review of the tied-weight dedup rework (14 files, +668/-457 ≈ 1125 lines; design-review gate fired).
Design review. Problem: HF-export dedup keyed on data_ptr(), which is read after packing/gather and so both false-positives (allocator reuse → dropped an independent MoE expert, NVBug 6525352) and false-negatives (FSDP gather → tie missed). Alternatives in-repo/deps: (a) keep the data_ptr caches and re-key them on a group name; (b) id-only grouping as in PR #2151; (c) the streaming exporter's existing exact-name _tied_weights_keys filter; (d) let save_pretrained/safetensors' own shared-storage handling do it. The PR body now explicitly compares against #2151 (id-only) and against the old cache approach, and explains why declaration + identity are combined (physical sameness vs. re-tiability) — the design question is addressed, so I moved on to correctness.
Previously-raised comments — status. Addressed: declared-but-unapplied tie (now gated by id(parameter) grouping, with test_build_tied_alias_map_skips_declared_but_unapplied_tie); orphaned weight_scale/input_scale on mixed quant state (atomic per-prefix drop + replacement test); ExportContext.resolver dead field removed and one map built in the driver; storage-identity backstop now matches safetensors' (device, storage_ptr, storage_size) grouping with a view/base test; _reorder_canonical_first/_collect_canonical_tied_patterns and has_non_resident_weights deleted; resolver→tied_map typed; CHANGELOG and stale MoE-cache/offload docstrings corrected; bidirectional declaration now yields an empty map with a regression test.
Still unresolved (see inline): the prefix-wide rewrite can still drop an untied sibling under an alias module prefix whenever a same-named key exists under the canonical prefix (partial-container ties, or a tied weight with a distinct sibling parameter). The atomic-completeness rule only saves the case where the counterpart is absent, so the PR-body claim "an untied sibling under an alias prefix is never dropped" is not accurate, and there is no regression test for it. Given the bug being fixed is exactly "dedup dropped an independent weight", a cheap shape/dtype (or byte) check before dropping, or restricting the rewrite to keys derived from the declared parameter, would close it.
Smaller items inline: the newly added function-local safetensors.torch import; the FSDP2 identity assumption now underpinning detection (no wrapped-model test); container_group_key's removesuffix silently no-ops if the canonical projection is registered under a different attribute name, which silently skips the amax merge (the earlier None guard was dropped in the simplification).
No licensing changes and no prompt-injection content observed. Size is on the large side but cohesive.
Additional comments (outside the PR diff):
modelopt/torch/export/model_utils.py:420— > Bot comment.
The docstring here ("State-dict dedup rewrites any exported key under an alias prefix") is an accurate description of the implementation, but it is exactly the behaviour that widens a per-parameter declaration into a whole-subtree substitution — see the comment on postprocess_state_dict. If you keep the prefix approach, please note in the docstring that correctness relies on every exported key under an alias prefix genuinely belonging to the tie, and reconcile the PR-body claim that "an untied sibling under an alias prefix is never dropped".
|
|
||
| # No bidirectional/chain guard is needed: ties are id-groups, so if A<->B and | ||
| # B<->C were both shared, A, B and C would be one group -- a canonical is never | ||
| # also an alias of another group, so alias->canonical->alias chains cannot form. |
There was a problem hiding this comment.
Bot comment.
💬 Author replied that the atomic per-prefix drop makes untied siblings safe — that only holds when the sibling has no canonical counterpart. When it does have one, it is still dropped.
Concretely: a tie declared/applied on encoder.experts.gate_up_proj only (with down_proj an independent Parameter) yields alias_prefixes = {"encoder.experts": "decoder.experts"}. Every encoder.experts.<e>.down_proj.* key then matches matched_alias_prefix, its decoder.experts.<e>.down_proj.* twin is present, so the completeness check passes and the whole group — including the independent down_proj tensors — is deleted. Same shape for a declared weight tie between two Linears with distinct biases: A.bias is dropped because B.bias exists.
That is the same failure class as NVBug 6525352 (an independent weight silently dropped), just reached by name instead of by address. Two cheap options:
- restrict the rewrite to keys derived from the declared parameter (packed weight + the known
weight_scale/weight_scale_2/input_scalesuffixes + the per-expert split pattern), or - keep the prefix rewrite but require
shape/dtype(ideallytorch.equal) agreement between alias key and canonical key before appending tokeys_to_delete, warning otherwise.
Either way please add a regression test: an alias prefix containing an untied sibling whose canonical twin is present, asserting the sibling survives.
There was a problem hiding this comment.
Addressed in 8568a86ea — thanks, this was a real one. I reproduced it first (an independent A.bias was being dropped just because B.bias existed), then took your option 1.
The dedup no longer rewrites by module prefix. Each tie now expands to only the tied parameter's own exported keys:
- Dense
X.weighttie →weight+weight_scale/weight_scale_2/input_scale. A siblingbiasisn't in that set, so it can't be swept. - Fused-MoE → only the per-expert keys of the tied projections, matched by projection name (
gate_proj/up_proj/down_proj). A partially-tied container dedups exactly what's tied, and an untieddown_projor a router under the same container survives.
Removed the now-unused alias_prefix_pairs / matched_alias_prefix / canonical_state_dict_key. Added regression tests for both cases (tied weight + independent bias; partially-tied container + router). Re-verified on HSG: DiffusionGemma 47067 tensors / 0 leaked, MiniMax 191211 / 15872 experts / 0 missing.
| f"tie_word_embeddings and was not caught by the name-based dedup." | ||
| ) | ||
| else: | ||
| seen_tensors[tensor_id] = key |
There was a problem hiding this comment.
Bot comment.
from safetensors.torch import storage_ptr, storage_size is a function-local import of a hard dependency (safetensors is already imported at module scope elsewhere in the export package) and there is no circular-import or heavy/optional-dependency reason given. Please move it to the top of the file per the repo convention; deferring it only delays an import error to export time.
There was a problem hiding this comment.
Addressed in 8568a86ea — moved from safetensors.torch import storage_ptr, storage_size to the top of the file.
| # resolves ties from the model's declarations (stable across FSDP resharding / offload | ||
| # / packing). | ||
| tied_map = TiedWeightMap(model) | ||
| _prepare_moe_inputs(model, dtype, is_modelopt_qlora) |
There was a problem hiding this comment.
Bot comment.
Detection now rests on id(parameter) observed here, i.e. after FSDP2 wrapping. If fully_shard hands the two tied modules distinct sharded parameter objects (e.g. when they land in different FSDP parameter groups), the group never forms, no alias is recorded, and both sides are written — the exact false-negative this PR sets out to fix, and the name-based map can no longer recover it because it was never populated.
The old name-only detection was immune to this. The PR body claims 4-GPU verification; could you state which of those runs actually had applied ties under FSDP2 (MiniMax has none), and add a CPU/2-proc test that builds the map on a wrapped model (or at least an assertion/warning when a declared alias matches a name but no id-group formed) so a future torch change can't silently disable dedup?
There was a problem hiding this comment.
Addressed in 8568a86ea, and thanks — this is the real trade-off of detecting ties by id().
To be straight about what the runs covered: MiniMax has no applied ties, and the DiffusionGemma run used accelerate device_map (resident), not fully_shard — so the specific 'fully_shard hands back distinct objects' case is genuinely untested. (The same applies to the id-based detection in #2151.)
What this does is make it fail loud instead of silent: when a declared tie doesn't form a shared id-group and the model is under FSDP2/offload, _build_tied_alias_map now warns that dedup may be incomplete and suggests gathering/unsharding before export. It's scoped to that context so it doesn't cry wolf on the normal declared-but-unapplied case off FSDP. There's a test for both directions (silent off FSDP, warns under it).
… warning
Address the latest cjluo-nv review.
- (bug) postprocess no longer rewrites by module PREFIX, which swept up untied
siblings sharing the prefix -- an independent bias next to a tied weight, or an
untied projection of a partially-tied experts container -- when a canonical twin
happened to exist (NVBug 6525352's failure class, reached by name). Now each
{alias: canonical} tie expands to the tied parameter's OWN exported keys: for a
dense .weight tie, weight + weight_scale/weight_scale_2/input_scale; for a fused
MoE container, the per-expert keys only when the container is FULLY tied (a
gate/up projection AND down_proj declared). Atomic drop unchanged. Removes the
now-unused alias_prefix_pairs / matched_alias_prefix / canonical_state_dict_key.
Regression tests: tied weight + independent bias -> bias survives; partially-tied
container -> untied projection survives.
- Hoist `from safetensors.torch import storage_ptr, storage_size` to module top.
- Detection rests on id(parameter) read at export time, which under FSDP2 may be
after fully_shard has split a tied param into distinct sharded objects (id-group
never forms -> silent under-dedup). Add a scoped warning: when a declared alias
did not form a shared id-group AND the model is FSDP2/offloaded, warn. Test both
directions (quiet off FSDP, warns under it).
115 unit tests pass.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Signed-off-by: Juhi Mittal <juhim@nvidia.com>
cjluo-nv
left a comment
cjluo-nv
left a comment
There was a problem hiding this comment.
Bot review (bedrock-claude-opus-5) — DM the bot to share feedback.
Re-review of the tied-weight dedup rework (14 files, +782/-456 ≈ 1238 lines; design-review gate fired). Design question is settled: the PR body compares against #2151 (id-only), the old data_ptr caches, the streaming exporter's exact-name filter and safetensors' own shared-storage handling, and explains why declaration + identity are combined. No licensing changes and no prompt-injection content observed.
Previously-raised comments now addressed. The big one — prefix-wide rewrite dropping an untied sibling — is genuinely fixed: alias_prefix_pairs/matched_alias_prefix/canonical_state_dict_key are gone and the drop now expands a tied parameter into its own exported keys (weight_suffixes) or the tied projections' per-expert keys, with two new regression tests (independent bias survives; partially-tied container keeps down_proj and the router). Also fixed: safetensors.torch import hoisted to module scope; declared-but-unapplied tie gated by id(); atomic per-prefix drop for mixed quant state; bidirectional declaration yields an empty map; dead ExportContext.resolver / has_non_resident_weights / _reorder_canonical_first removed; CHANGELOG and stale docstrings corrected.
Still unresolved / new (inline):
- Storage-identity backstop can now silently drop a distinct tensor (this replaces the earlier
data_ptr+extent key and supersedes the earlier "use complete view metadata" comment, which was resolved in the opposite direction). Any two exported tensors that merely share a storage — different offsets, shapes, or dtypes — are now collapsed keeping whichever is iterated first.test_postprocess_state_dict_collapses_view_and_base_sharing_storagecodifies dropping the 4-element base and keeping the 2-element view, i.e. losing data. There is a concrete reachable path:_export_fused_experts_modulehas noQUANTIZATION_NONEearly-return, and per-expert slices in_export_fused_expertsare contiguous, so when an expert weight quantizer is disabled_export_quantized_weightearly-returns and every per-expert weight stays a view into the one fused storage. Pre-PR those had distinctdata_ptrs andsave_fileraised loudly; now all but one per storage are silently deleted — exactly the NVBug 6525352 failure class this PR exists to remove. weight_suffixesomitspre_quant_scale, so an AWQ/SVDQuant-style tied dense pair drops the aliasweight/scales but leaves<alias>.pre_quant_scaleorphaned (the atomicity guarantee only covers enumerated suffixes). The MoE branch is fine because it matches by projection path component.- New function-local imports (
has_accelerate_offload,is_fsdp2_model) in_build_tied_alias_mapwith no stated reason;sync_tied_input_amaxstill importsdefaultdictlocally.
Smaller, non-blocking: the FSDP2 identity assumption is still only covered by a warning-path test (author acknowledges no wrapped-model repro); the ("dense_shared", id(weight)) amax merge still applies to quantized undeclared shares that both survive packing, widening both sides' input_scale unnecessarily; a tie declared on a fused projection whose alias/canonical last components differ falls through to a silent no-dedup with no warning, which sits oddly with the "fail loud" philosophy elsewhere in the diff.
| # Use tensor data pointer to identify tied weights | ||
| tensor_id = value.data_ptr() | ||
| if key in already_marked: | ||
| continue |
There was a problem hiding this comment.
Bot comment.
Keying the backstop on (device, storage_ptr, storage_size) matches safetensors' grouping, but not its action: save_file raises on a shared-storage group, whereas this drops every member but the first — even when the members are non-identical views (different storage_offset, shape, stride, or dtype). The new test enshrines that: {"short": base[:2], "long": base} keeps the 2-element view and deletes the 4-element base, i.e. loses data.
This is reachable, not hypothetical: _export_fused_experts_module has no get_quantization_format(...) == QUANTIZATION_NONE early-return (unlike _export_quant_linear), and in _export_fused_experts first_proj[idx, :expert_dim, :] / down[idx] are contiguous, so .contiguous() is a no-op and wrapper.weight aliases the fused storage. If the per-expert weight quantizers are disabled, _export_quantized_weight early-returns and every per-expert weight is a view into one storage → this loop keeps one key and silently deletes the rest. Pre-PR each view had a distinct data_ptr, so they all survived and save_file failed loudly.
Suggest only dropping when the two tensors are genuinely the same logical tensor (equal dtype/shape/stride/storage_offset, or torch.equal), and for a non-equivalent share either clone() the later tensor to break the sharing (what safetensors' own save_model/_remove_duplicate_names effectively does) or raise — silent deletion of a distinct weight is the failure class this PR is fixing. A regression test with two non-overlapping slices of one storage would pin it.
There was a problem hiding this comment.
Addressed in e185b01fb — good catch, this was a real one. Keying on (device, storage_ptr, storage_size) matched how safetensors groups shared storage, but not what it does: I was dropping all-but-first, which loses data when two tensors share a buffer but are different views of it (different offset/shape) — exactly the failure class this PR is fixing.
It now behaves like safetensors' own save_model:
- if a later key is the same view (
data_ptr, shape, stride,storage_offset, dtype all equal) → it's a true duplicate → drop it; - if it's a distinct view of the same buffer →
clone()it to break the share.
So the unquantized fused-experts case (each per-expert first_proj[idx] is a view into one buffer) now clones each expert rather than silently deleting it — no crash and no data loss. Added tests for distinct views (both survive), two non-overlapping slices, and a true duplicate; the old short/long test was updated to assert the non-lossy behavior.
| # `weight_suffixes` / `proj_splits` HERE -- nothing else needs to change. | ||
| # | ||
| # A dense quantized weight exports as these keys; a `X.weight` tie covers exactly them. | ||
| weight_suffixes = ("weight", "weight_scale", "weight_scale_2", "input_scale") |
There was a problem hiding this comment.
Bot comment.
weight_suffixes misses pre_quant_scale. postprocess_state_dict renames <module>.input_quantizer._pre_quant_scale → <module>.pre_quant_scale (see _KV_CACHE_REPLACEMENTS), so for an AWQ / NVFP4_AWQ / SVDQuant tied dense pair the alias's weight + weight_scale* + input_scale are dropped while <alias>.pre_quant_scale stays behind — an orphaned tensor for a weight that no longer exists, which is the same class of problem the atomic-drop rule was added to prevent (the atomicity check only ranges over the enumerated suffixes, so an unenumerated companion is invisible to it).
Adding "pre_quant_scale" to the tuple covers the known case; a cheap extra safety net would be to warn if, after dropping a dense group, any key still remains under the alias module prefix that is not a plain parameter (e.g. bias), so a future new companion key surfaces instead of silently leaking.
There was a problem hiding this comment.
Addressed in e185b01fb — thanks. Added pre_quant_scale to the suffix list, so an AWQ / NVFP4_AWQ / SVDQuant dense tie drops it along with the weight instead of leaving <alias>.pre_quant_scale orphaned.
I also took your safety-net idea: after a dense tie is dropped, if any non-bias key is still sitting under the alias's module prefix, it now warns — "likely an un-enumerated quantizer companion, add its suffix." That way a future new companion shows up as a warning instead of leaking silently. There's a regression test that a tied pre_quant_scale is dropped with the weight.
| unrealized = declared_aliases - set(alias_to_canonical) | ||
| if unrealized: | ||
| from modelopt.torch.quantization.utils.core_utils import has_accelerate_offload | ||
| from modelopt.torch.utils.distributed import is_fsdp2_model |
There was a problem hiding this comment.
Bot comment.
These two imports are function-local with no stated justification. modelopt/torch/export/quant_utils.py already imports modelopt.torch.quantization.model_calib at module scope and unified_export_hf.py imports both has_accelerate_offload and is_fsdp2_model at module scope, so a cycle looks unlikely — please move them to the top of model_utils.py, or add a one-line comment naming the concrete cycle if there is one. (Same for the from collections import defaultdict still inside sync_tied_input_amax in quant_utils.py, which this PR touches.)
There was a problem hiding this comment.
Addressed in e185b01fb — moved has_accelerate_offload / is_fsdp2_model to the top of model_utils.py and defaultdict to the top of quant_utils.py (no cycle, as you noted). One small consequence: since is_fsdp2_model is now imported into model_utils, the FSDP-warning test patches it there rather than at its source module.
…st imports Address the latest cjluo-nv review. - (1149) weight_suffixes missed `pre_quant_scale` (the AWQ/NVFP4_AWQ/SVDQuant companion, renamed from input_quantizer._pre_quant_scale). A dense tie dropped weight + weight_scale* + input_scale but orphaned `<alias>.pre_quant_scale`. Add it to the tuple, and warn after a dense drop if a non-`bias` key still remains under the alias prefix, so a future un-enumerated companion surfaces instead of silently leaking. - (1230, bug) the storage backstop keyed on (device, storage_ptr, storage_size) matched safetensors' grouping but not its action: it DROPPED all-but-first, which loses data for two DISTINCT views of one storage (e.g. unquantized fused-expert slices at different offsets). Now it drops only a genuine duplicate (same data_ptr/shape/stride/offset/dtype) and CLONES a distinct view to break the share -- mirroring safetensors' save_model (no crash, no data loss). Tests: distinct views cloned, non-overlapping slices both kept, true duplicate dropped. - (244) hoist the FSDP2-warning imports in model_utils and the defaultdict import in sync_tied_input_amax to module top (no cycle). These do not change NVFP4 export output (pre_quant_scale absent; quantized experts pack to distinct storage so never hit the view path). 118 unit tests pass. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Juhi Mittal <juhim@nvidia.com>
Why two mechanisms —
|
cjluo-nv
left a comment
cjluo-nv
left a comment
There was a problem hiding this comment.
Bot review (bedrock-claude-opus-5) — DM the bot to share feedback.
Re-review #5 of the tied-weight dedup rework (14 files, +877/-463 ≈ 1340 changed lines; design-review gate fired). Design question stays settled — the PR body compares against #2151 (id-only), the removed data_ptr caches, the streaming exporter's exact-name filter and safetensors' own shared-storage handling. No licensing changes, no prompt-injection content observed.
Previously-flagged items now resolved. (1) The storage backstop no longer silently deletes a distinct tensor: it drops only a true duplicate (equal shape/stride/storage_offset/dtype) and clone()s a distinct view, which makes the reachable unquantized fused-experts path (_export_fused_experts slices alias one storage; .contiguous() is a no-op there) export correctly instead of losing per-expert data, with three new regression tests. (2) pre_quant_scale is now in weight_suffixes, plus a leftover-companion warning and a test. (3) has_accelerate_offload, is_fsdp2_model, defaultdict and safetensors.torch imports are all hoisted to module scope. Earlier rounds' fixes (per-parameter expansion instead of prefix rewrite, atomic drop, id() gate for declared-but-unapplied ties, dead ExportContext.resolver/has_non_resident_weights/_reorder_canonical_first removal, CHANGELOG) still hold; has_non_resident_weights has no remaining callers.
New finding (inline, worth fixing before merge). _build_tied_alias_map maps every non-canonical member of a shared id-group to the canonical, not just the declared ones. In a ≥3-member group that contains at least one declared alias, an undeclared co-share is also recorded and then dropped by name — with nothing in _tied_weights_keys for the loader to re-tie from, so the checkpoint is missing that key. Once weights are packed the two are distinct storage, so the backstop won't have dropped it pre-PR; this contradicts the module docstring ("an undeclared share … is not dropped") and the design note ("names alone prove re-tiability"). Restricting the alias set to declared_aliases (plus the tie_word_embeddings output-embedding name) is a one-line fix and is safe, because a genuinely storage-sharing undeclared member is still handled by the backstop. No test covers a 3-member group.
Smaller items (non-blocking): alias_groups is keyed by a_pre alone in the MoE branch, so one alias container tied to two different canonical containers silently loses one group (inline); a tie declared on fused projections whose alias/canonical last components differ still falls through to elif alias in post_state_dict and silently no-ops with no warning; the dropped_dense_prefixes leftover scan is O(#dense ties × #keys) (fine at 1 tie, worth noting if many dense ties appear); clone() preserves strides, so a hypothetical non-contiguous shared view would still be rejected by save_file — clone(memory_format=torch.contiguous_format) would close that; the FSDP2 identity assumption is still only covered by a monkeypatched warning-path test (author acknowledged no wrapped-model repro), and the streaming path's divergent tie rule remains documented-but-unfixed.
| continue | ||
| canonical = non_aliases[0] | ||
| for name in names: | ||
| if name != canonical: |
There was a problem hiding this comment.
Bot comment.
This drops every non-canonical member of the id-group, not just the declared ones. If a group has ≥3 names and at least one of them matched a declared alias pattern (so non_aliases is non-empty and the continue above doesn't fire), an undeclared co-share also gets alias_to_canonical[name] = canonical and is then deleted by name in postprocess_state_dict — but there is no declaration for the loader to re-tie it from, so the checkpoint is missing that tensor.
This is not covered by the storage backstop: after packing, the two sides have distinct storage, so pre-PR that undeclared name survived. It also contradicts this function's own docstring ("an undeclared share … has no canonical the loader could re-tie from, so it is not dropped") and the design note in the PR body ("names alone prove re-tiability").
Suggest restricting the emitted aliases to names the model actually declared:
for name in names:
if name != canonical and (
name in declared_aliases or embedding_canonical.get(parameter_id) == canonical
):
alias_to_canonical[name] = canonical(for the tie_word_embeddings branch the output-embedding name is the declared side, so gate on that rather than on declared_aliases). A genuinely storage-sharing undeclared member is still collapsed by the backstop, so nothing regresses at save time. Please add a regression test for a 3-member group {declared alias, canonical, undeclared share} asserting the undeclared name survives.
| and any(part in tied_proj_names for part in key[len(prefix) :].split(".")) | ||
| ] | ||
| if members: | ||
| alias_groups[a_pre] = members |
There was a problem hiding this comment.
Bot comment.
alias_groups is keyed by a_pre only, while moe_containers is keyed by (a_pre, c_pre). If the same alias container has two projections tied to different canonical containers, the second iteration overwrites the first group and those alias keys are never dropped (silent partial dedup, duplicate keys left in the checkpoint). The earlier revision at least warned on that collision before alias_prefix_pairs was removed.
Cheapest fix is to key the group by the pair, e.g. alias_groups[f"{a_pre}->{c_pre}"] = members (group keys are only used for the log message and the dense_prefixes lookup), or warn when a_pre is already present.
Per human review, revert the address backstop to the original value.data_ptr() drop-on-collision pass rather than the storage-key/clone variant. It now runs after the name-based drop is applied, on the reduced dict, so it is byte-for-byte the pre-existing postprocess dedup -- declared ties are handled by name and never reach it. 6525352 is fixed by the tied_cache removal + name drop, not by this pass. Keeps pre_quant_scale + the leftover-companion warning (name pass). Tests updated to the raw-data_ptr behavior. No change for fully quantized exports (dormant). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Juhi Mittal <juhim@nvidia.com>
…t_hf_checkpoint(tied_map=...) Add a public build_tied_weight_map(model) that snapshots the tied-weight map while the model is resident, and a tied_map kwarg on export_hf_checkpoint that consumes it. FSDP2 shard / accelerate offload split the shared tied parameter's id-group, so a map built at export entry misses the tie; capturing pre-shard records it by name, which survives. Default (tied_map=None) is unchanged. Wire it in examples/hf_ptq/hf_ptq.py: capture on the resident model in load_model and pass it to export_hf_checkpoint. Tests: FSDP2 (GPU) and offload (CPU) prove the id-group is lost post-wrap and the pre-capture map still resolves the tie. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Juhi Mittal <juhim@nvidia.com>
|
Added optional load-time tie capture so the dedup is robust under FSDP2/offload. Why: the tie map groups params by What:
Tests: FSDP2 (GPU) and offload (CPU) prove the id-group is lost post-wrap and the pre-captured map still resolves the tie. Verified: DiffusionGemma (encoder+expert keys leaked: 0, 60 tied groups deduped) and MiniMax-M2 (15872 experts, 0 missing) exported through the Note: this is the caller-supplied (opt-in) mechanism. An alternative is auto-capturing inside |
Merge the two negative tied-alias-map cases (list-style + unapplied-tie) into one, merge the two offload characterization tests into the Option B contract test, and drop test_postprocess_backstop_drops_true_duplicate (a strict subset of ..._collapses_keys_sharing_a_dataptr). No coverage lost. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Juhi Mittal <juhim@nvidia.com>
4 participants
What does this PR do?
Type of change: Bug fix / robustness
Fixes NVBug 6525352 — MiniMax-M2.7
nvfp4_mlp_only-kv_fp8failed at TensorRT-LLM load withassert w1_weight is not None and w3_weight is not None, because the previousdata_ptr()-only postprocess dedup could falselydrop an independent MoE expert weight.
Rework tied-weight dedup during unified HF checkpoint export so it no longer depends on tensor
addresses. Ties are detected by object identity before packing, recorded as names, and the
duplicate is dropped by name on the final state dict.
Why the old
data_ptr()approach was wrong — it read the tie signal at the wrong time.The new flow — observe identity once, at the only moment it is true, and carry it forward as a name.
Design in words.
_build_tied_alias_mapgroups parameters byid(parameter)while the model is resident (before packing replaces shared Parameters withindependent tensors). A declared-but-unapplied tie (e.g.
lm_head↔embed_tokenswithtie_word_embeddings=False) is two distinct objects, so it never groups and is never dropped._tied_weights_keys/tie_word_embeddingschoose which member of a shared group to keep; they never create a tie. Only the alias
pattern is matched — the canonical-side regex is never parsed, so parallel-pattern / backref
declarations (DiffusionGemma) need no special handling.
postprocess_state_dict, atomically per module prefix (all of analias's keys, or none, and only when every key has a canonical counterpart) — so tied sides with
different quant state never orphan a
weight_scale/input_scale, and an untied sibling underan alias prefix is never dropped.
data_ptrsurvives only as a backstop, keyed on safetensors' own shared-storage identity(device, storage_ptr, storage_size). It collapses undeclared same-storage shares thatsave_filewould otherwise reject (e.g. a base tensor and a view of it), and only fires forunquantized/unpacked weights — quantized ties pack to distinct storage and are handled by name.
sync_tied_input_amaxmerges input amaxes so the kept side's singleinput_scalecovers everyside; it groups by name, plus an
id(weight)fallback for undeclared shares the backstop willcollapse (no silent activation clipping).
What changed
TiedWeightMap(model_utils.py) — a thin, immutable view over the id-derived{alias: canonical}map; the only stateful piece, built once in the driver and used bysync_tied_input_amaxandpostprocess_state_dict.the single dedup authority. The FSDP/offload
has_non_resident_weightsguard (which existed onlyto disable those caches) is removed with them.
Tradeoff (honest): with the caches gone, each side of a tied-MoE is re-packed instead of
aliased, so peak save-time memory is higher and the informational
total_parametersindexcounts tied experts per side (DiffusionGemma-26B reports 25.8B vs 14.4B). On-disk weights are
unchanged and correct — DiffusionGemma export is byte-identical to the known-good baseline.
Untied models and the streaming offload path are unaffected (streaming keeps its own
tie_word_embeddings-gated, exact-name drop).Usage
No API change. Existing export just works — and now dedups tied weights correctly under FSDP /
offload as well as on a single GPU:
Testing
Unit tests (CPU) cover: id-grouped alias map incl. per-layer backreference and DiffusionGemma
parallel-pattern; the FSDP case simulated on CPU (alias dropped by name across distinct
addresses); atomic per-expert MoE subtree drop; keep-both-sides when a canonical is absent, on a
bidirectional declaration, or when tied sides have mismatched quant state; declared-but-unapplied
tie is not dropped; name-based
sync_tied_input_amax+id()fallback for undeclared shares;storage-identity backstop collapsing a view/base share.
tests/unit/torch/export/+plugins/test_fused_experts.pypass.End-to-end verification (HSG):
nvfp4_mlp_only-kv_fp8, 4-GPU): 191211 tensors, 15872/15872 experts, 0 missing —the 6525352
w1_weight is Noneroot cause is absent.nvfp4_experts_only): tied encoder↔decoder experts +lm_head↔embeddingscollapse to the canonical side — 47067 tensors, 0 leaked encoder-expert keys, byte-identical to
the known-good baseline. Served in vLLM (
vllm/vllm-openai:gemma) and generates correctly(e.g.
17 × 23 = 391with a reasoning trace).Review
Addressed two rounds of maintainer review (
@cjluo-nv,@Edwardf0t1): atomic per-prefix drop formixed quant state,
id()amax merge for undeclared shares, and the storage-identity backstop.The detection core was then simplified to id-grouping (adopting the idea from #2151 while keeping
declared-only drop + the storage backstop), which let the interim guards — the object-identity
gate, the parallel-pattern heuristic and its
re.fullmatchcheck, thecontainer_group_keyfail-safe, and the postprocess bidirectional guard — all be removed as unreachable. Net ~120 fewer
lines than the pre-simplification revision.
Before your PR is "Ready for review"
Additional Information
Related bug: NVBug 6525352 (MiniMax-M2.7 TRT-LLM
load failure — root cause absent after this change). Related PR: #2151 (an id-only variant of the
same idea; this PR adds declared-only drop + the shared-storage backstop it lacks).
Summary by CodeRabbit
Bug Fixes
Testing