Repository: deepspeedai/DeepSpeed
Pull request: Fix Issue 4083
Fix #4083 

Problem statement:
[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?

Hints from issue discussion:

Commit message: removing bad check

Patch to review:
diff --git a/deepspeed/runtime/engine.py b/deepspeed/runtime/engine.py
index cc25ce69802c..0b2a713921bc 100644
--- a/deepspeed/runtime/engine.py
+++ b/deepspeed/runtime/engine.py
@@ -1137,10 +1137,6 @@ def _do_optimizer_sanity_check(self, basic_optimizer):
 
                 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:
-                return BFLOAT16
             return ZERO_OPTIMIZATION
         elif amp_enabled:
             if model_dtype != grad_accum_dtype:
diff --git a/deepspeed/runtime/zero/stage_1_and_2.py b/deepspeed/runtime/zero/stage_1_and_2.py
index bb578218b93c..df9c7139ef3d 100755
--- a/deepspeed/runtime/zero/stage_1_and_2.py
+++ b/deepspeed/runtime/zero/stage_1_and_2.py
@@ -902,7 +902,7 @@ def reduce_independent_p_g_buckets_and_remove_grads(self, param, i):
             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()


Code review comments:
diff_hunk: @@ -1137,10 +1137,6 @@ def _do_optimizer_sanity_check(self, basic_optimizer):
 
                 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(
path: deepspeed/runtime/engine.py
text: @user1:
This change could break backwards-compatibility for users who expect to use bf16_optimizer in this case.