Why does not the stage 3 implementation use broadcast and reduce-scatter?

[zarzen]

Hi there,
(Thanks for sharing your implementations!)

After reading the paper, and partially going throught the stage3.py implementation.
I have two questions in general:

1. is there any specific reason for picking all-gather instead of broadcast operation, which is mentioned in the paper.
2. by profiling the stage-3, using Nsight System, I didn't see any nccl reduce-scatter kernel on the timeline. As shown in the following. While, I do see the nccl-all-reduce kernel on the timeline, does that mean the ZeRO3 implementation using all-reduce for gradient reduction?

[samyam]

@zarzen Thank you for trying out ZeRO-3.

1. The communication collectives in the ZeRO-3 implementation is slightly different from what is described in the ZeRO paper. We are using all-gather instead of broadcast as you noticed. The reason for doing so is to effectively support offloading to CPU and NVMe. In the absence of offloading, both all-gather and broadcast would result in similar communication volume, but with offloading, all-gather supports a data partitioning scheme that enables better offload bandwidth utilization. For details, please take a look at Sec 6.1 in our ZeRO-Infinity paper : https://arxiv.org/pdf/2104.07857.pdf

2. ZeRO3 does use reduce-scatter for the gradients. It is strange that its not showing up in the profiler. Perhaps reduce-scatter is turned off in the config. You can trace the code path leading to the reduce scatter call from

params_to_reduce[0].reduce_gradients_at_owner( param_list=params_to_reduce, hierarchy=self.param_coordinator.hierarchy)

in stage3.py to the _reduce_scatter_gradient call in partition_parameters.py

Can you let us know if your execution is taking this code path? If not, please share how we can repro.

[zarzen]

@samyam
I found that the config must have "contiguous_gradients": true to enable reduce_scatter.
Using "reduce_scatter": true, does not enable reduce_scatter.

here is a condition check for invoking average_tensor: https://github.com/microsoft/DeepSpeed/blob/b3870363e026717752c119b21870e2886b9ad92d/deepspeed/runtime/zero/stage3.py#L2175

I am confused about this condition check, as the document for "contiguous_gradients" is related to memory fragmentation.

[samyam]

@zarzen ZeRO-3 should be run with contiguous gradients set to True. Just pushed a PR that makes contiguous gradients True by default for all cases. Originally when we added support for contiguous gradients there was a performance bug that showed up on small model sizes so we set to False by default but that bug has been fixed for a while now so there is no reason to not run with contiguous gradients set to True.

We always assumed that contiguous_gradients will be set to True for stage 3 and so only implemented the reduce_scatter under it. We should probably just deprecate the contiguous_gradients flags in stage 3 and just always have it on as there are no downsides to it.