DeepSpeed

0178e6cc - Fix unbalanced gradients bug in ZeRO-2 gradient accumulation (#545)

Fix unbalanced gradients bug in ZeRO-2 gradient accumulation (#545) * Use zero-tensors for missing gradients to avoid size mismatch * Unit test for unbalanced gradients in ZeRO * Formatting fixes