[BUG] LM converges much slower when gradient is accumulated in FP32.

[torshie]

[BUG] LM converges much slower when gradient is accumulated in FP32.

[torshie]

I'm not using megatron-deepspeed because some features are missing

[torshie]

I did further tests, it can be reproduced in a single node with multiple GPU's. The key to reproduce is to enable data parallel, so the problem probably lies somewhere in the gradient reduction code.

[torshie]

After reading stage_1_and_2.py for a couple of times, I think the problem may be in the backward() method, I may not fully understand the source code, any comments are welcome:

• In certain setup, DeepSpeedZeroOptimizer may call self.create_reduce_and_remove_grad_hooks() to register a gradient reduction hook.
  • in my case set overlap_comm to True will trigger this call, I didn't realize this until now.
• When loss.backward() is called, the registered hook is called.
• In the registered hook function, reduce_independent_p_g_buckets_and_remove_grads() will be called when is_gradient_accumulation_boundary is True.
• Inside reduce_independent_p_g_buckets_and_remove_grads(), get_gradient_for_reduction() will be called to access the accumulated gradients .grad_accum, which is updated by calling fill_grad_accum_attribute().
• However, in the backward() method, fill_grad_accum_attribute() is called after loss.backward().
• So basically when is_gradient_accumulation_boundary is True, get_gradient_for_reduction() will access the accumulated gradients before the final fill_grad_accum_attribute() call.

BTW, I think this is also the root cause of bug #5174 which is easily reproducible with the provided source code.