Ok, LAMB and LARS optimizer seems interesting, I’m not familiar with them. I’ll definitely give them a try.
Just for information, since last message:
I tried training with Adam and eps=1e-6 + amp, the training start well, decreasing at a good rate, until it goes to NaN after some time.
I tried training with Adam and default eps=1e-8 without amp, the training is fine.
I am gonna have a look to the param values and optimizer states before the training collapse soon (I need to finish some runs before).
Just a question, why do you advice me to try LAMB optimizer? Is it because you think the trust ratio may help to upscale the loss gradient in the layers where the said gradient is very small?
Oh also, the batch size I use is relatively small as I’m limited by my gpu memory and that I work with quite heavy input data and model. This may not be ideal with LAMB or LARS, as they seems to work better on huge batch. To give a rough idea:

With Adam optim without AMP, the max batch size I can use is only 3.
With Adam optim with AMP, the max batch size I can use is around 5.
With SGD or RMSPROP optim with AMP, the max batch size I can use is around 16. RMSPROP converge really slower and with significantly lower accuracy than Adam in my case. I did not tried SGD yet, but I plan to.

Still a lot of thing to try and test, I’ll give an update if I found something interesting.
Best regards,
Thomas