Repository: deepspeedai/DeepSpeed
Pull request: ZeRO Gradient Accumulation Dtype.
Fix #2352 
Adding configurable gradient accumulation data type to all ZeRO stages.
Added two gradient attributes to the parameter object:
1. `param.grad_reduc`: Used only as a pointer to reference which gradient tensor should be used in reduction. The tensor object pointed to by this attribute is used in all reduction operations.
2. `param.grad_accum`: Used as a pointer or standalone tensor representing the accumulated gradients. The tensor object pointed to by this attribute is used in all accumulation operations.

For stages 1 & 2.
Two modes:
1. `model_dtype == gradient_accumulation_dtype`
2. `model_dtype != gradient_accumulation_dtype`

**Stage 1, Mode 1**:
- `param.grad_reduc = param.grad`
- `param.grad_accum = param.grad`

**Stage 1, Mode 2**:
- `param.grad_reduc = param.accum_grad`
- `param.grad_accum = standalone tensor`

**Stage 2, Mode 1**:
- `param.grad_reduc = param.grad`
- `param.grad_accum = param.grad`

**Stage 2, Mode 2**:
- `param.grad_reduc = param.grad`
- `param.grad_accum = standalone tensor`


Problem statement:
[REQUEST] BF16 mixed precision => grad accum in fp32
**Is your feature request related to a problem? Please describe.**

We have proven with the BLOOM training that BF16 is by far more superior for mixed precision training than FP16, using Megatron-Deepspeed

But the latter is very complex, it'd be much easier for folks to use standalone ZeRO for training in bf16 mixed precision.

But for this to work we need ZeRO to support grad accumulation in fp32, similar to the recently added `BF16Optimizer`.

So this is the feature request to **backport `BF16Optimizer`'s fp32 grad accumulation to ZeRO-1,2,3.**

Once this is done I can safely tell those who are interested in an easier to solution to use ZeRO.

@tjruwase, @jeffra

Hints from issue discussion:
@tjruwase, would it be possible to implement this? We are ready to start using ZeRO-3/bf16 for the multi-modal training.

Thank you very much!
@stas00, is it better to close this or #2768? They are the same thing, right? 
Hi Tunji - you're the owner so it's up to you to decide. The new one is a duplicate of this one, so typically the earliest one stays.

And I don't agree with the other request that it should be hardcoded to fp32, it should be a user choice. Though most likely a sensible default should be fp32 for bf16 mixed precision training.
Makes sense, will close the newer one and reference this appropriately. 

Yes, the accumulation type will be configurable. Hopefully, we should have a WIP pushed later this week. It would be great to get your usual feedback as we iterate on a solution. 
Fantastic news, Tunji. Thank you.

And, yes, we would be happy to experiment with your WIP PR. 
Amazing, thanks!
hi I also met the same problem. @tjruwase have you found a solution?
Great, looking forward to see this new release!

Commit message: Merge branch 'master' into jomayeri/new-zero-accum

Patch to review:
diff --git a/deepspeed/runtime/engine.py b/deepspeed/runtime/engine.py
index 86daf283301b..875490ae4632 100644
--- a/deepspeed/runtime/engine.py
+++ b/deepspeed/runtime/engine.py
@@ -1208,13 +1208,9 @@ def _do_optimizer_sanity_check(self, basic_optimizer):
                         "**** You are using ZeRO with an untested optimizer, proceed with caution *****"
                     )
 
-            if model_dtype == torch.bfloat16 and grad_accum_dtype == torch.float32 and self.zero_optimization_stage(
-            ) == 1:
-                return BFLOAT16
-
-            if model_dtype != grad_accum_dtype:
+            if model_dtype != grad_accum_dtype and self.zero_optimization_stage() == 3:
                 raise NotImplementedError(
-                    "Model data type and gradient accumulation data type must be equal to use ZeRO"
+                    "Model data type and gradient accumulation data type must be equal to use ZeRO stage 3"
                 )
             return ZERO_OPTIMIZATION
         elif amp_enabled:
@@ -1497,7 +1493,7 @@ def _configure_bf16_optimizer(self, optimizer):
 
     def _configure_zero_optimizer(self, optimizer):
         zero_stage = self.zero_optimization_stage()
