remove compressed allreduce

Author: server-jack

deepspeed/runtime/engine.py

@@ -26,7 +26,6 @@
 from deepspeed.runtime.config import DeepSpeedConfig, DEEPSPEED_OPTIMIZERS, \
     ADAM_OPTIMIZER, ADAMW_OPTIMIZER, LAMB_OPTIMIZER, ONEBIT_ADAM_OPTIMIZER, ONEBIT_LAMB_OPTIMIZER, \
     TORCH_ADAM_PARAM, ADAM_W_MODE, ADAM_W_MODE_DEFAULT
-from deepspeed.runtime.comm import compressed_all_reduce
 
 from deepspeed.runtime.dataloader import DeepSpeedDataLoader
 from deepspeed.runtime.constants import \
@@ -139,7 +138,6 @@ def __init__(self,
         self.store_gradients = False
         self.store_gradients_cpu = False
         self.stored_gradients = None
-        self.bf16_compressed_allreduce = False # hardcode for now - it's not really working
 
         if dist_init_required is None:
             dist_init_required = not dist.is_initialized()
@@ -1292,29 +1290,23 @@ def allreduce_bucket(self, bucket):
 
         tensor_to_allreduce = tensor
 
-        if self.allreduce_always_fp32() and not self.bf16_compressed_allreduce:
+        if self.allreduce_always_fp32():
             tensor_to_allreduce = tensor.float()
 
         if self.postscale_gradients():
             if self.gradient_predivide_factor() != 1.0:
                 tensor_to_allreduce.mul_(1. / self.gradient_predivide_factor())
-            if self.bf16_compressed_allreduce and self.precision() == torch.bfloat16:
-                compressed_all_reduce(tensor_to_allreduce, group=self.data_parallel_group)
-            else:
-                dist.all_reduce(tensor_to_allreduce, group=self.data_parallel_group)
+            dist.all_reduce(tensor_to_allreduce, group=self.data_parallel_group)
 
             if self.gradient_average:
                 if self.gradient_predivide_factor() != self.dp_world_size:
                     tensor_to_allreduce.mul_(self.gradient_predivide_factor() /
                                              self.dp_world_size)
         else:
             tensor_to_allreduce.div_(self.dp_world_size)
-            if self.bf16_compressed_allreduce and self.precision() == torch.bfloat16:
-                compressed_all_reduce(tensor_to_allreduce, group=self.data_parallel_group)
-            else:
-                dist.all_reduce(tensor_to_allreduce, group=self.data_parallel_group)
+            dist.all_reduce(tensor_to_allreduce, group=self.data_parallel_group)
 
-        if self.allreduce_always_fp32() and tensor is not tensor_to_allreduce and not self.bf16_compressed_allreduce:
+        if self.allreduce_always_fp32() and tensor is not tensor_to_allreduce:
             tensor.copy_(tensor_to_allreduce)
 
         return tensor