feat: LOMO support

[KaiLv69]

Feature request

Support for the LOMO training strategy, which allows full-weights training of 7B models on consumer 24GB cards like RTX 4090.

ArXiv: https://arxiv.org/pdf/2306.09782.pdf https://arxiv.org/pdf/2310.10195.pdf
Implementation: https://github.com/OpenLMLab/LOMO (An installable package lomo-optim is also available here.)

Motivation

This is another step towards democratizing LLM. And the variant AdaLomo performs on par with AdamW.

Your contribution

I'm willing to submit a PR and modify the lomo-optim package if needed.
I may need some kind help for the integration on what I should do.

[amyeroberts]

cc @younesbelkada

[younesbelkada]

Nice and exciting work ! Happy to get to it after the work on #29588 !