-        model_dtype, grad_accum_dtype = self.get_data_types()
+        model_dtype, gradient_accumulation_dtype = self.get_data_types()
         timers = self.timers if self.wall_clock_breakdown() else None
 
         if optimizer is None:
@@ -1559,6 +1555,7 @@ def _configure_zero_optimizer(self, optimizer):
                 has_moe_layers=self.has_moe_layers,
                 fp16_master_weights_and_gradients=self.fp16_master_weights_and_gradients(
                 ),
+                gradient_accumulation_dtype=gradient_accumulation_dtype,
                 communication_data_type=self.communication_data_type,
                 elastic_checkpoint=self.zero_elastic_checkpoint())
 
diff --git a/deepspeed/runtime/zero/stage_1_and_2.py b/deepspeed/runtime/zero/stage_1_and_2.py
index 0b9d9c5a6fef..e8c0a58d5c6c 100755
--- a/deepspeed/runtime/zero/stage_1_and_2.py
+++ b/deepspeed/runtime/zero/stage_1_and_2.py
@@ -129,6 +129,7 @@ def __init__(self,
                  cpu_offload=False,
                  mpu=None,
                  clip_grad=0.0,
+                 gradient_accumulation_dtype=torch.float32,
                  communication_data_type=torch.float16,
                  postscale_gradients=True,
                  gradient_predivide_factor=1.0,
@@ -276,6 +277,16 @@ def __init__(self,
 
         self.all_reduce_print = False
         self.dtype = self.optimizer.param_groups[0]['params'][0].dtype
+        self.gradient_accumulation_dtype = gradient_accumulation_dtype
+
+        if self.dtype != self.gradient_accumulation_dtype:
+            self.use_separate_grad_accum = True
+        else:
+            self.use_separate_grad_accum = False
+        if self.use_separate_grad_accum and not self.partition_gradients:
+            self.use_grad_accum_for_reduction = True
+        else:
+            self.use_grad_accum_for_reduction = False
 
         self.round_robin_bit16_groups = []
         self.round_robin_bit16_indices = []
@@ -289,9 +300,12 @@ def __init__(self,
 
             # push this group to list before modify
             # TODO: Explore simplification that avoids the extra book-keeping by pushing the reordered group
-            trainable_parameters = [
-                param for param in param_group['params'] if param.requires_grad
-            ]
+            trainable_parameters = []
+            for param in param_group['params']:
+                if param.requires_grad:
+                    param.grad_accum = None
+                    param.grad_reduc = None
+                    trainable_parameters.append(param)
             self.bit16_groups.append(trainable_parameters)
 
             # not sure why apex was cloning the weights before flattening
@@ -495,6 +509,8 @@ def __init__(self,
         # resets the data structure value for the next backward propagation
         self.reset_partition_gradient_structures()
 
+        # Hooks execute in the order they are registered
+        self.create_set_grad_accum_hooks()
         # creates backward hooks for gradient partitioning
         if self.partition_gradients or self.overlap_comm:
             self.create_reduce_and_remove_grad_hooks()
@@ -672,7 +688,8 @@ def reduce_gradients(self, pipeline_parallel=False):
         if not self.overlap_comm:
             for i, group in enumerate(self.bit16_groups):
                 for param in group:
-                    if param.grad is not None:
+                    self.point_grad_reduc(param)
+                    if param.grad_reduc is not None:
                         self.reduce_ready_partitions_and_remove_grads(param, i)
         # reduce any pending grads in either hook/non-hook case
         self.overlapping_partition_gradients_reduce_epilogue()
@@ -738,7 +755,7 @@ def independent_gradient_partition_epilogue(self):
                         self.params_in_partition[i],
                         self.first_offset[i],
                         self.partition_size[i],
-                        dtype=self.dtype,
+                        dtype=self.gradient_accumulation_dtype,
                         device=get_accelerator().current_device_name(),
                         return_tensor_list=True)
                 else:
@@ -746,7 +763,7 @@ def independent_gradient_partition_epilogue(self):
                         self.params_in_partition[i],
                         self.first_offset[i],
                         self.partition_size[i],
-                        dtype=self.dtype,
+                        dtype=self.gradient_accumulation_dtype,
                         device=get_accelerator().current_device_name(),
                         return_tensor_list=True)
 
@@ -833,6 +850,64 @@ def increment_value(dictionary, key):
     def overlapping_partition_gradients_reduce_epilogue(self):
         self.independent_gradient_partition_epilogue()
 
+    # Set the accum grad pointer to the correct tensor
+    def set_grad_accum(self, param):
+        # Simply point to the native pytorch grad
+        if not self.use_separate_grad_accum:
+            param.grad_accum = param.grad
+        # Accum the gradient in a separate tensor
+        else:
+            if param.grad is not None:
+                if param.grad_accum is None:
+                    param.grad_accum = param.grad.to(self.gradient_accumulation_dtype)
+                else:
+                    param.grad_accum.add_(
+                        param.grad.to(self.gradient_accumulation_dtype).view(
+                            param.grad_accum.shape))
+                # Don't clear if you still need the native grad for reduction (stage 2)
+                if self.use_grad_accum_for_reduction:
+                    param.grad = None
+
+    # Clear the tensor the reduction gradient attribute is pointing to
+    def clear_grad_reduc(self, param):
+        param.grad_reduc = None
+        if self.use_grad_accum_for_reduction:
+            param.grad_accum = None
+        else:
+            param.grad = None
+
+    # Point to the gradient tensor to use in reduction
+    def point_grad_reduc(self, param):
+        if self.use_grad_accum_for_reduction:
+            param.grad_reduc = param.grad_accum
+        else:
+            param.grad_reduc = param.grad
+
+    # Called after a memory defrag moves tensor object via copy
+    def point_grad_to_reduc(self, param):
+        if self.use_grad_accum_for_reduction:
+            param.grad_accum = param.grad_reduc
+        else:
+            param.grad = param.grad_reduc
+
+    def create_set_grad_accum_hooks(self):
+        self.accum_accs = []
+        for param_group in self.bit16_groups:
+            for param in param_group:
+                if param.requires_grad:
+
+                    def wrapper(param):
+                        param_tmp = param.expand_as(param)
+                        grad_acc = param_tmp.grad_fn.next_functions[0][0]
+
+                        def accumulation_grad_setter(*notneeded):
+                            self.set_grad_accum(param)
+
+                        grad_acc.register_hook(accumulation_grad_setter)
+                        self.accum_accs.append(grad_acc)
+
+                    wrapper(param)
+
     def create_reduce_and_remove_grad_hooks(self):
         self.grad_accs = []
         for i, param_group in enumerate(self.bit16_groups):
@@ -868,6 +943,9 @@ def flatten_dense_tensors_aligned(self, tensor_list, alignment):
 
     ############### Independent Partition Gradient ########################
     def reduce_independent_p_g_buckets_and_remove_grads(self, param, i):
+        if param.grad_reduc is None:
+            self.point_grad_reduc(param)
+
         if self.elements_in_ipg_bucket + param.numel() > self.reduce_bucket_size:
             self.report_ipg_memory_usage("In ipg_remove_grads before reduce_ipg_grads",
                                          param.numel())
@@ -893,14 +971,15 @@ def reduce_independent_p_g_buckets_and_remove_grads(self, param, i):
                 0,
                 self.elements_in_ipg_bucket,
                 param.numel())
-            new_grad_tensor.copy_(param.grad.view(-1))
-            param.grad.data = new_grad_tensor.data.view_as(param.grad)
+            new_grad_tensor.copy_(param.grad_reduc.view(-1))
+            param.grad_reduc.data = new_grad_tensor.data.view_as(param.grad_reduc)
+            self.point_grad_to_reduc(param)
 
         self.elements_in_ipg_bucket += param.numel()
 
-        assert param.grad is not None, f"rank {dist.get_rank()} - Invalid to reduce Param {param_id} with None gradient"
+        assert param.grad_reduc is not None, f"rank {dist.get_rank()} - Invalid to reduce Param {param_id} with None gradient"
 
-        self.grads_in_ipg_bucket.append(param.grad)
+        self.grads_in_ipg_bucket.append(param.grad_reduc)
         self.params_in_ipg_bucket.append((i, param, param_id))
 
         #make sure the average tensor function knows how to average the gradients
@@ -972,7 +1051,7 @@ def average_tensor(self, tensor):
                     process_group = self.expert_dp_process_group[
                         param.group_name] if is_moe_param(
                             param) else self.dp_process_group
-                    param.grad.data.div_(dist.get_world_size(group=process_group))
+                    param.grad_reduc.data.div_(dist.get_world_size(group=process_group))
 
                 partition_ids = self.param_to_partition_ids[i][param_id]
                 assert all([p_id < dist.get_world_size(group=process_group) for p_id in partition_ids]), f"world size {dist.get_world_size(group=process_group)} and p_ids: {partition_ids}"
@@ -1068,7 +1147,7 @@ def get_grad_position(self, group_id, tensor_list, first_offset, partition_size)
             current_offset += num_elements
 
     def update_overflow_tracker_for_param_grad(self, param):
-        if param.grad is not None and self._has_inf_or_nan(param.grad.data):
+        if param.grad_accum is not None and self._has_inf_or_nan(param.grad_accum.data):
             self.local_overflow = True
 
     def async_accumulate_grad_in_cpu_via_gpu(self, param):
@@ -1095,19 +1174,19 @@ def buffer_to_accumulate_to_in_cpu():
                     dest_offset,
                     num_elements)
 
