Deepspeed communication pattern clarification

[nishantagrawalgit]

Queries regarding communcation calls in Deepspeed for Stage1(Zero1) and Stage2(Zero2):

1. As per paper, Stage1(Zero1) has Optimizer split implementation and we should see below calls on 2 GPU's(1 layer GPT2/345m model) with NCCL:
   1 call - Reduce_Scatter for Gradients
   1 call - All_Gather for optimizer

   As per details mentioned in one of prior github issues(Deepspeed) that reduce_scatter for stage 1 is not implemented. Thats why it calls all_reduce for stage1 Zero optimizer. Is it still not implemented?

   Based on the logs we are seeing All_Reduce calls in the backpass + Reduce_Scatter calls post backpass before optimizer step. So can you explain why reduce_scatter needed here as we can just do all_reduce for the gradient and all_gather for optimizer?

   We are seeing an overflow occur at the end of Optimizer step(we expected to see All_Gather comms at the end of the optimizer which we dont observe). Can you pls explain why overflow occurs as mentioned in the below logs:
   rank 0 detected overflow nan in tensor 0:1 shape torch.Size([50304, 1024])
   [deepspeed] OVERFLOW! Skipping step. Attempted loss scale: 4294967296, reducing to 4294967296

2. For Stage2(Zero2), Zero with Optimizer+gradient split- In the backpass the gradient should get scattered to gpu that owns that gradient chunk for reduce, and then perform all gather for optimizer(as per paper). In the run logs we are seeing:
   No calls in the backpass which means no comms and compute overlap in backpass because of enabling fusion buffer. As a result we have see only 2 All_Reduce calls post backpass for gradients. Pls confirm if our understanding is correct?

[samyam]

@nishantagrawalgit Thank you for your interest in DeepSpeed.

a) Reduce Scatter for Stage 1 is now already implemented. It is the default option for Stage 1, and we have removed the allreduce version from Stage 1

We use reduce_scatter instead of all_reduce to reduce the communication volume by half.

b) Buffer Overflow is normal at the beginning of the training. It is because of the loss scaling that is done to ensure that the gradients are produced within the range that fp16 datatype can handle. We scale the loss by a larger number in the beginning which causes the gradients to explode outside of the range of the fp16 data type. We slowly bring the loss scaling factor down such that the gradients barely fits within the upper range of fp16. By pushing the gradients to the upper end of the fp16 spectrum we collapsing of the gradients on the low end of the fp16 spectrum to 0.0.

2. I am not entirely sure what you mean by "No calls in the backward" and "because of enabling fusion buffer" but here is some explanation on whats happening in Stage 2:

The reduction is happening in the backward using hooks that gets triggered after the gradients are computed for each parameter. Depending on the reduce_buffer_size, multiple gradients maybe combined into a single reduction (is this what you mean by fusion?), and once there is enough gradients to fill the reduce_bucket those will be reduced in the middle the backward prop.

If you enable "comm_overlap" in deepspeed_config, most of the reduction should overlap with computation.

The post backpass reduce calls that you are seeing are only meant to reduce gradients that were not reduced during the backpass because the reduce_buckets were not fully filled to trigger reduction during the backward pass.

[nishantagrawalgit]

Thanks @samyam for the insights. We understood your clarifications on Stage2. We have few more queries on Stage1 & Stage2 communication pattern:

• Stage1 we are running the latest release and we still see 16 All_reduce call for gradient calculation, 2 Reduce_Scatter calls for optimizer and 1 all gather calls. We were expecting to see only, Reduce_Scatter for the gradients and All_Gather for optimizer based on our understanding of zero paper.. Can you please explain what we have misunderstood in terms of comms pattern?

• Stage2 communication calls should be same as Stage1 - Scatter_Reduce of the gradients and All_Gather for optimizer. But what we see from Stage2 comms profiling is below -
  All_Reduce for the gradients and All_Gather for the optimizer
  If we doing All_Reduce for the gradients(assuming with the large reduce_bucket and doing only 1 All_Reduce at the end of the backpass) then it means all GPU's should be owning the full copy of gradients and then post sync we will free up the memory from other GPU's which are not responsible for owning that part of the gradients. Is this how the flow implemented currently?

[samyam]

@nishantagrawalgit it is very strange that you are seeing all_reduce_calls for stage 1. Are you able to see which file and line number those all_reduce calls are coming from? Stage 1 should not have any all_reduce calls, only reduce scatter calls.

For Stage 2 there are several config parameters that will control the communication behavior:
i) reduce_scatter: If you set it to be false, it will use all_reduce instead of reduce
ii) reduce_bucket_size: If you set this to be as large as the entire model, then you will see the behavior you are describing above. This would not be the intended use of the reduce_bucket_size for zero stage 2. If you want the memory savings from stage 2, you should set reduce_bucket_size should be less than or equal to the partition size.

For stage 2, we generally see the best performance for large models when we set the reduce bucket size to about 50M and enable "comm_overlap".

[nishantagrawalgit]

@samyam Yes we are able to see Reduce calls(not Reduce_scatter) in Stage2 now by setting "Continuous_gradient=True" explicitly. We need to understand the significance of this flag which triggers Reduce calls instead of All_Reduce in Stage2.

For Stage1 - We are still seeing AllReduce calls in the backpass even by setting "continuous_gradient=True". As per our understanding, All_Reduce calls in backpass were getting triggered using gradient hooks in PyTorch DDP automatically. So we don't know how to change to Reduce_Scatter calls instead of All_Reduce(using latest version of Deepspeed lib). Can you pls help?

[samyam]

@nishantagrawalgit The contiguous_gradient flag defragments the memory during backward propagation. You can find details on this in our archive paper https://arxiv.org/abs/1910.02054 in section 6.3.

The reduce-scatter in stage 2 is implemented as a sequence of reduce calls, so there are no direct reduce_scatter calls as you noticed.

For Stage 1, are you using Pytorch DDP with DeepSpeed? If you are, can you try running with DeepSpeed without the Pytroch DDP. DeepSpeed has its own DDP implementation which will do the appropriate reduce_scatter at the end of the backward pass for Stage 1.