APOLLO optimizer

[fblgit]

⚠️ Please check that this feature request hasn't been suggested before.

• I searched previous Ideas in Discussions didn't find any similar feature requests.
• I searched previous Issues didn't find any similar feature requests.

🔖 Feature description

Any plan or timeline for APOLLO optimizer?

https://arxiv.org/abs/2412.05270

looks like.. very interesting

✔️ Solution

Implementing https://arxiv.org/abs/2412.05270

❓ Alternatives

buy me 4x gpu :D

📝 Additional Context

No response

Acknowledgements

• My issue title is concise, descriptive, and in title casing.
• I have searched the existing issues to make sure this feature has not been requested yet.
• I have provided enough information for the maintainers to understand and evaluate this request.

[winglian]

Thanks @fblgit! We'll look into this and see how feasible this is without a reference implementation

[fizzAI]

https://github.com/zhuhanqing/APOLLO/tree/main/apollo_torch there's a reference implementation now
would be cool to see :3

[ehartford]

FYI implementing an optimizer is quite independent from Axolotl. It's a huggingface transformers thing.

Check out how I implemented grokAdamW

huggingface/transformers#32521

[fizzAI]

Yes, but axo already has some optimizers of its own that it patches into TF's on its own, ie ADOPT, and it's probably much easier to get a PR doing that working here than dealing with upstream transformers