-        #accumulate gradients into param.grad or parts of it that belongs to this partition
+        #accumulate gradients into param.grad_accum or parts of it that belongs to this partition
         def accumulate_gradients():
             if not self.fp16_master_weights_and_gradients:
                 dest_buffer.copy_(self.accumulated_grads_in_cpu[param_id].view(-1),
                                   non_blocking=True)
-                param.grad.data.view(-1).add_(dest_buffer)
+                param.grad_accum.data.view(-1).add_(dest_buffer)
             else:
                 dest_buffer.narrow(0,
                                    source_offset,
                                    num_elements).copy_(
                                        self.accumulated_grads_in_cpu[param_id].view(-1),
                                        non_blocking=True)
-                param.grad.data.view(-1).narrow(
+                param.grad_accum.data.view(-1).narrow(
                     0,
                     source_offset,
                     num_elements).add_(dest_buffer.narrow(0,
@@ -1118,13 +1197,13 @@ def accumulate_gradients():
         def copy_gradients_to_cpu():
             if not self.fp16_master_weights_and_gradients:
                 self.accumulated_grads_in_cpu[param_id].data.copy_(
-                    param.grad.data.view(-1),
+                    param.grad_accum.data.view(-1),
                     non_blocking=True)
             else:
                 self.accumulated_grads_in_cpu[param_id].data.copy_(
-                    param.grad.data.view(-1).narrow(0,
-                                                    source_offset,
-                                                    num_elements),
+                    param.grad_accum.data.view(-1).narrow(0,
+                                                          source_offset,
+                                                          num_elements),
                     non_blocking=True)
 
         if param_id not in self.accumulated_grads_in_cpu:
@@ -1140,7 +1219,7 @@ def copy_gradients_to_cpu():
     def set_norm_for_param_grad(self, param):
         param_id = self.get_param_id(param)
         accumulated_grad = self.accumulated_grads_in_cpu[
-            param_id] if self.gradient_accumulation_steps > 1 else param.grad
+            param_id] if self.gradient_accumulation_steps > 1 else param.grad_accum
 
         [i, source_offset, dest_offset, num_elements] = self.grad_position[param_id]
 
@@ -1151,7 +1230,7 @@ def set_norm_for_param_grad(self, param):
 
     def set_norm_for_param_grad_in_gpu(self, param):
         param_id = self.get_param_id(param)
-        accumulated_grad = param.grad
+        accumulated_grad = param.grad_accum
 
         [i, source_offset, dest_offset, num_elements] = self.grad_position[param_id]
 
@@ -1170,12 +1249,14 @@ def async_inplace_copy_grad_to_fp32_buffer_from_gpu(self, param):
             dest_offset,
             num_elements)
 
-        src_tensor = param.grad.view(-1).narrow(0, source_offset, num_elements)
+        src_tensor = param.grad_accum.view(-1).narrow(0, source_offset, num_elements)
         if not self.fp16_master_weights_and_gradients:
             src_tensor = src_tensor.float()
 
         dest_tensor.copy_(src_tensor, non_blocking=True)
         param.grad = None  #offload only
+        param.grad_accum = None  #offload only
+        param.grad_reduc = None  #offload only
 
     def complete_grad_norm_calculation_for_cpu_offload(self, params):
         total_norm = 0.0
@@ -1257,8 +1338,9 @@ def copy_grads_in_partition(self, param):
             0,
             self.grads_in_partition_offset,
             param.numel())
