Repository: microsoft/deepspeed
Issue: [BUG] BF16 training with FP32 gradient accumulation - Zero Stage 1 CPU - DeepSpeedCPUAdam

I am trying to train a model with bf16,  fp32 gradient accumulation and zero stage 1. Using DeepSpeedCPUAdam as optimizer. 


**Config**
```json
{
    "steps_per_print":10000,
    "wall_clock_breakdown":true,
    "train_batch_size": 1024,
    "gradient_accumulation_steps": 16,
    "train_micro_batch_size_per_gpu": 2,
    "fp32_allreduce": true,
    "data_types": {
        "grad_accum_dtype":"fp32"
    },
    "bf16": {
        "enabled": true
    },
    "optimizer": {
        "type": "DeepSpeedCPUAdam",
        "params": {
            "lr": 0.001
        }
    },
    "gradient_clipping": 1.0,
    "scheduler": {
        "type": "WarmupLR",
        "params": {
            "warmup_min_lr": 0,
            "warmup_max_lr": 0.001,
            "warmup_num_steps": 1000
        }
    },

    "zero_optimization": {
        "stage": 1,
        "offload_optimizer": {
            "device": "cpu",
            "pin_memory": true
          },
      "allgather_partitions": true,
      "reduce_bucket_size": 5e8,
      "allgather_bucket_size": 5e8,
      "reduce_scatter": true,
      "overlap_comm": true,
      "round_robin_gradients":true,
      "contiguous_gradients": true,
      "sub_group_size": 1e9
    },
  }
```


**ds_report output**
```[2023-08-03 03:07:32,395] [INFO] [real_accelerator.py:133:get_accelerator] Setting ds_accelerator to cuda (auto detect)
--------------------------------------------------
DeepSpeed C++/CUDA extension op report
--------------------------------------------------
NOTE: Ops not installed will be just-in-time (JIT) compiled at
      runtime if needed. Op compatibility means that your system
      meet the required dependencies to JIT install the op.
--------------------------------------------------
JIT compiled ops requires ninja
ninja .................. [OKAY]
--------------------------------------------------
op name ................ installed .. compatible
--------------------------------------------------
 [WARNING]  async_io requires the dev libaio .so object and headers but these were not found.
 [WARNING]  async_io: please install the libaio-dev package with apt
 [WARNING]  If libaio is already installed (perhaps from source), try setting the CFLAGS and LDFLAGS environment variables to where it can be found.
async_io ............... [NO] ....... [NO]
cpu_adagrad ............ [YES] ...... [OKAY]
cpu_adam ............... [YES] ...... [OKAY]
fused_adam ............. [YES] ...... [OKAY]
fused_lamb ............. [YES] ...... [OKAY]
quantizer .............. [YES] ...... [OKAY]
random_ltd ............. [YES] ...... [OKAY]
 [WARNING]  sparse_attn requires a torch version >= 1.5 and < 2.0 but detected 2.1
 [WARNING]  using untested triton version (2.1.0+9e3e10c5ed), only 1.0.0 is known to be compatible
sparse_attn ............ [NO] ....... [NO]
spatial_inference ...... [YES] ...... [OKAY]
transformer ............ [YES] ...... [OKAY]
stochastic_transformer . [YES] ...... [OKAY]
transformer_inference .. [YES] ...... [OKAY]
--------------------------------------------------
DeepSpeed general environment info:
torch install path ............... ['/work2/conda/envs/env/lib/python3.11/site-packages/torch']
torch version .................... 2.1.0.dev20230725+cu121
deepspeed install path ........... ['/work2/conda/envs/env/lib/python3.11/site-packages/deepspeed']
deepspeed info ................... 0.10.0+unknown, unknown, unknown
torch cuda version ............... 12.1
torch hip version ................ None
nvcc version ..................... 12.0
deepspeed wheel compiled w. ...... torch 2.1, cuda 12.1
```

**Error message**
CPUAdam param is on cuda:7 and must be 'cpu', make sure you enabled 'offload_optimizer': 'cpu' in your ZeRO config.

But offload_optimizer is already set in our config. 
Same config works if I change the grad_accum_dtype to "bf16" with zero code changes.

