[Bug] Missing weight gradients from LinearWithGradAccumulationAndAsyncCommunication when Zero Bubble Pipeline Parallelism Is disabled

[mksit]

I have observed a recent change in LinearWithGradAccumulationAndAsyncCommunication to store the gradient of weights in WeightGradStore as a part of the new Zero Bubble Pipeline Parallelism feature (#396):

Megatron-DeepSpeed/megatron/core/tensor_parallel/layers.py

Line 370 in 1280f59

WeightGradStore.put(total_input, grad_output, weight, gradientUpdateFunction)

However, the stored gradients are only accessed in deepspeed_zbh1_engine:

Megatron-DeepSpeed/megatron/core/pipeline_parallel/deepspeed_zbh1_engine.py

Line 108 in 1280f59

WeightGradStore.pop()

If the Zero Bubble Pipeline Parallelism feature is not enabled, it seems that the gradients are not being returned. Is this an expected behavior?