-        new_grad_tensor.copy_(param.grad.view(-1))
-        param.grad.data = new_grad_tensor.data.view_as(param.grad)
+        new_grad_tensor.copy_(param.grad_reduc.view(-1))
+        param.grad_reduc.data = new_grad_tensor.data.view_as(param.grad_reduc)
+        self.point_grad_to_reduc(param)
         #print(f"Grad norm after copy to contiguous_buffer {param.grad.data.norm()}")
         self.grads_in_partition_offset += param.numel()
 
@@ -1308,7 +1390,7 @@ def reduce_ipg_grads(self):
                                 self.previous_reduced_grads = []
                             self.previous_reduced_grads.append(param)
                         else:
-                            param.grad = None  #only if self.partition_gradients
+                            self.clear_grad_reduc(param)
                     elif self.contiguous_gradients:
                         self.copy_grads_in_partition(param)
                 else:  # zero stage 1 - partition only optimizer state
@@ -1425,7 +1507,7 @@ def allreduce_bucket(self, bucket, rank=None, log=None):
     def _clear_previous_reduced_grads(self):
         if self.previous_reduced_grads is not None:
             for param in self.previous_reduced_grads:
-                param.grad = None  # overlap enabled
+                self.clear_grad_reduc(param)
             self.previous_reduced_grads = None
 
     # if rank is specified do a reduction instead of an allreduce
