[Performance] Excessive small reduce operations scheduled in backward pass with ZeRO2

Description
Consider a simple training script like https://gist.github.com/0x6b64/5782424ec7ad79b40d51dca7d4ce2b39

Here is the model definition:

printing model = SimpleMoeModel(
  (fc1): Linear(in_features=1000, out_features=1000, bias=True)
  (fc2): Linear(in_features=1000, out_features=1000, bias=True)
  (moeBertLayer): MoEBertLayer(
    (output): MoE(
      (deepspeed_moe): MOELayer(
        (gate): TopKGate(
          (wg): Linear(in_features=1000, out_features=16, bias=False)
        )
        (experts): Experts(
          (deepspeed_experts): ModuleList(
            (0-3): 4 x Sequential(
              (0): SimpleExpert(
                (fc1): Linear(in_features=1000, out_features=1000, bias=True)
                (fc2): Linear(in_features=1000, out_features=1000, bias=True)
              )
            )
          )
        )
      )
    )
  )
)

This model has 10026000 trainable parameter. Using ZeRO2 with training should partition the gradients in near equal size buffers. Since there are 8 participating ranks, 10026000/8 = 1253250 parameters are expected in each partition.

params_in_ipg_bucket in https://github.com/microsoft/DeepSpeed/blob/master/deepspeed/runtime/zero/stage_1_and_2.py#L917 are essentially all of the parameters which fit in one reduce bucket. During initialization (https://github.com/microsoft/DeepSpeed/blob/master/deepspeed/runtime/zero/stage_1_and_2.py#L739) the split of each parameter is predetermined and that information is stored in a combination of param_to_partition_ids and grad_start_offset variables.

However, after coalescing how the tensor splits are partitioned, rank_and_offsets can be printed here: https://github.com/microsoft/DeepSpeed/blob/master/deepspeed/runtime/zero/stage_1_and_2.py#L973 And we observe the following:

Here the columns are (dst_rank, offset, num elements). 
(0, 1002000.0, 252250.0), 
(0, 6006000.0, 4004000)
(1, 1254250.0, 252250.0), 
(1, 2002000.0, 4004000), 
(2, 1506500.0, 252250.0), 
(3, 1000, 8000.0), 
(3, 1001000.0, 1000), 
(3, 1758750.0, 243250.0), 
(4, 9000.0, 252250.0), 
(5, 261250.0, 252250.0), 
(6, 513500.0, 252250.0), 
(7, 0, 1000), 
(7, 765750.0, 235250.0), 

Why is it the case that we don't see equally segmented shards of size 1253250 amongst the 8or9 ranks, with a total of 8 reduce calls. However, what is observed in the above. The same is True for Stage 3 as well. Doing some empirical testing, we can see that have an equal size splits is schedules lesser number of kernels && also completes faster.

Would be great to get some insights into why the actual sharding is done seems to depend on the order in which the parameters are pushed into the params_in_ipg_bucket here: https://github.com/microsoft/DeepSpeed/blob/master/deepspeed/runtime/zero/stage_1_and_2.py#L856

There is parallel_partitioned_bit16_groups which seems to have the expected equally size (&& padded groups) of gradients. https://github.com/microsoft/DeepSpeed/blob/master/deepspeed/runtime/zero/stage_1_and_2.py#L335

But during the actual reduce call, there is a variable size reduce && on large models (~25B parameters), this short reduces can accumulate a lot a reduce the speed of the backwards pass.

--

What what I observe, the partitioning seems to work as such. If there is one parameter group with 1000 parameters, && 10 partitions, a fractional piece of a gradient can get placed in a partition. Tracing the code from here: https://github.com/microsoft/DeepSpeed/blob/master/deepspeed/runtime/zero/stage_1_and_2.py#L657

For 1.5B parameter model, here are the actual reduce calls - 325 of them, executing serially on a single stream. (Training on 128 data parallel ranks), this could be reduced to 129 calls of size 1500000000/128 = 11718750.0

