[QUESTION] overlap-grad-reduce in V-ZB

[LitLeo]

Your question
In version 1.1, the steps to implement the overlap-grad-reduce feature are as follows:

1、Within the DistributedDataParallel class, put delay weights in weight_grad_buffers, while other weight are stored in grad_buffers.
2、Add DistributedDataParallel.async_reduce_grad method, which can synchronize all grad_buffers or a specified weight_grad_buffer.
3、In the WeightGradStore.clear() method, first call model.async_reduce_grad to trigger start_grad_sync func for all weights in grad_buffers. Subsequently, perform the backward for the weights, synchronizing the gradient of each weight immediately after its gradient computation is completed.

As for how to ensure that all weights in weight_grad_buffers have executed the start_grad_sync operation within the WeightGradStore.clear() method？

As shown in the timeline below, for the last microbatch, not all weight backward computations are performed within the WeightGradStore.clear(). This means that not all delay weights will invoke the async_reduce_grad function.

[LiJunscs]

In the current version, they cancel the support of overlap_grad_reduce, which can be found in the file megatron.core.zbpp_utils.py. I feel so confused. I try to fix this in my environment with no grad_accumulation_fusion and overlap_grad_reduce, I find that it causes a bug because the back propagation' hook of one parameter is called multiple times. I'm not sure that is my error or not.

[QPHutu]

As for how to ensure that all weights in weight_grad_buffers have executed the start_grad_sync operation within the WeightGradStore.clear() method？

As shown in the timeline below, for the last microbatch, not all weight backward computations are performed within the WeightGradStore.clear(). This means that not all delay weights will invoke the async_reduce_grad function.

@LitLeo
If I understand correctly, you are asking how to invode the async_reduce_grad function for the white-colored weights in first device?

Note that we use different queues for each vitual pipeline stage, so they will do WeightGradStore.clear() separately.

https://github.com/sail-sg/zero-bubble-pipeline-parallelism/blob/zero-bubble-v0.1.0/megatron/core/weight_grad_store.py#L13

[QPHutu]

In the current version, they cancel the support of overlap_grad_reduce, which can be found in the file megatron.core.zbpp_utils.py. I feel so confused. I try to fix this in my environment with no grad_accumulation_fusion and overlap_grad_reduce, I find that it causes a bug because the back propagation' hook of one parameter is called multiple times. I'm not sure that is my error or not.

@LiJunscs Sorry for the inconvenience. After we updated to the latest version of Megatron, we just found so many features of our implementation cannot be compatible anymore. As Megatron is so messy and has thousands of commits per month (so we guess every team has their own stable version), we decided to stop supporting those trivial features and mainly show our core ideas. If anyone want to implement these features by self, we recommand to refer to our previous version.

[LiJunscs]

In the current version, they cancel the support of overlap_grad_reduce, which can be found in the file megatron.core.zbpp_utils.py. I feel so confused. I try to fix this in my environment with no grad_accumulation_fusion and overlap_grad_reduce, I find that it causes a bug because the back propagation' hook of one parameter is called multiple times. I'm not sure that is my error or not.

@LiJunscs Sorry for the inconvenience. After we updated to the latest version of Megatron, we just found so many features of our implementation cannot be compatible anymore. As Megatron is so messy and has thousands of commits per month (so we guess every team has their own stable version), we decided to stop supporting those trivial features and mainly show our core ideas. If anyone want to implement these features by self, we recommand to refer to our previous version.

I think your work is significant and difficult.

I have some suggestions for new megatron, not latest.

• The gradient_accumulation_fusion:
  if the flag is false, you can detach the grad_weight and add the grad_weight in place, because the weight is a leaf node in compute graph, like this:

  grad_weight.detach().add(grad_output.t() @ input)

• The overlap_grad_reduce:
  In version above, also latest, the ParamAndGradBuffer has a param called is_last_microbatch, Buffer registers the back hook when the is_last_microbatch is True. When the buffer initialize, is_last_microbatch is True. The flag is set to False when no_sync_grad is called in 1f1b or interleaved 1f1b. But in the zero bubble, you don't use the context manager, the back hook will be called after ColumnParallel backward when scheduling B or b, then register twice. This is my problem. I'm not sure whether this will do help to you.

[QPHutu]

In the current version, they cancel the support of overlap_grad_reduce, which can be found in the file megatron.core.zbpp_utils.py. I feel so confused. I try to fix this in my environment with no grad_accumulation_fusion and overlap_grad_reduce, I find that it causes a bug because the back propagation' hook of one parameter is called multiple times. I'm not sure that is my error or not.

@LiJunscs Sorry for the inconvenience. After we updated to the latest version of Megatron, we just found so many features of our implementation cannot be compatible anymore. As Megatron is so messy and has thousands of commits per month (so we guess every team has their own stable version), we decided to stop supporting those trivial features and mainly show our core ideas. If anyone want to implement these features by self, we recommand to refer to our previous version.

I think your work is significant and difficult.

I have some suggestions for new megatron, not latest.

• The gradient_accumulation_fusion:
  if the flag is false, you can detach the grad_weight and add the grad_weight in place, because the weight is a leaf node in compute graph, like this:
  grad_weight.detach().add(grad_output.t() @ input)
• The overlap_grad_reduce:
  In version above, also latest, the ParamAndGradBuffer has a param called is_last_microbatch, Buffer registers the back hook when the is_last_microbatch is True. When the buffer initialize, is_last_microbatch is True. The flag is set to False when no_sync_grad is called in 1f1b or interleaved 1f1b. But in the zero bubble, you don't use the context manager, the back hook will be called after ColumnParallel backward when scheduling B or b, then register twice. This is my problem. I'm not sure whether this will do help to you.

Thanks for your suggestions, I think it will be helpful when we plan to implement these features.