@@ -1539,14 +1621,23 @@ def zero_grad(self, set_to_none=False):
         """
         # FP32 grad should never exist.
         # For speed, set model fp16 grad to None by default
+        # zero all pointers to grad tensors
         for group in self.bit16_groups:
             for p in group:
                 if set_to_none:
                     p.grad = None  # epilogue and in step
+                    p.grad_accum = None
+                    p.grad_reduc = None
                 else:
                     if p.grad is not None:
                         p.grad.detach_()
                         p.grad.zero_()
+                    if p.grad_accum is not None:
+                        p.grad_accum.detach_()
+                        p.grad_accum.zero_()
+                    if p.grad_reduc is not None:
+                        p.grad_reduc.detach_()
+                        p.grad_reduc.zero_()
 
     def _model_parallel_all_reduce(self, tensor, op):
         """ Perform all reduce within model parallel group, if any.
@@ -1623,11 +1714,19 @@ def get_flat_partition(self,
                            return_tensor_list=False):
         flat_tensor_list = []
         current_size = 0
+        add_reduc_to_accum = self.use_separate_grad_accum and not self.use_grad_accum_for_reduction
+
         for i, tensor in enumerate(tensor_list):
-            if tensor.grad is None:
-                tensor.grad = torch.zeros_like(tensor)
+            if tensor.grad_accum is None:
+                tensor.grad_accum = torch.zeros_like(tensor, dtype=dtype)
+
+            if add_reduc_to_accum:
+                if tensor.grad_reduc is not None:
+                    tensor.grad_accum.add_(
+                        tensor.grad_reduc.to(self.gradient_accumulation_dtype).view(
+                            tensor.grad_accum.shape))
 
-            tensor = tensor.grad
+            tensor = tensor.grad_accum
             num_elements = tensor.numel()
             tensor_offset = 0
 
