Pretraining loss explosion

[mattgorb]

I have been trying to get this repo working for several months, but my loss keeps exploding between 30k and 100k iterations.

I have tried many things:
Turn flash attention off ( based on this issue: #524)
Using fp16 (based in this: #468)
Using GPT-4 tokenizer (based on #468)

At first the loss was going back up to about 8-10, now it is just going to NaN with fp16.

I have also tinkered with other setting such as gradient clipping, learning rate, etc. I keep my configuration at roughly 500k batch size.

I am lost on what to try next. Did anyone else fix this issue?

I have gotten GPT-2 Small down to about 3.0 loss.

[HarrisonUnifyAI]

My loss also exploded on an initial attempt with 8 A100 instance. I disabled flash attention and that seems to have fixed it, though it is much slower.

[beeeai]

@mattgorb @HarrisonUnifyAI I have not changed the default settings, other than the seed. (doing ddp on 8GPUs)
I get this explosion which starts around 30k as well. Do you think flash attention is the reason?
(The purple one is all default values, with the seed set to 1337. The blue one is another seed.)