[Gradient Compression] Replace the assertions in PowerSGD comm hook b…

…y stream syncrhonization (#49435)

Summary:
Pull Request resolved: #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

torch/distributed/algorithms/ddp_comm_hooks/powerSGD_hook.py

@@ -268,7 +268,8 @@ def decompress(fut):
 
         for p, q, tensor in zip(ps, qs, high_rank_tensors):
             torch.matmul(p, q.t(), out=tensor)
-            assert not torch.any(torch.isnan(tensor))
+        if torch.cuda.is_available():
+            torch.cuda.synchronize()
 
         if state.use_error_feedback:
             # Memorize the local errors.
@@ -414,7 +415,8 @@ def decompress(fut):
         if state.use_error_feedback:
             # Memorize the local errors.
             state.error_dict[bucket_index] = input_tensor_cp - input_tensor
-            assert not torch.any(torch.isnan(state.error_dict[bucket_index]))
+        if torch.cuda.is_available():
+            torch.cuda.synchronize()
         ret = input_tensor.resize_(total_length)
         return [ret]