Add support to APOLLO

[zhuhanqing]

Adding Support for APOLLO

This pull request introduces APOLLO, a memory-efficient optimizer tailored for pre-training large language models (LLMs) and full-parameter fine-tuning. APOLLO significantly reduces memory overhead compared to AdamW while maintaining comparable performance, making it ideal for LLM workflows.

Project page: APOLLO GitHub Repository

---

Features

• Low Memory Overhead: APOLLO achieves memory usage similar to SGD, even in rank-1 space, enabling efficient pre-training and fine-tuning of LLMs.
• AdamW-level performance: APOLLO achieves on-par or even better performance than AdamW, even at a much reduced rank.
• No SVD Requirement: Unlike GaLore, APOLLO does not require Singular Value Decomposition (SVD), eliminating the SVD delays commonly encountered when using GaLore.

---

Tests

We have added a test in the examples/extras/apollo directory. Additionally, we conducted a comparative evaluation with GaLore by fine-tuning models and testing on the MMLU task.

GaLore Performance

Average: 64.96
           STEM: 55.43
Social Sciences: 75.66
     Humanities: 59.72
          Other: 71.25

APOLLO Performance (Scaling Factor = 32)

With a scaling factor derived from the ratio of LLaMA-8B dimension (4096) to rank (128):

Average: 65.03
           STEM: 55.47
Social Sciences: 76.15
     Humanities: 59.60
          Other: 71.28

---

Why APOLLO?

APOLLO offers a compelling alternative to existing optimizers, especially for large-scale LLM pre-training and fine-tuning. Its reduced memory footprint and simplicity (no SVD) make it a practical choice for researchers and engineers.

[hiyouga]

LGTM