add grad accum support

[lchu6]

No description provided.

[lchu6]

@daviswer I think the reporting should be correct but help me double check as I typed really quick.

ddp_stats[1] stores gnorm and we add only every grad_accum_steps.
ddp_stats[2] stores "denominator" and we only count "real steps".
ddp_stats[0] stores loss and we count all steps (so grad_accum_steps times more value than above 1 and 2), but since our loss was anyway loss = loss / grad_accum_steps, so this value should also be correct with above denominator.

So if we have a total steps of 15 with grad accum step = 3:
we have 5 values for gnorm and denominator=5, while 15 values for loss. but those 15 values of loss was already original_loss/3, so it pair with denominator=5.

[daviswer]

That logic checks out, looks good!