Are gradients of persistent_parameters synced in backward with reduce-scatter?

[zarzen]

Hi there,
I noticed there is a all_gather step in _post_step function of stage3. The all-gather is used, instead of all-reduce, is it because the gradients of the persistent parameter is synchronized via reduce-scatter in the backward pass?
https://github.com/microsoft/DeepSpeed/blob/master/deepspeed/runtime/zero/stage3.py#L1707

[youth123]

I think this might be a bug. The persistent param should not participate in partition when it completes forward and backward, and it should do allreduce to the param itself after the whole step is complete.
Sorry, I was wrong. In the complete training phase, the param marked with ds_persist will not be released in the forward and backward processes, and then after the reduce scatter communication of grad, there is no need to save the complete weight on a single card, so it will be partitioned in the step, and then all the persistent weights will be allgathered in the post step.