Support saving and loading ShardedTensor. (#62242)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/62242
1) Add a state_dict hook to ensure ShardedTensors are
added to a state_dict.
2) Add a pre load state_dict hook to ensure ShardedTensor are added back to a
module at load time.
3) Add a `with_load_process_group` context manager for load time.
4) Added ser-de capability to ShardedTensor.
ghstack-source-id: 134860967
Test Plan:
1) unit tests
2) waitforbuildbot
Reviewed By: wanchaol
Differential Revision: D29927881
fbshipit-source-id: b1ef8872ed91e9cb0e2d5dd17d2764678ab89f0c
