Always getting NaNs in long training

[danbochman]

I've been experimenting with the LION optimizer in your other (great) Imagen repository. I can share my anecdotal experience and combinations:

• Models of different sizes 0.2B, 0.7B and 1B params.
• Betas such as beta1 0.95 and beta2 0.98
• Learning rates 1e-4, 3e-5 and 1e-5.
• Triton kernel turned both True and False.

Training was indeed fast but unfortunately in the end always ended up yielding NaNs.

I think a potential issue could be how LION interacts with a warmup schedule; I am not sure if you're supposed to do warmup with this optimizer or not (which I always did).

[ysesst93013]

I have same problem :(

[SergeySakharovskiy]

same NaN issue with CosineAnnealing scheduler after the first epoch.

[xiangning-chen]

I've been experimenting with the LION optimizer in your other (great) Imagen repository. I can share my anecdotal experience and combinations:

• Models of different sizes 0.2B, 0.7B and 1B params.
• Betas such as beta1 0.95 and beta2 0.98
• Learning rates 1e-4, 3e-5 and 1e-5.
• Triton kernel turned both True and False.

Training was indeed fast but unfortunately in the end always ended up yielding NaNs.

I think a potential issue could be how LION interacts with a warmup schedule; I am not sure if you're supposed to do warmup with this optimizer or not (which I always did).

May I know the learning rate schedule you are using?

[zjutzyl]

same issue, i set a big weight decay to avoid it. i suppose that 'update=symbol * lr' enlarging abs(parameter) while symbol not changing.

[lindakasabian]

same here. sudden nan losses during 100 e training with onecyclelr and clipping