[FSDP] Generic arguments for state_dict (#73323)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/73323
These arguments are just passed to nn.Module APIs, make them generic
args/kwargs so the main APIs can update their args and we won't have to.
ghstack-source-id: 149988160
Test Plan: CI
Reviewed By: fegin
Differential Revision: D34431023
fbshipit-source-id: 12e7cc47867655e7eeb76092e1e663c442594210
(cherry picked from commit 9c080894d950b0184887db8890d6036d8284519a)