pytorch
bd95cf44 - [DDP] Implement a hook which performs FunctionalSGD step. (#61678)

Commit
3 years ago
[DDP] Implement a hook which performs FunctionalSGD step. (#61678) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/61678 This diff makes the following changes: - Add `step_param` method to `_FunctionalSGD` class which is written similar to `step` but for a single param - Implement a communication hook wrapper that runs a given comm. hook and then applies functional SGD step - Verifies that this is equal to regular allreduce + SGD optimizerghstack-source-id: 133567598 ghstack-source-id: 134263399 Test Plan: CI Reviewed By: SciPioneer Differential Revision: D29701447 fbshipit-source-id: 183954593b82a092414623292f9b10e675fef96e
Author
Parents
Loading