```Traceback (most recent call last):
    self._configure_optimizer(optimizer, model_parameters)
           ^^^^^^^^^^^^^^^^^^
  File "/work2/dragoon_deepspeed/deepspeed_trainer.py", line 320, in <module>
  File "/work2/conda/envs/env/lib/python3.11/site-packages/deepspeed/runtime/engine.py", line 1220, in _configure_optimizer
AssertionError: CPUAdam param is on cuda:5 and must be 'cpu', make sure you enabled 'offload_optimizer': 'cpu' in your ZeRO config.```


Can someone help by letting me know how to solve this? 


Changed files:
deepspeed/runtime/engine.py
deepspeed/runtime/zero/stage_1_and_2.py

diff --git a/deepspeed/runtime/engine.py b/deepspeed/runtime/engine.py
index cc25ce69..5e5f9728 100644
--- a/deepspeed/runtime/engine.py
+++ b/deepspeed/runtime/engine.py
@@ -1137,9 +1137,8 @@ class DeepSpeedEngine(Module):
 
                 if self.global_rank == 0:
                     logger.warning("**** 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:
+            ) == 1 and not self.zero_cpu_offload():
                 return BFLOAT16
             return ZERO_OPTIMIZATION
         elif amp_enabled:
diff --git a/deepspeed/runtime/zero/stage_1_and_2.py b/deepspeed/runtime/zero/stage_1_and_2.py
index bb578218..e9e8fc16 100755
--- a/deepspeed/runtime/zero/stage_1_and_2.py
+++ b/deepspeed/runtime/zero/stage_1_and_2.py
@@ -836,7 +836,10 @@ class DeepSpeedZeroOptimizer(ZeROOptimizer):
                 param.grad_accum = param.grad
 
     def get_gradient_for_reduction(self, param):
-        return param.grad_accum.to(self.dtype) if self.use_grad_accum_for_reduction else param.grad
+        if self.use_grad_accum_for_reduction:
+            return param.grad_accum.to(self.dtype) if param.grad_accum is not None else None
+        else:
+            return param.grad
 
     # Clear the tensor the reduction gradient attribute is pointing to
     def clear_grad_reduc_pointer(self, param):
@@ -902,7 +905,7 @@ class DeepSpeedZeroOptimizer(ZeROOptimizer):
             else:
                 # keeping the gradients contiguous to prevent memory fragmentation, and avoid flattening
                 new_grad_tensor = self.ipg_buffer[self.ipg_index].narrow(0, self.elements_in_ipg_bucket, param.numel())
-                new_grad_tensor.copy_(param.grad.view(-1))
+                new_grad_tensor.copy_(grad_reduc.view(-1))
                 grad_reduc.data = new_grad_tensor.data.view_as(grad_reduc)
 
         self.elements_in_ipg_bucket += param.numel()
@@ -1262,7 +1265,8 @@ class DeepSpeedZeroOptimizer(ZeROOptimizer):
                 _, _, param_id = self.params_in_ipg_bucket[0]
                 assert self.get_param_id(self.extra_large_param_to_reduce
                                          ) == param_id, "param in ipg bucket does not match extra-large param"
-                self.average_tensor(self.extra_large_param_to_reduce.grad.view(-1))
+                extra_large_grad_reduc = self.get_gradient_for_reduction(self.extra_large_param_to_reduce)
+                self.average_tensor(extra_large_grad_reduc.view(-1))
                 self.extra_large_param_to_reduce = None
             else:
                 self.average_tensor(self.ipg_buffer[self.ipg_index])
@@ -1520,7 +1524,6 @@ class DeepSpeedZeroOptimizer(ZeROOptimizer):
                 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_()
@@ -1629,7 +1632,6 @@ class DeepSpeedZeroOptimizer(ZeROOptimizer):
     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):
@@ -1941,7 +1943,7 @@ class DeepSpeedZeroOptimizer(ZeROOptimizer):
             self.loss_scaler.backward(loss.float(), retain_graph=retain_graph)
 
         # Only for Stage 1, Mode 2
-        if self.use_separate_grad_accum and not self.partition_gradients:
+        if self.use_grad_accum_for_reduction:
             self.update_separate_grad_accum()
         else:
             self.set_grad_accum_pointer()