Fix allreduce for BF16 and ZeRO0 (#5170)
This PR fixes an issue with allreducing for ZeRO0 + BF16. (This replaces
#5154)
DeepSpeed uses `BF16_Optimizer` when ZeRO0 and BF16 are enabled. The
optimizer accumulates gradients on FP32 buffer soon after a backward
pass completes. However, DeepSpeed engine performs allreduce on BF16
gradients.
This PR fixes the issue by performing allreduce on the FP32 buffer. It
also eliminates an assertion that prohibits BF16+PP+Z1, which is
actually runnable.
This shows loss curves of the following conditions:
- BF16/Z0,Z1,Z2,Z3/NoPP
- BF16/Z0,Z1/PP(2 stages)
(all used 8GPUs, gradient accumulation step: 4)
---------
Co-authored-by: Logan Adams <114770087+loadams@users.noreply.github.com>
