pytorch/torch/distributed/optim
Rohan Varma 1dba329d20 Enable step_param for Adam functional optimizer (#62611)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/62611

Enables optimizer overlap with backwards in DDP for Adam. Additional optimizers, especially Adagrad will be done in follow up diffs.

1. Implement `step_param` method based on `step` in _FunctionalAdam (perf permitting we can later dedupe `step` to call `step_param`
2. Modify tests to test all current functional optimizers.
ghstack-source-id: 135207143

Test Plan: CI

Reviewed By: SciPioneer

Differential Revision: D29891783

fbshipit-source-id: 321915982afd5cb0a9c2e43d27550f433bff00d1
2021-08-06 10:53:55 -07:00
..
__init__.py [Model Averaging] Post-localSGD optimizer (#62131) 2021-07-28 18:42:06 -07:00
functional_adadelta.py [dist_optim] fix the bug of none grads on functional optimizers (#62249) 2021-07-27 18:10:51 -07:00
functional_adagrad.py [optim] take kw-only argument for functional optim APIs (#56185) 2021-04-15 20:08:04 -07:00
functional_adam.py Enable step_param for Adam functional optimizer (#62611) 2021-08-06 10:53:55 -07:00
functional_adamax.py [optim] take kw-only argument for functional optim APIs (#56185) 2021-04-15 20:08:04 -07:00
functional_adamw.py [optim] take kw-only argument for functional optim APIs (#56185) 2021-04-15 20:08:04 -07:00
functional_rmsprop.py [dist_optim] fix the bug of none grads on functional optimizers (#62249) 2021-07-27 18:10:51 -07:00
functional_rprop.py [dist_optim] fix the bug of none grads on functional optimizers (#62249) 2021-07-27 18:10:51 -07:00
functional_sgd.py [dist_optim] fix the bug of none grads on functional optimizers (#62249) 2021-07-27 18:10:51 -07:00
optimizer.py [doc] update distributed optimizer doc (#58084) 2021-05-13 23:37:00 -07:00
post_localSGD_optimizer.py [Model Averaging] Update model averager API to avoid the redundant params arg needed by post-localSGD optimizer (#62132) 2021-07-28 18:43:09 -07:00
zero_redundancy_optimizer.py Make _Join, _Joinable, _JoinHook public (#62605) 2021-08-03 12:20:11 -07:00
zero_redundancy_optimizer.pyi Make _Join, _Joinable, _JoinHook public (#62605) 2021-08-03 12:20:11 -07:00