How to make allreduce fully-overlapped in ZeRO-2

[TARGET]
I am trying to overlap inter-worker communications as much as possible, in order to boost the performance of ZeRO-2.

[ISSUE]
According to my understanding of ZeRO-2, reduce-scatter communications can be overlapped by backward computations, which is supported by utilizing different CUDA Streams in an asynchronized way.

F. I. Y.
related parameters in DS_CONFIG
related designs in runtime.zero.stage_1_and_2.py

Communication overlapping did work under the following setup, but not in a satisfying way: when profiling the actual communications I found that only part of the allreduce was covered by computation. According to my knowledge, allreduce in ZeRO-2 does not have any dependencies between adjacent nn modules (model parallelism was not turned on), so fully-overlapping should be achievable in theory.

#------ ZeRO-2 DS_CONFIG ------#

{
  "train_batch_size": 65536,
  "gradient_accumulation_steps": 1,
  "optimizer": {
      "type": "Adam",
      "params": {
      "lr": 0.00015
    }
  },
  "fp16": {
    "enabled": true,
    "loss_scale": 0
  },
  "zero_optimization": {
    "stage": 2,
    "overlap_comm": true,
    "contiguous_gradients": true,
    "reduce_scatter": true,
    "reduce_bucket_size": 1e6
  }
}

[BUG]
When diving into runtime.zero.stage_1_and_2.py, I found that there is a lock between computation stream and allreduce stream before bucketing and gradient communications. Given that CUDA device always returns False in get_accelerator().resolves_data_dependency(), the computation stream (get_accelerator().current_stream()) and communication stream (self.reduction_stream()) always have to wait for each other, until they are synchronized.
I wonder why the computation stream has to wait until the reduction stream finished, which prevented the following backward computations from overlapping the allreduce communications of previous gradients in runtime.zero.stage_1_and_2.py. I believe that the performance of ZeRO-2 can be further boosted if this restriction can be released without conflicts to other setups.

Start in deepspeed/runtime/zero/stage_1_and_2.py around the bucketing and gradient communication logic, then inspect get_accelerator().resolves_data_dependency() in accelerator/cuda_accelerator.py. Use the supplied ZeRO-2 configuration and profiling setup to trace the computation and reduction streams. Done means the allreduce overlaps subsequent backward computation without introducing conflicts or regressing other setups.