Resolve tied weights by declared name in HF export (FSDP/offload-correct)

[juhi10071998]

What does this PR do?

Type of change: Bug fix / robustness

Fixes NVBug 6525352 — MiniMax-M2.7
nvfp4_mlp_only-kv_fp8 failed at TensorRT-LLM load with assert w1_weight is not None and w3_weight is not None, because the previous data_ptr()-only postprocess dedup (no size
check) could falsely drop an independent MoE expert weight.

Resolve tied weights during unified HF checkpoint export by their declared name
(from the model's _tied_weights_keys / tie_word_embeddings) instead of by tensor
data_ptr(). This implements the TODO already noted on ExportContext
("dedup by tied-group name instead… which survives weight moves").

Why address identity is wrong here. data_ptr() answers "do these bytes start at the
same location?", not "are these the same declared parameter?". Those diverge:

• False positive: the fused-MoE source tensors are deleted mid-export, and PyTorch's
  allocator can recycle a freed address for a later, unrelated module — falsely aliasing it.
• False negative (the one that matters for scale): under the FSDP full_state_dict
  gather (and CPU/disk offload), tied weights are materialized at distinct addresses, so
  the address-based dedup misses a genuine tie and writes both copies. This is why tied-weight
  export was resident-path only. Name-based dedup runs on the final gathered state dict and is
  correct on 1 GPU and under multi-GPU FSDP (e.g. a 4-GPU MiniMax export that doesn't fit
  on one device).

What changed

• New TiedGroupResolver (model_utils.py): builds {alias → canonical} from dict-style
  _tied_weights_keys (with per-layer regex backreferences and the container-level fused-MoE
  tie) plus tie_word_embeddings. Uses 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 single authoritative dedup: drops each declared alias key
  whose canonical is present (address-independent). A (device, data_ptr, size) pass is kept
  as a backstop for undeclared / coincidental shares; deletion is idempotent.
• One resolver is threaded through sync_tied_input_amax (name-based grouping), the
  ExportContext MoE cache, and postprocess. sync_tied_input_amax still runs before
  packing so the retained weight's single input_scale covers every side's activation range.
• The per-module dense tied cache is removed — both sides pack identically and the
  duplicate is dropped by name. The fused-MoE cache is kept purely as a resident-path
  compute/memory optimization (skips re-unpacking 128 experts), name-keyed, and still disabled
  under has_non_resident_weights (FSDP2 / offload).

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 via _tied_weights_keys and
is unaffected.

Usage

No API change. Existing export just works — and now dedups tied weights correctly under FSDP /
offload as well as on a single GPU:

from modelopt.torch.export import export_hf_checkpoint

export_hf_checkpoint(model, export_dir="./exported")  # tied weights collapse by declared name

Testing

Unit tests (CPU): resolver alias-map + per-layer backreference; the FSDP case simulated on
CPU — declared alias dropped by name across distinct addresses; per-expert MoE subtree
drop; keep-alias-when-canonical-absent; name-based sync_tied_input_amax; MoE tied/untied
cache on the declared-tie + name-key contract; offload ExportContext residency guards.
tests/unit/torch/export/ + test_fused_experts.py + test_quant_embedding.py: 220 passed
(4 pre-existing test_quant_aware_conversion.py scoped-rule failures on main are unrelated).
End-to-end DiffusionGemma (single-GPU) and MiniMax (4-GPU FSDP) verification is in progress.

Before your PR is "Ready for review"

• Is this change backward compatible?: ✅ (on-disk output unchanged for declared ties / untied models)
• If you copied code from any other sources or added a new PIP dependency…: N/A
• Did you write any new necessary tests?: ✅
• Did you update Changelog?: ✅ (0.47 → Bug Fixes)
• Did you get Claude approval on this PR?: ❌ (draft)

Additional Information

Related bug: NVBug 6525352 (MiniMax-M2.7 TRT-LLM
load failure — root cause absent after this change).

Dedup happens in a single place — the name-based drop in postprocess_state_dict. A size-aware
(device, data_ptr, size) address pass is retained only as a backstop for undeclared
same-storage shares (which safetensors.save_file would otherwise reject); it only ever fires
for unquantized/unpacked shared weights, since quantized ties pack to distinct storage and are
collapsed by name.

End-to-end verification (4× GB200):

• MiniMax-M2.7 (nvfp4_mlp_only-kv_fp8): 191211 tensors, 15872/15872 experts, 0 missing — the
  6525352 w1_weight is None root cause is absent.
• DiffusionGemma-26B (nvfp4_experts_only): tied encoder↔decoder experts + lm_head↔embeddings
  collapse to the canonical side (0 leaked keys); matches the known-good baseline. Served in vLLM
  (vllm/vllm-openai:gemma) and generates correctly (e.g. 17 × 23 = 391 with a reasoning trace).

[copy-pr-bot]

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.

[coderabbitai]

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: 5077441b-d464-4a4d-9563-6d3e538af5e6

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

[codecov]

Codecov Report

❌ Patch coverage is 78.62595% with 28 lines in your changes missing coverage. Please review.
✅ Project coverage is 77.40%. Comparing base (22b6a14) to head (3c8cbcb).
⚠️ Report is 12 commits behind head on main.

Files with missing lines	Patch %	Lines
modelopt/torch/export/model_utils.py	76.13%	21 Missing ⚠️
modelopt/torch/export/quant_utils.py	89.65%	3 Missing ⚠️
modelopt/torch/export/unified_export_hf.py	40.00%	3 Missing ⚠️
modelopt/torch/export/hf_export_handlers.py	50.00%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2092      +/-   ##
==========================================
- Coverage   78.60%   77.40%   -1.20%     
==========================================
  Files         522      522              
  Lines       60167    62504    +2337     
==========================================
+ Hits        47294    48383    +1089     
- Misses      12873    14121    +1248     

Flag	Coverage Δ
unit	55.30% <78.62%> (-0.09%)	⬇️

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.