[eshaanpathak]

Gradient clipping wasn't working :(

The clip call (clip_grad_norm_) was inside training_step, before .backward() runs. Lightning calls training_step --> backward() --> clip --> optimizer.step(), so at the point this ran, p.grad was either None or leftover from the previous step, not the gradient about to be used. With gradient_accumulation: 1 (the default config), clipping did nothing. With accumulation > 1, it scales down the gradient mid-accumulation before later micro-batches even get added in.

Moved the clip into configure_gradient_clipping, which Lightning runs after backward (and after DDP grad sync), once per optimizer step.

Checked this against pytorch-lightning 1.9.5 (what's pinned) with a small repro, with and without accumulation. Before the fix, the optimizer's gradient was never capped no matter what gradient_clip_norm was set to. After the fix it is, every step.

Also added a gradient_norm_clipped log (post-clip) next to the existing gradient_norm (now correctly the current step's value, not the previous one).

All of the HF checkpoints may have been trained without clipping. Since it looks like clipping was enabled at 20K steps, it may have unlikely affected the checkpoints, but you would need to check your training run charts to confirm this!

Here are two charts to visualize the gradient getting clipped at 100 steps to 0.1: true_gradient_norm is the main branch right now before my changes, and gradient_norm_clipped is with my changes in this PR.

[knightnemo]

lgtm.