Outlier detection: catch more outliers by not updating moving average with skipped updates

[ademeure]

This is an improvement to the znorm/zgrad update skipping mechanisms (-sl and -sg) to avoid skipping updates for outliers. Note that znorm will still be updated if zgrad is an outlier that causes the update to be skipped (and vice-versa), but that shouldn't matter much.

Combined with "-sg 5" my intuition is that this should significantly improve stability (famous last words...) - we should still try to root cause the instability as much as possible after that of course.

[ademeure]

Updated the PR - very promising based on a short local run, going to use this for a longer run and for FP8 runs now...

The most surprisingly effective change seems to be that whenever grad z-norm is >2, beta2 is decreased from 0.95 to 0.9. As per various papers including https://arxiv.org/pdf/2304.13013, a smaller beta2 seems to improve stability, at a slight cost in how fast the loss improves (and the final value it converges to). I suspect that by selectively reducing beta2 based on the grad z-norm, we can get some of the stability benefits of a lower beta2 without actually hurting the performance as much or at all.

When grad z-norm is >8, it will "mostly skip" the update - i.e. learning rate is 0.1x, weight decay is 0.2x, and beta1 is increased from 0.90 to 0.95 in order to prevent outliers from "lingering" in the 1st momentum too much (reducing learning rate to 0 but not increasing beta1 resulted in some lingering high grad z-norms which went away after that change).