@@ -1667,6 +1766,8 @@ def get_flat_partition(self,
     def free_grad_in_param_list(self, param_list):
         for p in param_list:
             p.grad = None  # in step
+            p.grad_reduc = None
+            p.grad_accum = None
 
     def reset_cpu_buffers(self):
         self.norm_for_param_grads = {}
diff --git a/tests/unit/runtime/test_ds_initialize.py b/tests/unit/runtime/test_ds_initialize.py
index 5ca3469ad6c6..8c718c54f42e 100644
--- a/tests/unit/runtime/test_ds_initialize.py
+++ b/tests/unit/runtime/test_ds_initialize.py
@@ -116,7 +116,7 @@ def test(self, client_parameters):
         assert isinstance(ds_optimizer, FusedAdam)
 
 
-@pytest.mark.parametrize('optimizer_extension', ['zero1', 'zero2', 'amp', None])
+@pytest.mark.parametrize('optimizer_extension', ['zero1', 'zero2', 'zero3', 'amp', None])
 @pytest.mark.parametrize('model_dtype', ['fp16', 'bf16', 'fp32'])
 @pytest.mark.parametrize('grad_accum_dtype', [None, 'fp16', 'bf16', 'fp32'])
 class TestOptimizerImplementation(DistributedTest):
@@ -127,6 +127,8 @@ def test(self, optimizer_extension, model_dtype, grad_accum_dtype):
             zero_stage = 1
         elif optimizer_extension == 'zero2':
             zero_stage = 2
+        elif optimizer_extension == 'zero3':
+            zero_stage = 3
         else:
             zero_stage = 0
         amp = True if optimizer_extension == 'amp' else False
@@ -172,18 +174,36 @@ def test(self, optimizer_extension, model_dtype, grad_accum_dtype):
         # ZeRO 1 Wrapper
         is_supported[('zero1', 'fp16', None)] = True
         is_supported[('zero1', 'fp16', 'fp16')] = True
+        is_supported[('zero1', 'fp16', 'bf16')] = True
+        is_supported[('zero1', 'fp16', 'fp32')] = True
         is_supported[('zero1', 'bf16', None)] = True
+        is_supported[('zero1', 'bf16', 'fp16')] = True
         is_supported[('zero1', 'bf16', 'bf16')] = True
         is_supported[('zero1', 'bf16', 'fp32')] = True
         is_supported[('zero1', 'fp32', None)] = True
+        is_supported[('zero1', 'fp32', 'fp16')] = True
+        is_supported[('zero1', 'fp32', 'bf16')] = True
         is_supported[('zero1', 'fp32', 'fp32')] = True
         # ZeRO 2 Wrapper
         is_supported[('zero2', 'fp16', None)] = True
         is_supported[('zero2', 'fp16', 'fp16')] = True
+        is_supported[('zero2', 'fp16', 'bf16')] = True
+        is_supported[('zero2', 'fp16', 'fp32')] = True
         is_supported[('zero2', 'bf16', None)] = True
+        is_supported[('zero2', 'bf16', 'fp16')] = True
         is_supported[('zero2', 'bf16', 'bf16')] = True
+        is_supported[('zero2', 'bf16', 'fp32')] = True
         is_supported[('zero2', 'fp32', None)] = True
+        is_supported[('zero2', 'fp32', 'fp16')] = True
+        is_supported[('zero2', 'fp32', 'bf16')] = True
         is_supported[('zero2', 'fp32', 'fp32')] = True
+        # ZeRO 3 Wrapper
+        is_supported[('zero3', 'fp16', None)] = True
+        is_supported[('zero3', 'fp16', 'fp16')] = True
+        is_supported[('zero3', 'bf16', None)] = True
+        is_supported[('zero3', 'bf16', 'bf16')] = True
+        is_supported[('zero3', 'fp32', None)] = True
+        is_supported[('zero3', 'fp32', 'fp32')] = True
         # Amp Wrapper
         is_supported[('amp', 'fp32', None)] = True
         is_supported[('amp', 'fp32', 'fp32')] = True


Code review comments:
diff_hunk: @@ -833,6 +850,64 @@ def increment_value(dictionary, key):
     def overlapping_partition_gradients_reduce_epilogue(self):
         self.independent_gradient_partition_epilogue()
 
+    # Set the accum grad pointer to the correct tensor
+    def set_grad_accum(self, param):
+        # Simply point to the native pytorch grad
+        if not self.use_separate_grad_accum:
+            param.grad_accum = param.grad
+        # Accum the gradient in a separate tensor
+        else:
+            if param.grad is not None:
+                if param.grad_accum is None:
+                    param.grad_accum = param.grad.to(self.gradient_accumulation_dtype)
+                else:
+                    param.grad_accum.add_(
+                        param.grad.to(self.gradient_accumulation_dtype).view(
+                            param.grad_accum.shape))
+                # Don't clear if you still need the native grad for reduction (stage 2)
+                if self.use_grad_accum_for_reduction:
+                    param.grad = None
+
+    # Clear the tensor the reduction gradient attribute is pointing to
+    def clear_grad_reduc(self, param):
+        param.grad_reduc = None
+        if self.use_grad_accum_for_reduction:
+            param.grad_accum = None
+        else:
+            param.grad = None
+
+    # Point to the gradient tensor to use in reduction
+    def point_grad_reduc(self, param):
+        if self.use_grad_accum_for_reduction:
+            param.grad_reduc = param.grad_accum
+        else:
+            param.grad_reduc = param.grad
+
+    # Called after a memory defrag moves tensor object via copy
+    def point_grad_to_reduc(self, param):
+        if self.use_grad_accum_for_reduction:
+            param.grad_accum = param.grad_reduc
+        else:
+            param.grad = param.grad_reduc
+
+    def create_set_grad_accum_hooks(self):
path: deepspeed/runtime/zero/stage_1_and_2.py
text: @user1:
Hooks can incur runtime overheads that could be non-trivial for zero stage 1, so I would like to avoid them. Couldn't we call `set_grad_accum` for all params in the optimizer constructor?

@author:
It could be replaced by nested for loops in the backward pass.

@user1:
Let's chat about this.

diff_hunk: @@ -1208,13 +1208,9 @@ def _do_optimizer_sanity_check(self, basic_optimizer):
                         "**** You are using ZeRO with an untested optimizer, proceed with caution *****"
                     )
 
