pytorch
ecb5ac90 - [Gradient Compression] Add get_per_parameter_tensors method to GradBucket class (#53009)

Commit
3 years ago
[Gradient Compression] Add get_per_parameter_tensors method to GradBucket class (#53009) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/53009 It can be a common operation to apply layer-wise operations over per-parameter tensors in a DDP communication hook. Create a util method in GradBucket class before publishing GradBucket APIs. ghstack-source-id: 122833594 Test Plan: buck test mode/dev-nosan caffe2/test/distributed:c10d -- test_powerSGD_ddp_comm_hook_nccl f254364097 Reviewed By: rohan-varma Differential Revision: D26717893 fbshipit-source-id: 916db319de8b85dd22bc4e35db5671bf4e34740f
Author
Yi Wang
Parents
Loading