make fp32 allreduce optional

Author: sdtblck

deepspeed/runtime/engine.py

@@ -803,7 +803,6 @@ 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:
@@ -819,7 +818,8 @@ 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())
+                precision=self.precision(),
+                fp32_allreduce=self.allreduce_always_fp32)
         elif zero_stage == ZERO_OPTIMIZATION_GRADIENTS:
             optimizer = FP16_DeepSpeedZeroOptimizer(
                 optimizer,
@@ -839,7 +839,8 @@ 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())
+                precision=self.precision(),
+                fp32_allreduce=self.allreduce_always_fp32)
         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
@@ -867,7 +868,8 @@ 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())
+                aio_config=self.aio_config(),
+                fp32_allreduce=self.allreduce_always_fp32)
 
         else:
             raise NotImplementedError("ZeRO stage {} not implemented".format(zero_stage))

deepspeed/runtime/zero/stage1.py

@@ -121,7 +121,8 @@ def __init__(self,
                  clip_grad=0.0,
                  max_elements_per_comm=5e8,
                  elastic_checkpoint=True,
-                 precision=torch.half):
+                 precision=torch.half,
+                 fp32_allreduce=False):
 
         # Load pre-built or JIT compile (un)flatten ops
         util_ops = UtilsBuilder().load()
@@ -130,10 +131,7 @@ def __init__(self,
 
         # set precision
         self.precision = precision
-        if self.precision == torch.bfloat16:
-            self.fp32_allreduce = True
-        else:
-            self.fp32_allreduce = False
+        self.fp32_allreduce = fp32_allreduce
 
         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 = True if self.precision == torch.bfloat16 else allreduce_always_fp32
+        self.fp32_allreduce = allreduce_always_fp32
 
 
         # Load pre-built or JIT compile (un)flatten ops