bug about Zero-optimizer

[gbxu]

In Section 6 of paper Rajbhandari, Samyam, et al. "ZeRO: Memory Optimization Towards Training A Trillion Parameter Models.", authors explain that there is no more communication overhead for P_os and P_g.

I think the right algorithm is that for each weight, Zero-Adam will read a tensor after reduce-scatter op, then output the new 1/n part of this weight with 1/n part of optimizer states. This 1/n-size output of Zero-Adam will be passed to all_gather() or broadcast(), syncing with other processes.

However in deepspeed/pt/deepspeed_light.py#L666, the gradients will do an all_reduce before passed to optimizer. Is it right?

[jeffra]

Hi @gbxu, thanks for your interest in DeepSpeed/ZeRO. Good catch regarding all_reduce vs. reduce_scatter. When we first implemented ZeRO it didn't appear reduce_scatter was implemented in torch.distributed so did our initial implementation using all_reduce. We later discovered that reduce_scatter is indeed implemented in torch.distributed but not documented here. We currently have an updated implementation that uses reduce_scatter and has reduced communication time compared with the current implementation. We will be releasing it here soon (along with other updated performance features) after we finish testing/performance evaluation.

Keep an eye on our news site and watch our github repo for updates: https://www.deepspeed.ai/news/