Skip to content

Simplify tied-weight export dedup - #2151

Draft
chadvoegele wants to merge 1 commit into
NVIDIA:mainfrom
chadvoegele:tt-simplify-tied-weight-export-dedup
Draft

Simplify tied-weight export dedup#2151
chadvoegele wants to merge 1 commit into
NVIDIA:mainfrom
chadvoegele:tt-simplify-tied-weight-export-dedup

Conversation

@chadvoegele

Copy link
Copy Markdown
Contributor

What does this PR do?

Type of change: ?

Proposal for a simplified version of #2092

Usage

# Add a code snippet demonstrating how to use this

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.).

  • Is this change backward compatible?: ✅ / ❌ / N/A
  • If you copied code from any other sources or added a new PIP dependency, did you follow guidance in CONTRIBUTING.md: ✅ / ❌ / N/A
  • Did you write any new necessary tests?: ✅ / ❌ / N/A
  • Did you update Changelog?: ✅ / ❌ / N/A
  • Did you get Claude approval on this PR?: ✅ / ❌ / N/A

Additional Information

Signed-off-by: Chad Voegele <cvoegele@nvidia.com>
@copy-pr-bot

copy-pr-bot Bot commented Aug 11, 2026

Copy link
Copy Markdown

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

@coderabbitai

coderabbitai Bot commented Aug 11, 2026

Copy link
Copy Markdown
Contributor

Important

Review skipped

Draft detected.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 21fff3a0-cdb5-44d8-ab97-57126a40cdc7

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

  • 🔍 Trigger review

Comment @coderabbitai help to get the list of available commands.

"""
return has_accelerate_offload(module) or is_fsdp2_model(module)


Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@Fridah-nv you recently added this functionality, does this break any of your logic? or is it safe to remove this?

@codecov

codecov Bot commented Aug 11, 2026

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 84.09091% with 14 lines in your changes missing coverage. Please review.
✅ Project coverage is 78.71%. Comparing base (e4fe1e5) to head (33292bc).
⚠️ Report is 1 commits behind head on main.

Files with missing lines Patch % Lines
modelopt/torch/export/model_utils.py 86.95% 6 Missing ⚠️
modelopt/torch/export/quant_utils.py 83.33% 6 Missing ⚠️
modelopt/torch/export/hf_export_handlers.py 50.00% 1 Missing ⚠️
modelopt/torch/export/unified_export_hf.py 50.00% 1 Missing ⚠️
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     
Flag Coverage Δ
unit 55.30% <84.09%> (+0.02%) ⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

juhi10071998 added a commit to juhi10071998/Model-Optimizer that referenced this pull request Aug 12, 2026
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>
juhi10071998 added a commit to juhi10071998/Model-Optimizer that referenced this pull request Aug 12, 2026
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>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants