[FSDP] Provide a utility API to allow users easily to set state_dict_type (#73716)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/73716
fsdp_state_dict_type() traverses the modules and calls state_dict_type when it is necessary, eliminating the burden from the users.
ghstack-source-id: 150775052
Test Plan: CI
Reviewed By: rohan-varma
Differential Revision: D34532321
fbshipit-source-id: 339dcb14a692d5d83c266892644393464d77db40
(cherry picked from commit bb4634c77d6b46644a7daad005ff52738753221a)
