add inverse_sqrt lr decay style

[NouamaneTazi]

adds support for inverse square decay style
P.S: ignores min-lr param for now

To use it just add these flags to your command:

--lr-decay-style=inverse_sqrt --lr-warmup-style=linear/constant

[thomasw21]

LGTM! You may want to also add the normalized version as well depending on what was used for PaLM for example.

[thomasw21]

Oh I don't have write access to this repo. You'll have to ask someone that does. cc @spyysalo

[thomasw21]

Lgtm!