[FSDP][Easy] Minor cleanup (#84761)
This PR simply pulls out some minor changes from the next (monolithic) PR.
Differential Revision: [D39392147](https://our.internmc.facebook.com/intern/diff/D39392147)
Pull Request resolved: https://github.com/pytorch/pytorch/pull/84761
Approved by: https://github.com/zhaojuanmao
