[Gradient Compression] Replace the assertions in PowerSGD comm hook by stream syncrhonization (#49435)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/49435
Previously the assertion that prevents illegal memory access is because of the torch.any that returns a boolean value, which initiates a data transfer from the device to the host and forces a synchronization.
An explicit synchronization is more to the point.
Original PR issue: Investigate Applying PowerSGD to Communication Hook for Gradient Compression #47202
ghstack-source-id: 118664204
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: D25573484
fbshipit-source-id: 516d0d502da2863b516c15332702335ee662f072
