ZeRO-1 tune max-elems + bug fix (microsoft#532)

* zero-1 memory fix

* auto-tune max elems per comm to reduce padding/comm intervals

* clean-up and added previously missing reduction options

* fix testing backing to work with torch1.7

deepspeed/runtime/engine.py

@@ -661,7 +661,7 @@ def _configure_fp16_optimizer(self, optimizer):
     def _configure_zero_optimizer(self, optimizer):
         zero_stage = self.zero_optimization_stage()
         logger.info('Creating fp16 ZeRO stage {} optimizer'.format(zero_stage))
-
+        assert not self.allreduce_always_fp32(), "ZeRO does not support 'fp32_allreduce': true"
         if zero_stage == ZERO_OPTIMIZATION_OPTIMIZER_STATES:
             assert self.zero_reduce_scatter(), 'Stage 1 only supports reduce scatter mode'
             optimizer = FP16_DeepSpeedZeroOptimizer_Stage1(

deepspeed/runtime/utils.py

@@ -68,12 +68,6 @@ def check_using_norm(self, norm_group, reduce_overflow=True):
         return bool(overflow)
 
     def check(self, param_groups=None):
-
-        #TODO: what's the equivalent here? do we need this?
-        # for group in self.fp32_from_fp32_groups:
-        #     for param in group:
-        #         params.append(param)
-
         params = []
         if param_groups is None:
             params = self.params