[Bug] grad_weight can't be NoneType when running with DeepSpeed on Zero3.

[ys950902]

When running the Megatron-DeepSpeed with DeepSpeed on zero3, the grad_weight is set to None by default, that will cause the error issue below:
https://github.com/microsoft/DeepSpeed/blob/master/deepspeed/runtime/zero/stage3.py#L1253
AttributeError: 'NoneType' object has no attribute 'numel'
This pr is to fix this error issue, the weight.grad equals to grad_weight.

[ys950902]

@tjruwase could you please take a look on this pr, it will block the users to use the zero3 stage.

[ys950902]

Any concern for this pr?

[tjruwase]

@ys950902, apologies for the delay on this PR. I am confused about this issue since Megatron-DeepSpeed and ZeRO3 have always been compatible. Can you share some repro details to help my understanding? Thanks!

[ys950902]

@ys950902, apologies for the delay on this PR. I am confused about this issue since Megatron-DeepSpeed and ZeRO3 have always been compatible. Can you share some repro details to help my understanding? Thanks!

Hi @tjruwase, sorry for late response, for zero-bubble added in latest Megatron-DeepSpeed, the backward computation will be divided in two parts(weight gradient and output), and the weight gradient will not compute at once and will set to be None, when in running zero2/3 will divided the gradient may cause some unexpected error, this pr is to add the flag when doing pipeline-parallelism go the current path, when doing zero2/3 go the former path.

[ys950902]

Hi @tjruwase, if you still have some concern for this pr, please let me know.

[tjruwase]

Hi @tjruwase, sorry for late response, for zero-bubble added in latest Megatron-DeepSpeed, the backward computation will be divided in two parts(weight gradient and output),

Thanks for the explanation. In that case, shouldn't this behavior depend on whether zero-bubble is enabled? In other words, check for args.enable_zbh1_pipeline. Can you please clarify?

[ys950902]

Hi @tjruwase, sorry for late response, for zero-bubble added in latest Megatron-DeepSpeed, the backward computation will be divided in two parts(weight gradient and output),

Thanks for the explanation. In that case, shouldn't this behavior depend on whether zero-bubble is enabled? In other words, check for args.enable_zbh1_pipeline. Can you please clarify?

Thanks for your reply, yes, maybe check for args.enable_zbh1_pipeline is more reasonable, because for pipelien-parallelism, only supported for ZERO1/ZERO0 on Deepspeed, it won't divided the gradient, so is also okay set to NONE on weight gradient calculation. I will modify the check to args.enable_zbh1_pipeline to avoid the confuse.

[ys950902]

Hi @tjruwase, I noticed the bug issue #442 in zero-bubble, so I did a little modified to better understanding for customers, and this pr is the solution for this bug, Only enable args.enable_zbh1_pipeline go this path calculate the weight gradient when schedule is pop, when is not enable go the former path to calculate the weight gradient. I think this is more clear. And if approved, could you please also merge this pr, then we can upgrade the Megatron-DeepSpeed to latest for next release.