Simplify tied-weight export dedup - #2151
Conversation
Signed-off-by: Chad Voegele <cvoegele@nvidia.com>
|
Important Review skippedDraft detected. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
Comment |
| """ | ||
| return has_accelerate_offload(module) or is_fsdp2_model(module) | ||
|
|
||
|
|
There was a problem hiding this comment.
@Fridah-nv you recently added this functionality, does this break any of your logic? or is it safe to remove this?
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #2151 +/- ##
==========================================
- Coverage 78.73% 78.71% -0.02%
==========================================
Files 522 522
Lines 60357 60334 -23
==========================================
- Hits 47523 47494 -29
- Misses 12834 12840 +6
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:
|
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>
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>
What does this PR do?
Type of change: ?
Proposal for a simplified version of #2092
Usage
# Add a code snippet demonstrating how to use thisTesting
Before your PR is "Ready for review"
Make sure you read and follow Contributor guidelines and your commits are signed (
git commit -s -S).Make sure you read and follow the Security Best Practices (e.g. avoiding hardcoded
trust_remote_code=True,torch.load(..., weights_only=False),pickle, etc.).CONTRIBUTING.md: ✅ / ❌ / N/AAdditional Information