AdaBelief and AMP seem incompatible about lightweight-gan

[juntang-zhuang]

@Mut1nyJD @FilipAndersson245 Hi, I just tried a by-pass to deal with the mixed precision issue, that is to cast weight and gradient to float32, update, then cast to float16. In this way the float32 burden is only applied to the weight update, but not the backward, so the computation overload would not be too much. See the link to code juntang-zhuang/Adabelief-Optimizer#31 (comment). Please let me know if you have other suggestions.

[juntang-zhuang]

I guess it might be a problem with low-precision with epsilon. Because default eps=1e-8 for Adam (1e-16 for AdaBelief), this number will be just rounded to 0 in 16 bit precision. In this case, the update is pretty close to divide by 0 (or a very very small number), and the eps completely loses its impact as in float32 format.

[lucidrains]

Hmm, I don't really know, you should ask @FilipAndersson245

[FilipAndersson245]

I have not tested Adabelief together with mixed-precision mode, maybe submit an issue at the authors github?

[Mut1nyJD]

I have not tested Adabelief together with mixed-precision mode, maybe submit an issue at the authors github?

Nope not yet wanted to ask here first if anyone had some success, but will do, thank you!

[juntang-zhuang]

Hi, thanks for testing with AdaBelief. AdaBelief might be incompatible with mixed-precision, because low precision might cause difference in gradient to be 0, hence division by 0. I'm not quite familiar with low-precision training, could you point to the code where low-precision takes effect? We will look into it and perhaps find some solution for the next release of AdaBelief. Thanks a lot.

[Mut1nyJD]

Hmm okay I am not so sure if that it is a pure AdaBelief problem it probably is just far worse there.
But if with normal Adam it seems to become unstable as well it just takes a lot longer after a bit more than 50000 iterations it crashed out with the same error. I am still running PyTorch1.6 maybe worth upgrading to 1.7

[Mut1nyJD]

But if with normal Adam it seems to become unstable as well it just takes a lot longer after a bit more than 50000 iterations it crashed out with the same error. I am still running PyTorch1.6 maybe worth upgrading to 1.7

Strangely enough this does not seem to happen when I increase the number of attention layers, maybe just pure coincidence?

[Mut1nyJD]

@juntang-zhuang
Great! Sounds like a reasonable workaround to me so should work. Happy to give it a try again when I have some time.