Galore ignores provided learning rate
Author: jakubgajski
Pull request: no
Labels: name: optimization

### System Info

- `transformers` version: 4.39.3
- Platform: Linux-6.8.0-76060800daily20240311-generic-x86_64-with-glibc2.35
- Python version: 3.10.13
- Huggingface_hub version: 0.20.3
- Safetensors version: 0.4.2
- Accelerate version: 0.26.1
- Accelerate config: 	not found
- PyTorch version (GPU?): 2.2.0 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: True
- Using distributed or parallel set-up in script?: False

### Who can help?

@younesbelkada

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [X] My own task or dataset (give details below)

### Reproduction

I'm trying to finetune (extended with new tokens) Mistral 7B with Galore, but no matter where I pass the learning rate, trainer is using 0.001.

model setup:
```python
# source model is a path to Mistral 7B with some extra indices in embedding matrix.
model = MistralForCausalLM.from_pretrained(source_model,
                                           torch_dtype=torch.bfloat16,
                                           attn_implementation='flash_attention_2')

# I don't think it is important, but I will leave it here just in case:
# freeze embeddings that do not need gradient, train only new tokens
input_emb = model.get_input_embeddings()
input_emb.weight.register_hook(lambda grad: zero_grad(grad, list(range(0, 3200))))

model.set_input_embeddings(input_emb)

training_args = TrainingArguments(output_dir=output_model,
                                  evaluation_strategy="steps", warmup_ratio=0.05,
                                  overwrite_output_dir=True, gradient_accumulation_steps=1,
                                  per_device_train_batch_size=4, per_device_eval_batch_size=1,
                                  num_train_epochs=3, bf16_full_eval=True, bf16=True,
                                  eval_steps=1000, save_strategy='steps', save_steps=1000,
                                  save_total_limit=5, dataloader_num_workers=8,
                                  load_best_model_at_end=True,
                                  neftune_noise_alpha=0, 
                                  prediction_loss_only=True,
                                  gradient_checkpointing=True,
                                  optim="galore_adamw_8bit_layerwise",
                                  optim_target_modules=[
                                      'q_proj', 'k_proj', 'down_proj', 'up_proj',
                                      'gate_proj', 'v_proj', 'o_proj', 'lm_head'
                                  ],
                                  optim_args="""rank=512, update_proj_gap=500, scale=0.25, lr=0.0002""",
                                  learning_rate=0.0002,
                                  )

trainer = Trainer(
    model=model,
    args=training_args,
    eval_dataset=validation,
    train_dataset=dataset,
    data_collator=identity_collator)
```

Yet, stdout logs shows: {'loss': 4.6255, 'grad_norm': 0.0, 'learning_rate': 0.001, 'epoch': 0.01}, while I'm passing learning rate to both directly TrainingArguments and Galore's optim_args.

### Expected behavior

Learning rate shall be the same as passed to TrainingArguments via learning_rate argument.

Comments:
Hey @jakubgajski, I looked a bit into it myself and found the issue to be in the creation of the optimizer. More specifically, the learning rate isn't passed but the default value `1e-3` is always used (see [this line](https://github.com/huggingface/transformers/blob/76fa17c1663a0efeca7208c20579833365584889/src/transformers/trainer_pt_utils.py#L1282)).

To fix it locally, just change the line from
`super().__init__([dummy_tensor], {"lr": 1e-03})` to
`super().__init__([dummy_tensor], {"lr": kwargs.get("lr", 1e-03)})`.

I'll submit a PR later today, but not sure when I have the time. Hope that helps.

Thanks @vasqu for your time, it helps :)

Thanks again @vasqu for the fix ! 🤩