-            if model_dtype == torch.bfloat16 and grad_accum_dtype == torch.float32 and self.zero_optimization_stage(
-            ) == 1:
-                return BFLOAT16
-
-            if model_dtype != grad_accum_dtype:
+            if model_dtype != grad_accum_dtype and self.zero_optimization_stage() == 3:
path: deepspeed/runtime/engine.py
text: @user1:
This change makes bf16_optimizer unusable, right?

@author:
Correct, I believe it replaces all use cases of the current BF16 Optimizer.

@user1:
@user2, what do you think? We probably need more testing to validate correctness parity before disabling bf16_optimizer. 

Also, bf16_optimizer has superior memory efficiency, and probably throughput. We need to validate that as well. 

@user2:
I'm not sure what you are asking, Tunji. 

We know for sure that some people out there use Megatron-Deepspeeed (as we get Issues), so it's almost certain someone will want the BF16Optimizer in a working condition. If that's the question.



@user1:
Yes, that is the question. How to ensure BF16Optimizer users retain access?

@user2:
write code in such a way so it's usable?

why is there a problem in the first place? 



@user1:
Sorry for not being clear. Let me back up a bit. With this PR, both zero_stage_1 now supports `model_dtype=bf16` and `accum_dtype=fp32` similar to bf16_optimizer.  And so, with that dtypes combination and `stage=1`, zero_stage_1 optimizer will now be instantiated instead of bf16_optimizer. I don't think this is desirable, at least not until we validated the correctness and performance equivalence. So, @author, can you please restore the old lines 1211--1213 and place it above the new lines 1211--1213. In other words, we need both logics. We can chat more for clarifications. Thanks!

@user2:
I thought this was the case, thank you for elaborating, Tunji. 

As I was suggesting a year ago we needed some specific flag in the config that activates `BF16Optimizer`. It keeps on biting us the way it is now.

Schedule a call for Monday? Please ping me on slack when it's good for you