Fix FSDP2 handling for tied embeddings

[realAsma]

What does this PR do?

Type of change: Bug fix

Groups distinct input and output embedding modules that share a weight into one FSDP2 parameter group before Hugging Face Trainer and Accelerate apply their normal auto-wrapping. This prevents PyTorch from rejecting the shared parameter when distillation prepares both the student and teacher models.

The path is limited to FSDP2 models with genuinely tied input/output weights; untied models and other distributed backends are unchanged.

Usage

No user-facing API change is required. Existing FSDP2 distillation commands work without additional configuration.

Testing

• pytest_pwd tests/unit/torch/opt/plugins/test_transformers_fsdp.py tests/unit/torch/distill/plugins/test_huggingface_kd.py -q (6 passed)
• Changed-file pre-commit suite (passed)
• Two-rank, one-step offline distillation in nvcr.io/nvidia/tensorrt-llm/release:1.3.0rc22 with PyTorch 2.12.0a0, Accelerate 1.14.0, and Transformers 5.5.4:
  • Baseline reproduced the duplicate tied-parameter error.
  • The production change completed successfully with the standalone workaround disabled.

Before your PR is "Ready for review"

• Is this change backward compatible?: ✅
• 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?: ✅
• Did you update Changelog?: N/A
• Did you get Claude approval on this PR?: N/A

Additional Information

The failure occurs before the first training step when FSDP2 separately wraps the two modules that own a tied embedding weight.

[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: f166cfc7-437f-41b7-88ca-359f424334b1

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

[realAsma]

BB: this does not look like a ModelOpt side error to me. This is related to Huggingface and PyTorch. Can you check or search how Huggingface handles tied weights with FSDP2 and report here?

[github-actions]

PR Preview Action v1.8.1
🚀 View preview at https://NVIDIA.github.io/Model-Optimizer/pr-preview/pr-2059/
Built to branch gh-pages at 2026-08-03 23:46 UTC. Preview will be ready when the GitHub Pages deployment is complete.

[realAsma]

🤖 Bot comment.

🐝 I’m checking how Hugging Face handles tied weights with FSDP2 and will report back here.

[realAsma]

🤖 Bot comment.

Checked upstream handling. This is not a ModelOpt-originated error:

• PyTorch FSDP2 deliberately requires shared/tied parameters to be assigned to the same fully_shard instance; its RFC states that shared parameters must be in the same group, and the runtime error prescribes fully_shard([module_a, module_b]): [RFC] Per-Parameter-Sharding FSDP pytorch/pytorch#114299 (comment)
• Accelerate's current FSDP2 preparation walks auto-wrap matches bottom-up and then calls fully_shard(model); it does not detect and group the distinct input/output embedding modules that own one shared parameter: https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/fsdp_utils.py
• Hugging Face's recent tied-embedding fix addresses a different case: the config says tied but the checkpoint contains separate weights. Their guidance there is tie_word_embeddings=False: FSDP2 fails due to KeyError: 'lm_head.weight' huggingface/accelerate#3870 and Fix FSDP2 tied embedding errors with targeted ValueError guidance huggingface/accelerate#3878
• Our reproducer has genuinely tied weights (the two attributes are the same Parameter), so disabling tying would change the model and is not the right remedy.

Conclusion: the grouping in this PR matches PyTorch's required workaround, but it is compensating for an upstream Accelerate/FSDP2 integration gap. I recommend filing/fixing this in Accelerate; we should keep the ModelOpt workaround only if we explicitly want a compatibility bridge until an upstream release contains the fix.

[realAsma]

🤖 Bot comment.

Confirmed: this is an upstream Hugging Face/PyTorch integration issue, not ModelOpt-owned behavior.

Accelerate's FSDP2 preparation independently calls fully_shard(module, ...) for each auto-wrap match before sharding the root (source). PyTorch explicitly rejects a tied/shared parameter split across FSDP groups and directs callers to fully_shard([module_a, module_b]) (source). The workaround in this PR follows that PyTorch prescription.

Hugging Face also recently merged native FSDP2 handling that rewrites the input-embedding/output-head plan when they are tied before applying fully_shard (transformers#46707); see _resolve_tied_embed_lm_head_plan). That is further evidence the ownership belongs in the HF distributed integration layer.

Recommendation: report/fix the legacy Trainer→Accelerate FSDP2 path upstream. I would not merge a permanent ModelOpt-specific policy copy unless we deliberately want a version-gated compatibility workaround for affected HF releases. No code changes prepared in this pass.

[codecov]

Codecov Report

❌ Patch coverage is 82.35294% with 3 lines in your changes missing coverage. Please review.
✅ Project coverage is 67.00%. Comparing base (14b20c0) to head (d7601cf).
⚠️ Report is 1 commits behind head on main.

Files with missing lines	Patch %	Lines
modelopt/torch/opt/plugins/transformers.py	82.35%	3 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2059      +/-   ##
==========================================
- Coverage   67.00%   67.00%   -0.01%     
==========================================
  Files         520      520              
  Lines       59545    59562      +17     
==========================================
+ Hits        39900    39908       +8     
- Misses      19645    19654       +9     

Flag	Coverage Δ
unit	55.25% <82.35%> (-0.01%)	⬇️

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.