Fix/tied weight export identity - #2081
Conversation
📝 WalkthroughWalkthroughExport deduplication now uses shared ChangesExport duplicate-weight handling
Estimated code review effort: 3 (Moderate) | ~25 minutes Sequence Diagram(s)sequenceDiagram
participant ExportContext
participant hf_export_handlers
participant moe_utils
participant StateDict
ExportContext->>hf_export_handlers: Provide canonical names and formats
hf_export_handlers->>ExportContext: Mark duplicate weights skipped
hf_export_handlers->>moe_utils: Export fused experts without tied caches
moe_utils-->>hf_export_handlers: Return expert weights
hf_export_handlers->>StateDict: Remove skipped weights and related buffers
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (5 passed)
✨ Finishing Touches 💡 1⚔️ Resolve merge conflicts 💡
🧪 Generate unit tests (beta)
Comment |
57cf8da to
728828b
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: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
modelopt/torch/export/moe_utils.py (1)
77-81: 🗄️ Data Integrity & Integration | 🟠 Major | 🏗️ Heavy liftKeep source tensors alive in
_moe_tied_cache.
_moe_tied_cachestores rawdata_ptr()values and the packed module. After_delete_fused_moe_source_attrsremoves the source tensors, their storage can be released when_export_fused_expertsreturns. PyTorch can reuse the same addresses for a later untied module. That module can then take the cache-hit path and alias the earlier module’s packed experts.Retain both source tensors in each cache entry. Include device and tensor range metadata in the key. Add a regression for sequential untied exports after source-attribute deletion.
Based on the PR objective, the export must alias only genuine tied expert weights.
Also applies to: 109-115
🤖 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/moe_utils.py` around lines 77 - 81, Update `_export_fused_experts` and the `_moe_tied_cache` lookup so cache entries retain the original source tensors, not just raw `data_ptr()` values and the packed module. Expand the cache key to include device and tensor range metadata, and make the cache-hit path validate those fields before reusing a packed expert. Preserve the tied-weights fast path for genuine shared expert tensors, and add a regression covering sequential untied exports after `_delete_fused_moe_source_attrs` to ensure later modules cannot alias earlier packed experts.
🤖 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 1066-1070: Update the deduplication key in the loop over
post_state_dict within the tied-weight removal logic to include complete view
metadata: device, data pointer, dtype, shape, and stride, rather than only the
logical byte count. Ensure entries are removed only when their tensor views are
exactly equivalent, preserving distinct tensors such as same-start views with
different shapes or strides, and add a regression covering that case.
In `@tests/unit/torch/quantization/plugins/test_fused_experts.py`:
- Around line 696-707: Extend the tied and untied fused-expert tests around
_export_fused_experts to assert aliasing for every contracted buffer: weight,
weight_scale, weight_scale_2, and input_scale. Preserve the existing weight
assertions, and compare each corresponding produced scale buffer between the
relevant expert instances so both tests validate the complete alias contract.
---
Outside diff comments:
In `@modelopt/torch/export/moe_utils.py`:
- Around line 77-81: Update `_export_fused_experts` and the `_moe_tied_cache`
lookup so cache entries retain the original source tensors, not just raw
`data_ptr()` values and the packed module. Expand the cache key to include
device and tensor range metadata, and make the cache-hit path validate those
fields before reusing a packed expert. Preserve the tied-weights fast path for
genuine shared expert tensors, and add a regression covering sequential untied
exports after `_delete_fused_moe_source_attrs` to ensure later modules cannot
alias earlier packed experts.
🪄 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: 9df26287-d88e-40f4-841d-2a1fa3d555fb
📒 Files selected for processing (5)
modelopt/torch/export/hf_export_handlers.pymodelopt/torch/export/moe_utils.pymodelopt/torch/export/quant_utils.pytests/unit/torch/export/test_unified_export_hf.pytests/unit/torch/quantization/plugins/test_fused_experts.py
| # Remove any tied weights if found. Device and size distinguish independent tensors whose | ||
| # allocator addresses happen to match. Zero-pointer tensors are left for serialization to reject. | ||
| for key, value in post_state_dict.items(): | ||
| if isinstance(value, torch.Tensor): | ||
| # Use tensor data pointer to identify tied weights | ||
| tensor_id = value.data_ptr() | ||
| 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.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
Use complete tensor metadata for deduplication.
value.numel() * value.element_size() is only the logical byte count. It does not identify dtype, shape, or stride. It is also not the storage span for a non-contiguous view. Two state-dict entries can therefore share the current key while representing different tensors. The loop then deletes the later key at Line 1072. The preceding squeeze(0) can also change shape without changing the pointer or byte count.
Include dtype, shape, and stride, or require exact view metadata before removing a duplicate key. Add a regression for same-start views with different shape or stride.
Based on the PR objective, deduplication must prevent false-positive removal without dropping distinct exported tensors.
Suggested key
- tensor_id = (value.device, value.data_ptr(), value.numel() * value.element_size())
+ tensor_id = (
+ value.device,
+ value.data_ptr(),
+ value.dtype,
+ tuple(value.shape),
+ tuple(value.stride()),
+ )📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| # Remove any tied weights if found. Device and size distinguish independent tensors whose | |
| # allocator addresses happen to match. Zero-pointer tensors are left for serialization to reject. | |
| for key, value in post_state_dict.items(): | |
| if isinstance(value, torch.Tensor): | |
| # Use tensor data pointer to identify tied weights | |
| tensor_id = value.data_ptr() | |
| if isinstance(value, torch.Tensor) and value.data_ptr() != 0: | |
| tensor_id = (value.device, value.data_ptr(), value.numel() * value.element_size()) | |
| # Remove any tied weights if found. Device and size distinguish independent tensors whose | |
| # allocator addresses happen to match. Zero-pointer tensors are left for serialization to reject. | |
| for key, value in post_state_dict.items(): | |
| if isinstance(value, torch.Tensor) and value.data_ptr() != 0: | |
| tensor_id = ( | |
| value.device, | |
| value.data_ptr(), | |
| value.dtype, | |
| tuple(value.shape), | |
| tuple(value.stride()), | |
| ) |
🤖 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 1066 - 1070, Update the
deduplication key in the loop over post_state_dict within the tied-weight
removal logic to include complete view metadata: device, data pointer, dtype,
shape, and stride, rather than only the logical byte count. Ensure entries are
removed only when their tensor views are exactly equivalent, preserving distinct
tensors such as same-start views with different shapes or strides, and add a
regression covering that case.
| # Module-level dedup preserves genuine fused-expert ties. Transient | ||
| # per-expert wrapper addresses are intentionally not cached. | ||
| moe_tied_cache: dict = {} | ||
| _export_fused_experts( | ||
| parent.encoder.experts, | ||
| torch.float16, | ||
| _moe_tied_cache=moe_tied_cache, | ||
| _tied_cache=tied_cache, | ||
| ) | ||
| _export_fused_experts( | ||
| parent.decoder.experts, | ||
| torch.float16, | ||
| _moe_tied_cache=moe_tied_cache, |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Check every buffer covered by the alias contract.
_alias_per_expert_subtree_from_prior aliases weight, weight_scale, weight_scale_2, and input_scale. The tied test does not check input_scale. The untied test checks only weight. Add assertions for the produced scale buffers in both tests.
As per path instructions, tests must exercise the behavior they claim to validate.
Also applies to: 734-744
🤖 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 696
- 707, Extend the tied and untied fused-expert tests around
_export_fused_experts to assert aliasing for every contracted buffer: weight,
weight_scale, weight_scale_2, and input_scale. Preserve the existing weight
assertions, and compare each corresponding produced scale buffer between the
relevant expert instances so both tests validate the complete alias contract.
Source: Path instructions
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #2081 +/- ##
===========================================
+ Coverage 67.15% 78.04% +10.88%
===========================================
Files 521 521
Lines 59857 59860 +3
===========================================
+ Hits 40199 46715 +6516
+ Misses 19658 13145 -6513
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:
|
728828b to
8aa14e7
Compare
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
There was a problem hiding this comment.
Warning
CodeRabbit couldn't request changes on this pull request because it doesn't have sufficient GitHub permissions.
Please grant CodeRabbit Pull requests: Read and write permission and re-run the review.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tests/unit/torch/export/test_unified_export_hf.py`:
- Around line 218-225: Update postprocess_state_dict’s tensor
alias-deduplication key to include dtype and layout metadata, specifically shape
and stride, so equal-size views sharing a data pointer remain distinct when they
reference different elements. Modify
test_postprocess_state_dict_preserves_tensors_with_different_byte_ranges to use
two equal-size strided views such as storage[:2] and storage[::2], and retain
assertions that both state-dict entries survive.
🪄 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: 7e9cb918-b967-4680-891c-6f999ba3a6e3
📒 Files selected for processing (7)
modelopt/torch/export/hf_export_handlers.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_unified_export_hf.pytests/unit/torch/quantization/plugins/test_fused_experts.py
🚧 Files skipped from review as they are similar to previous changes (3)
- modelopt/torch/export/quant_utils.py
- tests/unit/torch/quantization/plugins/test_fused_experts.py
- modelopt/torch/export/moe_utils.py
| def test_postprocess_state_dict_preserves_tensors_with_different_byte_ranges(): | ||
| storage = torch.arange(4) | ||
| state_dict = {"short": storage[:2], "long": storage} | ||
| assert state_dict["short"].data_ptr() == state_dict["long"].data_ptr() | ||
|
|
||
| processed = postprocess_state_dict(state_dict, maxbound=448, quantization=None) | ||
|
|
||
| assert set(processed) == set(state_dict) |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
Cover equal-size strided views before deduplicating state-dict tensors.
storage[:2] and storage[::2] have the same data_ptr() and byte size, but they reference different elements. postprocess_state_dict currently treats them as tied and removes one entry. Include dtype and layout metadata, such as shape and stride, in the alias key. Change this test to use two such views.
As per coding guidelines, “Tests must exercise the behavior they claim to validate.”
Proposed regression case
def test_postprocess_state_dict_preserves_tensors_with_different_byte_ranges():
storage = torch.arange(4)
- state_dict = {"short": storage[:2], "long": storage}
- assert state_dict["short"].data_ptr() == state_dict["long"].data_ptr()
+ state_dict = {"contiguous": storage[:2], "strided": storage[::2]}
+ assert state_dict["contiguous"].data_ptr() == state_dict["strided"].data_ptr()
+ assert state_dict["contiguous"].numel() == state_dict["strided"].numel()
processed = postprocess_state_dict(state_dict, maxbound=448, quantization=None)
assert set(processed) == set(state_dict)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| def test_postprocess_state_dict_preserves_tensors_with_different_byte_ranges(): | |
| storage = torch.arange(4) | |
| state_dict = {"short": storage[:2], "long": storage} | |
| assert state_dict["short"].data_ptr() == state_dict["long"].data_ptr() | |
| processed = postprocess_state_dict(state_dict, maxbound=448, quantization=None) | |
| assert set(processed) == set(state_dict) | |
| def test_postprocess_state_dict_preserves_tensors_with_different_byte_ranges(): | |
| storage = torch.arange(4) | |
| state_dict = {"contiguous": storage[:2], "strided": storage[::2]} | |
| assert state_dict["contiguous"].data_ptr() == state_dict["strided"].data_ptr() | |
| assert state_dict["contiguous"].numel() == state_dict["strided"].numel() | |
| processed = postprocess_state_dict(state_dict, maxbound=448, quantization=None) | |
| assert set(processed) == set(state_dict) |
🤖 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 218 - 225,
Update postprocess_state_dict’s tensor alias-deduplication key to include dtype
and layout metadata, specifically shape and stride, so equal-size views sharing
a data pointer remain distinct when they reference different elements. Modify
test_postprocess_state_dict_preserves_tensors_with_different_byte_ranges to use
two equal-size strided views such as storage[:2] and storage[::2], and retain
assertions that both state-dict entries survive.
Sources: Coding guidelines, Path instructions
Signed-off-by: Chad Voegele <cvoegele@nvidia.com>
8aa14e7 to
6da5ff9
Compare
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
There was a problem hiding this comment.
Warning
CodeRabbit couldn't request changes on this pull request because it doesn't have sufficient GitHub permissions.
Please grant CodeRabbit Pull requests: Read and write permission and re-run the review.
Actionable comments posted: 1
🧹 Nitpick comments (3)
modelopt/torch/export/registry.py (1)
55-61: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDocument the new context fields.
The class docstring describes
duplicate_weight_maponly. Add one line each forweight_locations,weight_formats, andskipped_weight_namesso handler authors know which stage populates each field.weight_formatsin particular is filled by_process_quantized_modules, not by__post_init__.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@modelopt/torch/export/registry.py` around lines 55 - 61, Update the containing class docstring to document weight_locations, weight_formats, and skipped_weight_names, including which processing stage populates each field; explicitly state that _process_quantized_modules populates weight_formats rather than __post_init__. Leave the field definitions and behavior unchanged.Source: Coding guidelines
modelopt/torch/export/unified_export_hf.py (1)
857-861: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueSimplify the module-name derivation.
ctx.weight_locations[full_name]already returns the owning module and the weight attribute name. The arithmetic slice reconstructs the module path from string lengths.rsplitexpresses the same result and removes the length bookkeeping.♻️ Proposed simplification
for full_name in ctx.skipped_weight_names: - _, weight_name = ctx.weight_locations[full_name] - module_name = full_name[: -(len(weight_name) + 1)] if "." in full_name else "" - prefix = f"{module_name}." if module_name else "" + module_name, _, weight_name = full_name.rpartition(".") + prefix = f"{module_name}." if module_name else "" attrs = quantizer_attr_names(weight_name)🤖 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 857 - 861, In the loop over ctx.skipped_weight_names, replace the length-based module_name derivation with rsplit on full_name using weight_name as the suffix separator, preserving an empty module name when no module prefix exists. Keep the existing prefix and quantizer_attr_names logic unchanged.tests/unit/torch/quantization/plugins/test_fused_experts.py (1)
696-705: 🗄️ Data Integrity & Integration | 🔵 Trivial | ⚡ Quick winAssert the skipped-weight names in the tied test.
The test checks
duplicate_ofand the resulting module structure. It does not checkctx.skipped_weight_names, which is the field_remove_skipped_duplicate_weightsconsumes to strip the tied entries from the exported state dict. Add that assertion so a regression inmark_skippedfails here instead of surfacing later as duplicated expert weights in a checkpoint.♻️ Proposed assertion
assert ctx.duplicate_of("decoder.experts.down_proj") == "encoder.experts.down_proj" + assert ctx.skipped_weight_names == { + "decoder.experts.gate_up_proj", + "decoder.experts.down_proj", + } assert hasattr(parent.encoder.experts, "0")🤖 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 696 - 705, Add an assertion in the tied-experts test around _process_quantized_modules that ctx.skipped_weight_names contains the expected encoder/decoder tied expert weight names consumed by _remove_skipped_duplicate_weights, covering both gate_up_proj and down_proj entries and their relevant expert paths.Source: Path instructions
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@modelopt/torch/export/hf_export_handlers.py`:
- Around line 43-53: Initialize weight_formats in ExportContext.__post_init__
alongside duplicate_weight_map and weight_locations, using the existing
weight-format population logic so contexts are fully usable at construction.
Keep _is_duplicate_with_same_format unchanged and verify handlers invoked
through _export_transformers_checkpoint and PrepareMoEInputsRegistry are safe
with the initialized context.
---
Nitpick comments:
In `@modelopt/torch/export/registry.py`:
- Around line 55-61: Update the containing class docstring to document
weight_locations, weight_formats, and skipped_weight_names, including which
processing stage populates each field; explicitly state that
_process_quantized_modules populates weight_formats rather than __post_init__.
Leave the field definitions and behavior unchanged.
In `@modelopt/torch/export/unified_export_hf.py`:
- Around line 857-861: In the loop over ctx.skipped_weight_names, replace the
length-based module_name derivation with rsplit on full_name using weight_name
as the suffix separator, preserving an empty module name when no module prefix
exists. Keep the existing prefix and quantizer_attr_names logic unchanged.
In `@tests/unit/torch/quantization/plugins/test_fused_experts.py`:
- Around line 696-705: Add an assertion in the tied-experts test around
_process_quantized_modules that ctx.skipped_weight_names contains the expected
encoder/decoder tied expert weight names consumed by
_remove_skipped_duplicate_weights, covering both gate_up_proj and down_proj
entries and their relevant expert paths.
🪄 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: 232c31ad-1ad5-49f6-8336-3fff83d368a2
📒 Files selected for processing (8)
modelopt/torch/export/hf_export_handlers.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_unified_export_hf.pytests/unit/torch/quantization/plugins/test_fused_experts.py
🚧 Files skipped from review as they are similar to previous changes (1)
- modelopt/torch/export/quant_utils.py
| def _is_duplicate_with_same_format( | ||
| module_name: str, | ||
| ctx: ExportContext, | ||
| weight_name: str, | ||
| ) -> bool: | ||
| """Whether this source weight can be omitted in favor of its canonical tied name.""" | ||
| full_name = _full_weight_name(module_name, weight_name) | ||
| canonical_name = ctx.duplicate_of(full_name) | ||
| if canonical_name is None: | ||
| return False | ||
| return ctx.weight_formats[canonical_name] == ctx.weight_formats[full_name] |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
Guard against an unpopulated weight_formats.
ExportContext.__post_init__ fills duplicate_weight_map and weight_locations, but weight_formats is filled later, inside _process_quantized_modules (modelopt/torch/export/unified_export_hf.py, Lines 814-816). Any handler invoked with a context that has not passed through _process_quantized_modules therefore sees a populated duplicate_weight_map and an empty weight_formats. Line 53 then raises KeyError instead of returning False.
_export_transformers_checkpoint already passes export_ctx to the PrepareMoEInputsRegistry handlers at Line 923, before weight_formats exists. Confirm that no handler on that path reaches this helper. The durable fix is to populate weight_formats in __post_init__ next to the other fields, so the context is fully initialized at construction.
#!/bin/bash
# Description: Find every handler that receives an ExportContext and check which reach _is_duplicate_with_same_format.
set -euo pipefail
rg -nP --type=py -C3 '_is_duplicate_with_same_format|PrepareMoEInputsRegistry\.register|ExportModuleRegistry\.register'
rg -nP --type=py -C4 'def .*\(.*ctx: ExportContext'🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@modelopt/torch/export/hf_export_handlers.py` around lines 43 - 53, Initialize
weight_formats in ExportContext.__post_init__ alongside duplicate_weight_map and
weight_locations, using the existing weight-format population logic so contexts
are fully usable at construction. Keep _is_duplicate_with_same_format unchanged
and verify handlers invoked through _export_transformers_checkpoint and
PrepareMoEInputsRegistry are safe with the initialized context.
| # Use tensor data pointer to identify tied weights | ||
| tensor_id = value.data_ptr() | ||
| 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.
why do we need data_ptr still?
…al_parameters) Independent commit (easy to revert). Recovers the peak-memory and total_parameters regression that removing moe_tied_cache introduced, while keeping the name-based identity that makes dedup FSDP/offload-correct. Complements the approach in PR NVIDIA#2081 (which omits at source but groups by id()); here the omit is driven by the declared name resolver, so it is both memory-efficient AND declaration-gated. When a fused-experts container is the *alias* side of a declared tie (its canonical partner is a different container, exported independently), skip splitting and packing it entirely and delete its fused source, so it emits no state-dict keys. The loader re-ties it from the model's _tied_weights_keys. This avoids materializing a second packed copy of every expert -- the memory cost of pack-then-drop -- and restores the correct total_parameters (tied experts counted once). Gated to non-FSDP: under FSDP2 the container is packed normally and the duplicate is dropped by name in postprocess_state_dict on the gathered state dict, avoiding module-buffer surgery on resharded modules. Dense ties are unchanged (cheap to pack; postprocess drops the alias key by name). - TiedGroupResolver.container_is_alias(container_name, first_proj): True only for the omittable alias side of a declared tie. - _export_fused_experts_module: omit-at-source for the alias container (non-FSDP). - Tests: resolver container_is_alias; handler-level omit (alias container emits no keys, canonical exported in full). On-disk output unchanged vs the pack-then-drop path (same keys dropped, just never materialized); verified by the existing DiffusionGemma identity checks. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Juhi Mittal <juhim@nvidia.com>
2 participants
What does this PR do?
Type of change: ? Bug fix
Fixing false positive duplicate removal in experts
Usage
Testing
Before your PR is "Ready for review"
Make sure you read and follow Contributor guidelines and your commits are signed (
git commit -s -S).Make sure you read and follow the Security Best Practices (e.g. avoiding hardcoded
trust_remote_code=True,torch.load(..., weights_only=False),pickle, etc.).CONTRIBUTING.md: ✅Additional Information
Summary by CodeRabbit
Bug Fixes
Tests