[FSDP] Implement sharded_optim_state_dict and flatten_sharded_optim_state_dict. (#77628)
As title
Differential Revision: [D36436496](https://our.internmc.facebook.com/intern/diff/D36436496/)
Pull Request resolved: https://github.com/pytorch/pytorch/pull/77628
Approved by: https://github.com/awgu
