FEAT / Optim: Add GaLore optimizer
Author: younesbelkada
Pull request: yes

# What does this PR do?

As per title, adds the GaLore optimizer from https://github.com/jiaweizzhao/GaLore 

Fixes: https://github.com/huggingface/transformers/issues/29512

This is how I am currently testing the API:
```python
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-galore",
    max_steps=100,
    per_device_train_batch_size=2,
    optim="galore_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()
```

cc @pacman100 @muellerzr 

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

see https://github.com/OpenAccess-AI-Collective/axolotl/pull/1370#issuecomment-1988422268 for intermediate results

Is it possible to integrate per-layer weight updates as described in the galore paper?

![image](https://github.com/huggingface/transformers/assets/94501378/b28940c9-4baf-4686-9e8b-5a39ac075b25)

It seems to reduce memory usage significantly:

![image](https://github.com/huggingface/transformers/assets/94501378/e95995a1-7fac-4089-8dd9-3593bb06e561)

But the [original implementation](https://github.com/jiaweizzhao/GaLore/blob/master/torchrun_main.py#L334-L368) is somewhat complex.

> Is it possible to integrate per-layer weight updates as described in the galore paper?
> 
> ![image](https://private-user-images.githubusercontent.com/94501378/311920992-b28940c9-4baf-4686-9e8b-5a39ac075b25.png?jwt=eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJpc3MiOiJnaXRodWIuY29tIiwiYXVkIjoicmF3LmdpdGh1YnVzZXJjb250ZW50LmNvbSIsImtleSI6ImtleTUiLCJleHAiOjE3MTAzMTQ5NzYsIm5iZiI6MTcxMDMxNDY3NiwicGF0aCI6Ii85NDUwMTM3OC8zMTE5MjA5OTItYjI4OTQwYzktNGJhZi00Njg2LTllOGItNWEzOWFjMDc1YjI1LnBuZz9YLUFtei1BbGdvcml0aG09QVdTNC1ITUFDLVNIQTI1NiZYLUFtei1DcmVkZW50aWFsPUFLSUFWQ09EWUxTQTUzUFFLNFpBJTJGMjAyNDAzMTMlMkZ1cy1lYXN0LTElMkZzMyUyRmF3czRfcmVxdWVzdCZYLUFtei1EYXRlPTIwMjQwMzEzVDA3MjQzNlomWC1BbXotRXhwaXJlcz0zMDAmWC1BbXotU2lnbmF0dXJlPTNlYWNmOWNlYmNkNDQ1YTE0YzRlMjU3YzU0YzJhNWVmMDJhODE2NWUyYTkzYmMzNzY2OGY3YmQyY2VkNDY3MjMmWC1BbXotU2lnbmVkSGVhZGVycz1ob3N0JmFjdG9yX2lkPTAma2V5X2lkPTAmcmVwb19pZD0wIn0.TdqveKK_1bWtOLLOnFJs3Y5QPmHZfLvH-xy8JJP9tcY)
> 
> It seems to reduce memory usage significantly:
> 
> ![image](https://private-user-images.githubusercontent.com/94501378/311923308-e95995a1-7fac-4089-8dd9-3593bb06e561.png?jwt=eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJpc3MiOiJnaXRodWIuY29tIiwiYXVkIjoicmF3LmdpdGh1YnVzZXJjb250ZW50LmNvbSIsImtleSI6ImtleTUiLCJleHAiOjE3MTAzMTQ5NzYsIm5iZiI6MTcxMDMxNDY3NiwicGF0aCI6Ii85NDUwMTM3OC8zMTE5MjMzMDgtZTk1OTk1YTEtN2ZhYy00MDg5LThkZDktMzU5M2JiMDZlNTYxLnBuZz9YLUFtei1BbGdvcml0aG09QVdTNC1ITUFDLVNIQTI1NiZYLUFtei1DcmVkZW50aWFsPUFLSUFWQ09EWUxTQTUzUFFLNFpBJTJGMjAyNDAzMTMlMkZ1cy1lYXN0LTElMkZzMyUyRmF3czRfcmVxdWVzdCZYLUFtei1EYXRlPTIwMjQwMzEzVDA3MjQzNlomWC1BbXotRXhwaXJlcz0zMDAmWC1BbXotU2lnbmF0dXJlPTdmZWU2ZjQyODFjNTNlMjI2NDM2NDE0OTllMjU2N2UzOWMwM2VjNDhmMzU0YjJmMDZhMTZiMmNlOTQ5NDIxY2UmWC1BbXotU2lnbmVkSGVhZGVycz1ob3N0JmFjdG9yX2lkPTAma2V5X2lkPTAmcmVwb19pZD0wIn0.oXhjQ0wCZlQTJ33cwoohgZfgbtJVA4feJ6pKws-Tv9I)
> 
> But the [original implementation](https://github.com/jiaweizzhao/GaLore/blob/master/torchrun_main.py#L334-L368) is somewhat complex.

Sure, per-layer weight update is crucial to the current implementation of GaLore (may not be in the future). Without that, it would require an additional 14GB GRAM for the gradients for a 7B model, making full-parameter fine-tuning infeasible with a 24GB GPU.

You may refer to our discussion https://github.com/jiaweizzhao/GaLore/issues/6 for more empirical results.

Thanks everyone ! 
@hiyouga I had a look at your implementation and it looks super elegant ! It gave me some inspiration to add layer-wise GaLoRE in this PR !
With the recent changes I am able to full-fine tune Mistral-7B in fp32 on a single A10 80GB GPU taking appox 44GB VRAM and the loss goes down quite smoothly:
```bash
{'loss': 10.6212, 'grad_norm': 0.0, 'learning_rate': 0.001, 'epoch': 0.0}                                                                               
{'loss': 9.1162, 'grad_norm': 0.0, 'learning_rate': 0.001, 'epoch': 0.0}                                                                                
{'loss': 8.5984, 'grad_norm': 0.0, 'learning_rate': 0.001, 'epoch': 0.0}                                                                                
{'loss': 8.339, 'grad_norm': 0.0, 'learning_rate': 0.001, 'epoch': 0.0}                                                                                 
{'loss': 8.8677, 'grad_norm': 0.0, 'learning_rate': 0.001, 'epoch': 0.0}                                                                                
{'loss': 8.4181, 'grad_norm': 0.0, 'learning_rate': 0.001, 'epoch': 0.0}                                                                                
{'loss': 8.3885, 'grad_norm': 0.0, 'learning_rate': 0.001, 'epoch': 0.0}                                                                                
{'loss': 8.0438, 'grad_norm': 0.0, 'learning_rate': 0.001, 'epoch': 0.0}                                                                                
{'loss': 8.0605, 'grad_norm': 0.0, 'learning_rate': 0.001, 'epoch': 0.0}                                                                                
{'loss': 8.1086, 'grad_norm': 0.0, 'learning_rate': 0.001, 'epoch': 0.0}                                                                                
{'loss': 8.0195, 'grad_norm': 0.0, 'learning_rate': 0.001, 'epoch': 0.0}                                                                                
{'loss': 8.4299, 'grad_norm': 0.0, 'learning_rate': 0.001, 'epoch': 0.0}                                                                                
{'loss': 7.9611, 'grad_norm': 0.0, 'learning_rate': 0.001, 'epoch': 0.01}                                                                               
{'loss': 7.9822, 'grad_norm': 0.0, 'learning_rate': 0.001, 'epoch': 0.01}                                                                               
{'loss': 7.9144, 'grad_norm': 0.0, 'learning_rate': 0.001, 'epoch': 0.01}                                                                               
{'loss': 7.5685, 'grad_norm': 0.0, 'learning_rate': 0.001, 'epoch': 0.01}                                                                               
{'loss': 7.9422, 'grad_norm': 0.0, 'learning_rate': 0.001, 'epoch': 0.01}                                                                               
{'loss': 7.7231, 'grad_norm': 0.0, 'learning_rate': 0.001, 'epoch': 0.01}                                                                               
{'loss': 7.6636, 'grad_norm': 0.0, 'learning_rate': 0.001, 'epoch': 0.01}                                                                               
{'loss': 7.789, 'grad_norm': 0.0, 'learning_rate': 0.001, 'epoch': 0.01}   
```
The script I used:
```python
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-galore-new",
    max_steps=100,
    per_device_train_batch_size=2,
    optim="galore_adamw_layerwise",
    optim_target_modules=[r".*.attn.*", r".*.mlp.*"],
    gradient_checkpointing=True,
    logging_strategy="steps",
    logging_steps=5,
    learning_rate=2e-3,
    save_strategy="no",
    run_name="galore-imdb",
    lr_scheduler_type="cosine"
)

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()
```
This PR is ready for review! 

Nice work @younesbelkada !
Note that the layerwise approach may not support DDP currently, please see https://github.com/jiaweizzhao/GaLore?tab=readme-ov-file#train-7b-model-with-a-single-gpu-with-24gb-memory
I guess it may be because the torch's all-reduce operator collects all the gradients over multiple GPUs in each optimization step.

OK thanks @hiyouga ! i'll put that in the doc page! 

The layer-wise update step seems like it could be applied to any optimizer to reduce the memory usage during training, right? But at the cost of speed perhaps?

Hi @kiddyboots216 @peterjc123 
I followed the implementation of @hiyouga and read a bit the source code of galore, I think the optimizers from the Galore will automatically fall back to the classic optimizers for non-target modules. e.g. if one passes `galore_adaw_8bit_layerwise` the target modules will get optimized with galore whereas other layers will get optimized with adam_8bit, is this understanding correct?

> Hi @kiddyboots216 @peterjc123 I followed the implementation of @hiyouga and read a bit the source code of galore, I think the optimizers from the Galore will automatically fall back to the classic optimizers for non-target modules. e.g. if one passes `galore_adaw_8bit_layerwise` the target modules will get optimized with galore whereas other layers will get optimized with adam_8bit, is this understanding correct?

Yes, I think that is correct but it seems that you has an additional check that disallows per-layer optimizers without galore.

I wonder what's the tradeoff for layerwise optimization in general? Surely this cannot be a free lunch, because you can just use it with AdamW. If this just reduced the VRAM requirement for your gradient in general then everyone would do it. It might be helpful to have some understanding of whether this makes optimization slower, or whether layerwise optimization is fundamentally incompatible with ZeRO.

Testing this for ~10 steps isn't enough, if the default update_proj_gap is 200 you need to test for 200 steps.

You will see that it freezes up for a long time at the 200 step mark because SVD is extremely slow.

Hi everyone! 
Thanks so much for your review! I have addressed the final comments and this is ready for a final review! 

@younesbelkada something is not quite right.  setting `optim_target_modules=["self_attn", "mlp"]` in the training kwargs doesn't match them. setting `optim_target_modules='all_linear'` does seem to work though, so it seems the behavior of `check_target_module_exists` does not work with the `"self_attn", "mlp"` module names.

```
ValueError: None of the target modules were found! (['mlp', 'self_attn']). Please make sure to pass a valid `target_modules`.
```

Nice catch @winglian ! [6ff3762](https://github.com/huggingface/transformers/pull/29588/commits/6ff37620ecac75b4f986cb3dde451ef1ab32ef0e) should fix it, can you try now?

> Nice catch @winglian ! [6ff3762](https://github.com/huggingface/transformers/pull/29588/commits/6ff37620ecac75b4f986cb3dde451ef1ab32ef0e) should fix it, can you try now?

It progresses to start training now, but I've been getting an odd error and I'm not sure if it's a transformers, axolotl, or galore_torch issue.

```
  File "/workspace/transformers/src/transformers/trainer.py", line 1774, in train
    return inner_training_loop(
  File "/workspace/transformers/src/transformers/trainer.py", line 2170, in _inner_training_loop
    self.optimizer.step()
  File "/root/miniconda3/envs/py3.10/lib/python3.10/site-packages/accelerate/optimizer.py", line 145, in step
    self.optimizer.step(closure)
  File "/root/miniconda3/envs/py3.10/lib/python3.10/site-packages/torch/optim/lr_scheduler.py", line 68, in wrapper
    return wrapped(*args, **kwargs)
  File "/root/miniconda3/envs/py3.10/lib/python3.10/site-packages/torch/optim/optimizer.py", line 373, in wrapper
    out = func(*args, **kwargs)
  File "/root/miniconda3/envs/py3.10/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
  File "/root/miniconda3/envs/py3.10/lib/python3.10/site-packages/galore_torch/adamw.py", line 96, in step
    grad = state["projector"].project(grad, state["step"])
  File "/root/miniconda3/envs/py3.10/lib/python3.10/site-packages/galore_torch/galore_projector.py", line 17, in project
    self.ortho_matrix = self.get_orthogonal_matrix(full_rank_grad, self.rank, type='right')
  File "/root/miniconda3/envs/py3.10/lib/python3.10/site-packages/galore_torch/galore_projector.py", line 87, in get_orthogonal_matrix
    A = U[:, :rank] @ torch.diag(s[:rank])
TypeError: slice indices must be integers or None or have an __index__ method
```

So setting `optim_args` doesn't work and results in that error. Am I setting it incorrectly?

```
args = TrainingArguments(
    output_dir="./test-galore",
    max_steps=100,
    per_device_train_batch_size=1,
    optim="galore_adamw",
    optim_args="rank=128",
    optim_target_modules=["attn", "mlp"]
)
```

I believe this is the correct way according to https://github.com/huggingface/transformers/blob/838b87abe231fd70be5132088d0dee72a7bb8d62/src/transformers/trainer.py#L1047-L1051

@younesbelkada the bug is in the `optim_args` string splitting. it casts the value to a string, and galore expects an int. I expect it's going to have to need some fancy checks to determine if it's an int type or float type, since rank expects an int (doesn't work with float), but scale is a float. An easier solution might be to have TrainingArguments.optim_args be a `Union[str, Dict[str, Any]]` instead.

Nice catch @winglian , in [57e7096](https://github.com/huggingface/transformers/pull/29588/commits/57e7096eb9f397dc27e98d2d31470f1fc18436ce) i should properly take care of everything and I can confirm it works fine. The correct way of passing galore args is the following: 
```python
            args = TrainingArguments(
                tmpdir,
                learning_rate=1e-9,
                logging_steps=5,
                optim="galore_adamw",
                optim_args="rank=64, update_proj_gap=100, scale=0.10",
                optim_target_modules=[r".*attn.*", r".*mlp.*"],
            )
```

Hi,
According to the GaLore README:

> We are working on the offical release of GaLore. In the meanwhile, please feel free to try the pre-release version and provide feedback to us. Currently, the pre-release version (e.g., GaLore optimizers) should provide a decent memory reduction and accurate simulation of GaLore algorithm.

When the official version of GaLore is released, will it be integrated into Transformers?

Is there anyway to get this optimizer fully offloaded to the cpu? Similar to how deepspeed does it with a large 32bit adam optimizer. That way the entire system can be working together with little interference through the pcei bus (compared to going into shared memory). Optimizing gpu memory to the max, allowing the gpu to handle the highest batch sizes possible, while the cpu deals with the optimizer.

> Is there anyway to get this optimizer fully offloaded to the cpu? Similar to how deepspeed does it with a large 32bit adam optimizer. That way the entire system can be working together with little interference through the pcei bus (compared to going into shared memory). Optimizing gpu memory to the max, allowing the gpu to handle the highest batch sizes possible, while the cpu deals with the optimizer.

If you use ZeRO stage 2/3 (offload all optimizer state to CPU) as done in deepspeed/FSDP then you will not have any VRAM usage from optimizer state. There's not much point in using GaLoRE if you're already offloading the optimizer state to the cpu, because GaLoRE only reduces the optimizer state. The layerwise backprop isn't part of GaLoRE, it can be done by any optimizer. 

> > Is there anyway to get this optimizer fully offloaded to the cpu? Similar to how deepspeed does it with a large 32bit adam optimizer. That way the entire system can be working together with little interference through the pcei bus (compared to going into shared memory). Optimizing gpu memory to the max, allowing the gpu to handle the highest batch sizes possible, while the cpu deals with the optimizer.
> 
> If you use ZeRO stage 2/3 (offload all optimizer state to CPU) as done in deepspeed/FSDP then you will not have any VRAM usage from optimizer state. There's not much point in using GaLoRE if you're already offloading the optimizer state to the cpu, because GaLoRE only reduces the optimizer state. The layerwise backprop isn't part of GaLoRE, it can be done by any optimizer.

Wait you can do this without having to use adam 32bit cpu? So I could do like adam 8bit fully offloaded to cpu? Last I tried this it didn’t work but I must’ve not done it the way you are saying. cause I’m trying to get the cpu to do the work, not just hold the memory in system ram, (the way paged adam works).


Apologies for this being off topic. To make up for it, I started testing GaLoRE with Unsloth/hf trainer (on windows), I can confirm it works, and i’m getting great results. The loss is what I would normally get, with the added bonus that Its letting me do a batch size of 8 where previously with paged adam 8bit I was at batch size of 4 training a 7b model.

> > > Is there anyway to get this optimizer fully offloaded to the cpu? Similar to how deepspeed does it with a large 32bit adam optimizer. That way the entire system can be working together with little interference through the pcei bus (compared to going into shared memory). Optimizing gpu memory to the max, allowing the gpu to handle the highest batch sizes possible, while the cpu deals with the optimizer.
> > 
> > 
> > If you use ZeRO stage 2/3 (offload all optimizer state to CPU) as done in deepspeed/FSDP then you will not have any VRAM usage from optimizer state. There's not much point in using GaLoRE if you're already offloading the optimizer state to the cpu, because GaLoRE only reduces the optimizer state. The layerwise backprop isn't part of GaLoRE, it can be done by any optimizer.
> 
> Wait you can do this without having to use adam 32bit cpu? So I could do like adam 8bit fully offloaded to cpu? Last I tried this it didn’t work but I must’ve not done it the way you are saying. cause I’m trying to get the cpu to do the work, not just hold the memory in system ram, (the way paged adam works).

Could you try doing this the way it's referenced here https://github.com/TimDettmers/bitsandbytes/issues/89#issuecomment-2006957063 I'm not sure why it wouldn't work

(sorry to go off-topic on a PR with many people and commits on it)