Repository: deepspeedai/DeepSpeed
Pull request: Add ZeRO 1 support to PP for BF16.
Addressing issue #2994 

Problem statement:
[BUG] Does bf16 support Zero stage 1 with pipeline?
**Describe the bug**
I'm using Deepspeed-Megatron although,  using pipeline parallelism and setting 
```
"bf16": {
    "enabled": "auto"
  }
```
will step into the `NotImplementedError` in 
```python
#/opt/conda/lib/python3.8/site-packages/deepspeed/runtime/pipe/engine.py
def _exec_reduce_grads(self):
      self._force_grad_boundary = True
      if self.pipeline_enable_backward_allreduce:
          if self.bfloat16_enabled():
              if self.zero_optimization_stage() == 0:
                  self._bf16_reduce_grads()
              else:
                  assert self.zero_optimization_stage() == 1, "only bf16 + z1 are supported"
                  raise NotImplementedError()
          else:
              self.allreduce_gradients(bucket_size=MEMORY_OPT_ALLREDUCE_SIZE)
      self._force_grad_boundary = False
```
But when using transformer integrated deepspeed with Zero Stage 1/2/3. It work fine
The only diff I found is that, in megatron-deepspeed the model was a subclass of  PipelineModule whereas in transformer is not

I wonder whether deepspeed now support pipeline bf16 with zero stage 1 or just my code mistakes

Hints from issue discussion:
@lyj201002, thanks for reporting this bug. We are working on a fix.

Commit message: Add ZeRO 1 support to PP for BF16.

Patch to review:
diff --git a/deepspeed/runtime/pipe/engine.py b/deepspeed/runtime/pipe/engine.py
index f3370cd0e908..d9d04675b84b 100644
--- a/deepspeed/runtime/pipe/engine.py
+++ b/deepspeed/runtime/pipe/engine.py
@@ -241,11 +241,10 @@ def _exec_reduce_grads(self):
         self._force_grad_boundary = True
         if self.pipeline_enable_backward_allreduce:
             if self.bfloat16_enabled():
-                if self.zero_optimization_stage() == 0:
+                if self.zero_optimization_stage() < 2:
                     self._bf16_reduce_grads()
                 else:
-                    assert self.zero_optimization_stage() == 1, "only bf16 + z1 are supported"
-                    raise NotImplementedError()
+                    raise NotImplementedError("PP+BF16 only work for ZeRO Stage 1")
             else:
                 self.allreduce_gradients(bucket_size=MEMORY_OPT_ALLREDUCE_SIZE)
         self._force_grad_boundary = False


Code review comments:
diff_hunk: @@ -241,11 +241,10 @@ def _exec_reduce_grads(self):
         self._force_grad_boundary = True
         if self.pipeline_enable_backward_allreduce:
             if self.bfloat16_enabled():
-                if self.zero_optimization_stage() == 0:
+                if self.zero_optimization_stage() < 2:
path: deepspeed/runtime/pipe/engine.py
text: @user1:
Replace constant 2 with `ZeroStageEnum.gradients` as in  https://github.com/microsoft/DeepSpeed/blob/f7d71ec15376bc8188e5e8eca9c79177536c6e5d/deepspeed/runtime/zero/config.py#L67