Warn when FSDP auto-wrap policy splits tied weights

[c-pozzi]

Summary

• Adds detection of shared (tied) parameters that would be placed in separate FSDP units by the auto_wrap_policy
• Emits a rank_zero_warn before wrapping, turning a cryptic RuntimeError: size mismatch into an actionable message
• Applies to both Fabric (setup_module) and PyTorch Lightning (_setup_model) FSDP strategies

Motivation

Models like Llama, GPT-2, and Mistral tie their input embedding and output head weights. When users include torch.nn.Embedding in their FSDP auto-wrap policy, the embedding gets its own FSDP unit while the tied lm_head stays in the root unit. FSDP shards each unit independently, so lm_head sees a flat/sharded tensor instead of the expected 2D weight — causing a size mismatch deep in torch with no indication of the real cause.

Test plan

• Tied weights across FSDP units → warning emitted
• Tied weights in same FSDP unit → no warning
• No shared params → no warning
• No policy set → no warning

Closes #21403

🤖 Generated with Claude Code

---

📚 Documentation preview 📚: https://pytorch-lightning--21613.org.readthedocs.build/en/21613/

[c-pozzi]

@justusschock

[codecov]

Codecov Report

❌ Patch coverage is 97.50000% with 1 line in your changes missing coverage. Please review.
✅ Project coverage is 87%. Comparing base (bb7820f) to head (ac456ab).
⚠️ Report is 9 commits behind head on master.

Additional details and impacted files

@@           Coverage Diff           @@
##           master   #21613   +/-   ##
=======================================
  Coverage      87%      87%           
=======================================
  Files         270      270           
  Lines       23934    23974   +40     
=======================================
+ Hits        20713    20749   +36     
- Misses       3221     3225    +4