Optim: APOLLO optimizer integration
Author: zhuhanqing
Pull request: yes

# What does this PR do?

This PR integrates APOLLO (Approximated Gradient Scaling for Memory-Efficient LLM Optimization) into Hugging Face's Transformers.
APOLLO is a memory-efficient optimizer designed for LLM pre-training and full-parameter fine-tuning, offering SGD-like memory cost with AdamW-level performance.

📜 Paper: https://arxiv.org/abs/2412.05270
💻 Code: https://github.com/zhuhanqing/APOLLO

## Why APOLLO?
APOLLO introduces a new level of memory efficiency for LLM optimization:
✅ Ultra-Low Memory Usage → Achieves significant savings, even beyond GaLore, approaching SGD-level efficiency.
✅ Adam(W)-Level Performance → Maintains or surpasses Adam(W) performance, validated on LLaMA models up to 7B scale.
✅ No Expensive SVD Computation → Unlike GaLore, APOLLO leverages lightweight random projection, avoiding training stalls in large-scale LLM fine-tuning.

## Third-Party Validation of APOLLO

APOLLO has been merged into [LLaMA-Factory](https://github.com/hiyouga/LLaMA-Factory/pull/6617), [FluxML](https://github.com/FluxML/Optimisers.jl/pull/196), and with a validated performance in the [post](https://bsky.app/profile/benjmurrell.bsky.social/post/3lcyfrf5b7k2u).

With these validations, merging APOLLO into Transformers would offer Hugging Face users an efficient, memory-friendly optimizer for training LLMs—reducing GPU memory requirements and making large-scale model training more accessible! 🚀

## Test of the integration

Following the approach in [PR #29588](https://github.com/huggingface/transformers/pull/29588) for GaLore, I have successfully integrated APOLLO into Hugging Face Transformers.

✅ Ensuring Correctness
To verify the integration, I have:
1️⃣ Added multiple unit tests in `tests/trainer/test_trainer.py`.
2️⃣ Manually tested the API using the following script:
```
import torch
import datasets
from transformers import TrainingArguments, AutoConfig, AutoTokenizer, AutoModelForCausalLM
import trl

train_dataset = datasets.load_dataset('imdb', split='train')

args = TrainingArguments(
    output_dir="./test-apollo",
    max_steps=100,
    per_device_train_batch_size=2,
    optim="apollo_adamw",
    optim_target_modules=["attn", "mlp"]
)

model_id = "mistralai/Mistral-7B-v0.1"

config = AutoConfig.from_pretrained(model_id)

tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_config(config).to(0)

trainer = trl.SFTTrainer(
    model=model, 
    args=args,
    train_dataset=train_dataset,
    dataset_text_field='text',
    max_seq_length=512,
)

trainer.train()
```

## Who can review?

@muellerzr and @SunMarc  – Would love your feedback on this integration! 😊
Let me know if any modifications are needed. Thanks! 🚀



Comments:
Hello @muellerz, 

The auto code check has some failures as above, but it seems not from our code but from other code, 
e.g., [tests_torch](https://app.circleci.com/pipelines/github/huggingface/transformers/117737/workflows/626a7625-ef6d-4aac-8914-0b3127e9988f/jobs/1584963?utm_campaign=vcs-integration-link&utm_medium=referral&utm_source=github-checks-link&utm_content=summary) and [check_repository_consistency](https://app.circleci.com/jobs/github/huggingface/transformers/1584943)

Hi @SunMarc and @muellerz,

I have fixed your comment and updated our commit! 
Thank you so much for your quick response and efforts in reviewing our codes.

Yep, I closely follow the GaLore implementation to add APOLLO support to ***ensure our code quality meets the standard of hugging face***
APOLLO shares a similar method to using low-rank space to reduce memory costs like GaLore, making it feasible to follow GaLore implementation. APOLLO uses the low-rank space to approximate the gradient scaling factor for full-rank raw gradient, while our previous GaLore fundamentally differs from the new APOLLO as they use a more projected gradient descent method and perform updates in low-rank space, leading to accuracy loss due to the usage of low-rank updates.

**Fix your comments**
I have fixed your comments in the new commit.

**Fix redundancy**
I believe the redundancy lies in the `src/transformers/trainer.py` since we construct the APOLLO optimizer using a similar logic as GaLore by separating modules into low-rank and non-low-rank ones and appending those params as params groups.  I modularize the construction logic such that GaLore and APOLLO only need to call a function to share the same codes.

Thank you again for your time and help!
Let me know if any modifications are needed. Thanks! 🚀

Best wishes,
Hanqing

Hi @SunMarc,

Thank so much for your feedback!

I submit two commits, with the first one adding the typing in the `setup_low_rank_optimizer` function.

For the second commit, I agree that our current description of APOLLO in `trainer.md` is lazy and simply follows the GaLore one. I rewrote the description to make sure it is clean and clear and address your comments. In particular, I show how to precisely configure APOLLO-Mini by passing the right arguments. 

Thank you again for your time and help! Looking forward to your feedback!

Best wishes,
Hanqing

Letting @ArthurZucker having a quick look if possible. Otherwise, I'll merge in 2/3 days ! 

The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_36062). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

Hi @SunMarc, thank you so much for your help! Btw our APOLLO has been accepted by [MLSys 2025](https://mlsys.org/), just announced today. I am eager to see APOLLO contribute more to the open-source community, especially democratizing LLM training with the integration with HF trainer!

Congrats for your acceptance !