Yi Wang 71f3399e19 [Gradient Compression] Implement the original layerwise PowerSGD (#49417) ... Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/49417 The existing implementation applies PowerSGD to a batch of flattened tensors, which is a coarse-grained compression. This hook now is renamed as "batched_powerSGD_hook". Now implement the original implementation in the paper, which applies PowerSGD to each per-parameter tensor. This is a layerwise fine-grained compression. Although this original implementation is slower, it is expected to achieve a higher accuracy, especially when the shapes of per-param tensors cannot be aligned. Also add a test in distributed_test.py. Original PR issue: Investigate Applying PowerSGD to Communication Hook for Gradient Compression #47202 ghstack-source-id: 118921275 Test Plan: buck test mode/dev-nosan caffe2/test/distributed:c10d -- test_powerSGD_ddp_comm_hook_nccl buck test mode/dev-nosan caffe2/test/distributed:distributed_nccl_fork -- test_DistributedDataParallel_powerSGD_ddp_comm_hook Reviewed By: rohan-varma Differential Revision: D25511543 fbshipit-source-id: 19ef188bc2d4c7406443c8fa233c1f2c2f27d93c		2020-12-18 18:02:15 -08:00
..
ddp_comm_hooks	[Gradient Compression] Implement the original layerwise PowerSGD (#49417)	2020-12-18 18:02:15 -08:00
__init__.py	[Gradient Compression] Add unit tests that test default Python comm hook implementations (#47158)	2020-11-06 00:28:09 -08:00