Revert "make fp32 allreduce optional"

This reverts commit 027485a.

It was breaking neox training for a reason I can't yet figure out

Author: sdtblck

deepspeed/runtime/engine.py

@@ -803,6 +803,7 @@ def _configure_fp16_optimizer(self, optimizer):
     def _configure_zero_optimizer(self, optimizer):
         zero_stage = self.zero_optimization_stage()
         log_dist('Creating fp16 ZeRO stage {} optimizer'.format(zero_stage), ranks=[0])
+        assert not self.allreduce_always_fp32(), "ZeRO does not support 'fp32_allreduce': true"
         timers = self.timers if self.wall_clock_breakdown() else None
 
         if zero_stage == ZERO_OPTIMIZATION_OPTIMIZER_STATES:
@@ -818,8 +819,7 @@ def _configure_zero_optimizer(self, optimizer):
                 dp_process_group=self.data_parallel_group,
                 elastic_checkpoint=self.zero_elastic_checkpoint(),
                 mpu=self.mpu,
-                precision=self.precision(),
-                fp32_allreduce=self.allreduce_always_fp32)
+                precision=self.precision())
         elif zero_stage == ZERO_OPTIMIZATION_GRADIENTS:
             optimizer = FP16_DeepSpeedZeroOptimizer(
                 optimizer,
@@ -839,8 +839,7 @@ def _configure_zero_optimizer(self, optimizer):
                 postscale_gradients=self.postscale_gradients(),
                 gradient_predivide_factor=self.gradient_predivide_factor(),
                 gradient_accumulation_steps=self.gradient_accumulation_steps(),
-                precision=self.precision(),
-                fp32_allreduce=self.allreduce_always_fp32)
+                precision=self.precision())
         elif zero_stage == ZERO_OPTIMIZATION_WEIGHTS:
             print("Initializing ZeRO Stage 3") if dist.get_rank() == 0 else None
             from deepspeed.runtime.zero.stage3 import FP16_DeepSpeedZeroOptimizer_Stage3
@@ -868,8 +867,7 @@ def _configure_zero_optimizer(self, optimizer):
                 postscale_gradients=self.postscale_gradients(),
                 gradient_predivide_factor=self.gradient_predivide_factor(),
                 gradient_accumulation_steps=self.gradient_accumulation_steps(),
-                aio_config=self.aio_config(),
-                fp32_allreduce=self.allreduce_always_fp32)
+                aio_config=self.aio_config())
 
         else:
             raise NotImplementedError("ZeRO stage {} not implemented".format(zero_stage))

deepspeed/runtime/zero/stage1.py

@@ -121,8 +121,7 @@ def __init__(self,
                  clip_grad=0.0,
                  max_elements_per_comm=5e8,
                  elastic_checkpoint=True,
-                 precision=torch.half,
-                 fp32_allreduce=False):
+                 precision=torch.half):
 
         # Load pre-built or JIT compile (un)flatten ops
         util_ops = UtilsBuilder().load()
@@ -131,7 +130,10 @@ def __init__(self,
 
         # set precision
         self.precision = precision
-        self.fp32_allreduce = fp32_allreduce
+        if self.precision == torch.bfloat16:
+            self.fp32_allreduce = True
+        else:
+            self.fp32_allreduce = False
 
         if dp_process_group is not None and partition_size is not None:
             raise ValueError("Cannot specify both dp_process_group "

deepspeed/runtime/zero/stage2.py

@@ -115,7 +115,7 @@ def __init__(self,
             raise SystemError("Cannot use fp16 without CUDA.")
         self.optimizer = init_optimizer
         self.precision = precision
-        self.fp32_allreduce = allreduce_always_fp32
+        self.fp32_allreduce = True if self.precision == torch.bfloat16 else allreduce_always_fp32
 
 
         # Load pre-built or JIT compile (un)flatten ops