Horrifying grad bug present in v0.21.0-pre.1/2

[khoek]

A simple, handwritten token transformer I wrote for some mathematics research has found a grad bug in v0.21.0-pre.1. Good/bad loss plot, for identical model settings between v0.20.1 and v0.21.0-pre.1:

I'm in the process of trying to hunt down what the problem is, but thought I should let you know.

[khoek]

I've confirmed it on v0.21.0-pre.2, too. It's hard for me to produce a reproducer for you because the dataset I am training on is a 50/50 classification problem with millions of items in the corpus, and as observed the broken version kind-of does learn the 50-50 split (observe that the loss decreased below log 2=0.693), and there is an embedded feature which helps it further. Nonetheless, of course it fails to actually learn the task as you can see, and something is seriously wrong.

[khoek]

Progress update: I'm trying to run git bisect now but random initialization and a modest number of steps has been too noise, now I'm trying to init from a known good checkpoint and try to detect instability from there. Seems promising.

[laggui]

Thanks for flagging the issue, that's one we need to make sure doesn't make its way to the official 0.21 release 🙏

Please let me know if you manage to narrow it down.

It's not immediately clear what change might have triggered this, which wasn't caught by our tests (and some example training we use to validate functionalities).