torch.cuda.synchronize() takes a lot of time for the 10B model training with Zero 2

I recently run some experiments for the 10B model training with Zero 2 on 128 A100 GPUs with 400Gb/s bandwidth, and I observed that communication dominates the total training time, which is about 18 times more expensive than the backward computation. I managed to use nsight system to do some profiling, and the following picture shows the profiling result:

It can be seen that the torch.cuda.synchronize() in
https://github.com/microsoft/DeepSpeed/blob/a4763f5516c0a9bb8f32d2e8a48c618f0fe88e37/deepspeed/runtime/zero/stage2.py#L1084 is the most expensive one and is called multiple times due to the bucketing. I wonder if there is any potential problem if we remove this line? And if we can adopt multiple CUDA streams as in APEX to parallelize the allreduce calls?

Start by reviewing deepspeed/runtime/zero/stage2.py at the referenced line and the surrounding bucketing and allreduce flow. Use the reported 10B-model, Zero 2, 128-A100 profiling context to investigate whether removing torch.cuda.synchronize() is safe and how multiple CUDA streams could affect the calls. Done means the synchronization and stream behavior are evaluated without regressing distributed training correctness or performance.