[Model Averaging] Add a few member methods of PostLocalSGDOptimizer (#63340)
authorYi Wang <wayi@fb.com>
Mon, 16 Aug 2021 23:33:21 +0000 (16:33 -0700)
committerFacebook GitHub Bot <facebook-github-bot@users.noreply.github.com>
Mon, 16 Aug 2021 23:39:01 +0000 (16:39 -0700)
commit068d6fec5cb36cecd757f769c2b245a2c77555af
tree5b9c840c1c5f178404be4e9817fff216f858d703
parentaa63c0d9df54ea1fe65791143c8cee6b34cab4f2
[Model Averaging] Add a few member methods of PostLocalSGDOptimizer (#63340)

Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/63340

Some methods are needed such as accessing optimizer states. These are necessary for integration with PyTorch Lightning.

Proposal: https://github.com/pytorch/pytorch/issues/59699
ghstack-source-id: 135912246

Test Plan: buck test mode/dev-nosan caffe2/test/distributed:distributed_nccl_fork -- test_ddp_hook_parity_post_localSGD

Reviewed By: rohan-varma

Differential Revision: D30328794

fbshipit-source-id: e585b874313bd266fdc7c79936e2af98700c7bad
torch/distributed/optim/post_localSGD_optimizer.py