[Gradient Compression] Refactor default_hooks.py and powerSGD_hook.py by creating a util function that make a vanilla allreduce future

torch/distributed/algorithms/ddp_comm_hooks/default_hooks.py

@@ -2,6 +2,20 @@
 import torch.distributed as dist
 
 
+def allreduce_fut(
+    process_group: dist.ProcessGroup, tensor: torch.Tensor
+) -> torch.futures.Future:
+    group_to_use = process_group if process_group is not None else dist.group.WORLD
+
+    "Averages the input gradient tensor by allreduce and returns a future."
+    fut = dist.all_reduce(tensor, group=group_to_use, async_op=True).get_future()
+
+    def div_by_group_size(fut):
+        return [fut.value()[0].div_(group_to_use.size())]
+
+    return fut.then(div_by_group_size)
+
+
 def allreduce_hook(
     process_group: dist.ProcessGroup, bucket: dist._GradBucket
 ) -> torch.futures.Future:
@@ -17,16 +31,7 @@ def allreduce_hook(
     Example::
         >>> ddp_model.register_comm_hook(process_group, allreduce_hook)
     """
-    group_to_use = process_group if process_group is not None else dist.group.WORLD
-    world_size = group_to_use.size()
-
-    tensor = bucket.get_tensors()[0]
-    fut = dist.all_reduce(tensor, group=group_to_use, async_op=True).get_future()
-
-    def then_callback(fut):
-        return [fut.value()[0].div_(world_size)]
-
-    return fut.then(then_callback)
+    return allreduce_fut(process_group, bucket.get_tensors()[0])
 
 
 def fp16_compress_hook(

torch/distributed/algorithms/ddp_comm_hooks/powerSGD_hook.py

@@ -4,6 +4,7 @@
 import numpy as np
 import torch
 import torch.distributed as dist
+import torch.distributed.algorithms.ddp_comm_hooks.default_hooks as default
 
 
 def _orthogonalize(matrix, epsilon=1e-8):
@@ -127,7 +128,7 @@ def maybe_increase_iter(self, bucket):
 
         if self.iter == self.start_powerSGD_iter:
             logging.info(
-                "Starting to apply PowerSGD after {} iterations.".format(self.iter)
+                "Start to apply PowerSGD after {} iterations.".format(self.iter)
             )
 
 
@@ -183,15 +184,8 @@ def powerSGD_hook(state: PowerSGDState, bucket) -> torch.futures.Future:
 
     # Run vanilla allreduce in the first `start_powerSGD_iter` iterations.
     if state.iter < state.start_powerSGD_iter:
-        fut = dist.all_reduce(
-            input_tensor, group=group_to_use, async_op=True
-        ).get_future()
-
-        def div_callback(fut):
-            return [fut.value()[0].div_(world_size)]
-
         state.maybe_increase_iter(bucket)
-        return fut.then(div_callback)
+        return default.allreduce_fut(group_to_use, input_tensor)
 
     # Apply PowerSGD after `start_powerSGD_iter` iterations.
     device = input_tensor.device