Feature request: GaLore optimizer

[gil2rok]

Feature request for the Gradient Low-Rank Projection (GaLore) optimizer.

The GaLore optimizer computes low-rank gradients way to dramatically reduce memory. The ArXiv paper is here and the original Github implementation is here.

GaLore is quite popular: HuggingFace has implemented it here, PyTorch Lightning is trying to implement it here. The original implementation has 1K+ Github stars.

A good starting point may be this clean PyTorch implementation here in the PyTorch Optimizers library.

Lastly, readers should also be aware of the improved Q-GaLore paper and repository here and here respectively.

[vroulet]

Hello @gil2rok,

Thanks for pointing this out! Would you be willing to contribute with an implementation?
If not, we'll add it to our todo list but any contributions are super welcome! :)

[gil2rok]

At the moment, I do not have the bandwidth to add this optimizer. If I have some unexpected time, will try to come back to this cause I think it'd be super cool to implement.

Thanks for all you guys do for this library!

[viralvgupta]

Hi All,
I am working with my colleague and we have a working prototype of Galore implementation. We are doing some final tests. Can you suggest if we should follow PR diff. Please point us to any clear guidelines. super appreciate the help.

[vroulet]

That's great news!
Yes, the workflow currently would be to

1. add this optimizer in the contrib folder, note that the PR you pointed out introduced some errors (like in the dtypes at initializations), so take a look at the latest version of e.g. prodigy or schedule_free.
2. add this optimizer to the common_test of this folder, in particular making sure it can support some of optax functionalities like MultiSteps.
3. if some specific functionalities of this optimizer requires some tests, feel free to add a file that make these tests (schedule_free for example has its own specific tests)

Feel free to start a PR and send it to us so that we can help reviewing it and integrating it.

Thank you for this!