[Reland][FSDP] Another fix for DTensor, use_orig_params=True

[awgu]

Stack from ghstack:

• [FSDP][Perf] Pre-allocate full prec sharded param #90619 [FSDP][Perf] Pre-allocate full prec sharded param
• [FSDP][Easy][BE] Minor cleanup of post-backward logic #90618 [FSDP][Easy][BE] Minor cleanup of post-backward logic
• [FSDP][Perf] Pre-allocate sharded grad in default stream #90617 [FSDP][Perf] Pre-allocate sharded grad in default stream
• [FSDP][Perf] Pre-allocate sharded grad in default stream; save one copy when downcasting grad #90616 [FSDP][Perf] Save one copy when downcasting grad
• [FSDP][Perf] Pre-allocate padded unsharded grad in default stream #90614 [FSDP][Perf] Pre-allocate padded unsharded grad in default stream
• [FSDP] Sanitize HandleConfig for mixed precision #90631 [FSDP] Sanitize HandleConfig for mixed precision
• [FSDP] Tighten post-bwd cast to reduce_dtype #90615 [FSDP] Tighten post-bwd cast to reduce_dtype
• [FSDP][Easy] Move to _storage() in test file #90622 [FSDP][Easy] Move to _storage() in test file
• [FSDP] Save _stream_to_name for debugging #90611 [FSDP] Save _stream_to_name for debugging
• [Reland][FSDP] Another fix for DTensor, use_orig_params=True #90562 [Reland][FSDP] Another fix for DTensor, use_orig_params=True

This is a reland of #89845 with nothing changed. This should avoid the internal breakage now that DTensor does not import torchgen (#90106).

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/90562

• 📄 Preview Python docs built from this PR
• 📄 Preview C++ docs built from this PR
• ❓ Need help or want to give feedback on the CI? Visit our office hours

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit 42395ae:
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[awgu]

@awgu has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

[fduwjj]

LGTM

[awgu]

@pytorchbot merge

[pytorchmergebot]

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here