[FSDP] Generalize fsdp_modules() (#73553)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/73553
Generalize fsdp_modules():
1. Make fsdp_modules() a static method so that it can be called even when users have no reference to the FSDP module -- it can be wrapped inside an non-FSDP module.
2. Add root_only option to returns only the root FSDP module(s).
ghstack-source-id: 150690955
Test Plan: CI
Reviewed By: rohan-varma
Differential Revision: D34537408
fbshipit-source-id: 1bf057704954e8177d0a28f2b3d0dae4da882249
(cherry picked from commit 437aad7a3b612ca275213c32c2fa6ac37f2ed3ab)
