pytorch/torch/distributed/algorithms/model_averaging
Yi Wang c1415a0a72 [Reland] [Model Averaging] Simplify PostLocalSGD Optimizer API (#65197)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/65197

1. The constructor accepts a local optimizer instance instead of the inputs of local optimizer constructor and the class type.
2. The parameters are read from local optimizer's param_groups instead of a separate input.

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

Test Plan: buck test mode/dev-nosan //caffe2/test/distributed:distributed_nccl_spawn -- test_post_localSGD_optimizer_parity

Reviewed By: rohan-varma

Differential Revision: D31007439

fbshipit-source-id: bbb0526e6763ef76775b85088571506b3942c722
2021-09-17 10:31:58 -07:00
..
__init__.py [Model Averaging] Provide a util function for model averaging (#60303) 2021-06-23 15:41:15 -07:00
averagers.py [Model Averaging] Fix docstring of PeriodicModelAverager (#62392) 2021-07-29 17:26:27 -07:00
utils.py [Reland] [Model Averaging] Simplify PostLocalSGD Optimizer API (#65197) 2021-09-17 10:31:58 -07:00