, 24, 48, 118, 120, 214, 404, 428, 452, 476, 500, 642, 666, 690, 714, 762, 904, 928, 952, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024, 1096, 1120, 1166, 1190, 1762, 2000, 2024, 2048, 2048, 2048, 2048, 2048, 2048, 2048, 2048, 2048, 2048, 2048, 2048, 2048, 2048, 2048, 2048, 2048, 2048, 2048, 2048, 2048, 2238, 3072, 3072, 3072, 3072, 3072, 3072, 3310, 3452, 4096, 4096, 4096, 4096, 4096, 4096, 4096, 4096, 4096, 16352, 76000, 155240, 193600, 214888, 294128, 307600, 332488, 378664, 424840, 484848, 524288, 563728, 577200, 623736, 656440, 716088, 740976, 754448, 833688, 854976, 893336, 939512, 972576, 1018752, 1048576, 1048576, 1048576, 1048576, 1048576, 1048576, 1048576, 1048576, 1048576, 1048576, 1048576, 1048576, 1048576, 1048576, 1048576, 1048576, 1048576, 1048576, 1048576, 1048576, 1048576, 1048576, 1048576, 1048576, 1048576, 1048576, 1048576, 1048576, 1048576, 1048576, 1048576, 1048576, 1048576, 1048576, 1048576, 1048576, 1048576, 1048576, 1048576, 1048576, 1048576, 1048576, 1048576, 1048576, 1048576, 1048576, 1048576, 1048576, 1048576, 1048576, 1048576, 1048576, 1048576, 1048576, 1048576, 1048576, 1048576, 1048576, 1048576, 1217288, 1579600, 1857376, 1974616, 2011176, 2219688, 2336928, 2614704, 2977016, 3175552, 3254792, 3537864, 3617104, 3815640, 4177952, 4194304, 4194304, 4194304, 4194304, 4194304, 4194304, 4194304, 4194304, 4194304, 4194304, 4194304, 4194304, 4194304, 4194304, 4194304, 4194304, 4194304, 4194304, 4198400, 7154968, 7154968, 7154968, 7154968, 7154968, 7154968, 7154968, 7154968, 7154968, 7154968, 7154968, 7154968, 7154968, 7154968, 7154968, 7154968, 7154968, 7154968, 7154968, 7154968, 7154968, 7154968, 7154968, 7154968, 7154968, 7154968, 7154968, 7154968, 7154968, 7154968, 7154968, 7154968, 7154968, 7154968, 7154968, 7154968, 7154968, 7154968, 7154968, 7154968, 7154968, 7154968, 7154968, 7154968, 7154968, 7154968, 7154968, 7154968, 7154968, 7154968, 7154968, 7154968, 7154968, 7154968, 7154968, 7154968, 7154968, 7154968, 7154968, 7154968, 7154968, 7154968, 7154968, 7154968, 7154968, 7154968, 7154968, 7154968, 7154968, 7154968, 7154968, 7154968, 7154968, 16787456, 16787456, 16787456, 16787456, 16787456

To Reproduce

Download the source: https://gist.github.com/0x6b64/5782424ec7ad79b40d51dca7d4ce2b39

mpirun -np 8 --allow-run-as-root python deepspeed_simple_moe_training_zero2.py

torch                         2.0.0
torchaudio                    2.0.0
torchbenchmark                0.0.7
torchdata                     0.6.0
torchtext                     0.15.1
torchtnt                      0.0.7
torchvision                   0.15.0
deepspeed                     0.6.1+06f2048

Any pointers will be greatly appreciated!

Start with the reproduction gist and run the provided mpirun command using the listed DeepSpeed and PyTorch versions. Read deepspeed/runtime/zero/stage_1_and_2.py around the partition initialization, params_in_ipg_bucket, and rank_and_offsets references, then compare the observed reduce calls with parallel_partitioned_bit16_groups. Done means the cause of the uneven reductions is established and a validated improvement is demonstrated without breaking ZeRO2 behavior.