adalomo and deepspeed zero3 offload error
Author: YooSungHyun
Pull request: no
Labels: name: bug

### System Info

python==3.11.11
transformers==4.48.1
torch==2.5.1 
deepspeed==0.16.3

### Who can help?

@muellerzr @SunMarc

### Information

- [x] The official example scripts
- [ ] My own modified scripts

### Tasks

- [x] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

use adalomo & deepspeed zero3 param, optimizer offload then,
[rank7]: Traceback (most recent call last):
[rank7]:   File "/data/shyu/LLM-train/train_sft.py", line 443, in <module>
[rank7]:     main(script_args, training_args, model_args)
[rank7]:   File "/data/shyu/LLM-train/train_sft.py", line 423, in main
[rank7]:     trainer.train()
[rank7]:   File "/data/shyu/.train/lib/python3.11/site-packages/transformers/trainer.py", line 2171, in train
[rank7]:     return inner_training_loop(
[rank7]:            ^^^^^^^^^^^^^^^^^^^^
[rank7]:   File "/data/shyu/.train/lib/python3.11/site-packages/transformers/trainer.py", line 2531, in _inner_training_loop
[rank7]:     tr_loss_step = self.training_step(model, inputs, num_items_in_batch)
[rank7]:                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank7]:   File "/data/shyu/.train/lib/python3.11/site-packages/transformers/trainer.py", line 3712, in training_step
[rank7]:     self.accelerator.backward(loss, **kwargs)
[rank7]:   File "/data/shyu/.train/lib/python3.11/site-packages/accelerate/accelerator.py", line 2238, in backward
[rank7]:     self.deepspeed_engine_wrapped.backward(loss, **kwargs)
[rank7]:   File "/data/shyu/.train/lib/python3.11/site-packages/accelerate/utils/deepspeed.py", line 261, in backward
[rank7]:     self.engine.backward(loss, **kwargs)
[rank7]:   File "/data/shyu/.train/lib/python3.11/site-packages/deepspeed/utils/nvtx.py", line 18, in wrapped_fn
[rank7]:     ret_val = func(*args, **kwargs)
[rank7]:               ^^^^^^^^^^^^^^^^^^^^^
[rank7]: TypeError: DeepSpeedEngine.backward() got an unexpected keyword argument 'learning_rate'

### Expected behavior

run training step well

Comments:
I think the `if` condition in the `backward` section of `accelerate` should also check for the presence of `learning_rate`, or the issue should already be handled within `self.deepspeed_engine_wrapped.backward(loss, **kwargs)`. It seems like either the `if` condition is missing or the handling is omitted inside `deepspeed_engine_wrapped`. Is there a possible solution for this?

```python3
    def backward(self, loss, **kwargs):
        """
        Scales the gradients in accordance to the `GradientAccumulationPlugin` and calls the correct `backward()` based
        on the configuration.

        Should be used in lieu of `loss.backward()`.

        Example:

        ```python
        >>> from accelerate import Accelerator

        >>> accelerator = Accelerator(gradient_accumulation_steps=2)
        >>> outputs = model(inputs)
        >>> loss = loss_fn(outputs, labels)
        >>> accelerator.backward(loss)
        ```
        """
        learning_rate = kwargs.get("learning_rate")

        if self.distributed_type != DistributedType.DEEPSPEED:
            # deepspeed handles loss scaling by gradient_accumulation_steps in its `backward`
            loss = loss / self.gradient_accumulation_steps
        if self.distributed_type == DistributedType.DEEPSPEED:
            self.deepspeed_engine_wrapped.backward(loss, **kwargs)
        elif self.distributed_type == DistributedType.MEGATRON_LM:
            return
        elif self.scaler is not None:
            self.scaler.scale(loss).backward(**kwargs)
        elif learning_rate is not None and self.has_lomo_optimizer:
            self.lomo_backward(loss, learning_rate)
        else:
            loss.backward(**kwargs)
```


hi @YooSungHyun 

Now, DeepSpeedEngine_backward() function doesn't support the 'learning_rate' keyword argument. You need to remove this argument from your code. 

Here are the steps to resolve this issue:

Step 1: Identify the lines of your code where you are calling the DeepSpeedEngine_backward() function with 'learning_rate' keyword argument. 

Step 2: Remove the 'learning_rate' keyword argument from these lines of code. For example, if you have a line like this:
```
self.accelerator.backward(loss, learning_rate=self.lr)
```
You should change it to:
```
self.accelerator.backward(loss)
```
Step 3: Once you've removed the 'learning_rate' keyword argument, your code should run without throwing the TypeError.

If you are using Hugging Face's Trainer class, you might have to set the learning rate outside of the training loop. For example, you can use the AdamW optimizer like this:
```
from transformers import AdamW

optimizer = AdamW(model.parameters(), lr=1e-5)
```
In this case, you need to make sure that the 'lr' argument in AdamW is not being passed to DeepSpeedEngine_backward() function.

> hi [@YooSungHyun](https://github.com/YooSungHyun)
> 
> Now, DeepSpeedEngine_backward() function doesn't support the 'learning_rate' keyword argument. You need to remove this argument from your code.
> 
> Here are the steps to resolve this issue:
> 
> Step 1: Identify the lines of your code where you are calling the DeepSpeedEngine_backward() function with 'learning_rate' keyword argument.
> 
> Step 2: Remove the 'learning_rate' keyword argument from these lines of code. For example, if you have a line like this:
> 
> ```
> self.accelerator.backward(loss, learning_rate=self.lr)
> ```
> 
> You should change it to:
> 
> ```
> self.accelerator.backward(loss)
> ```
> 
> Step 3: Once you've removed the 'learning_rate' keyword argument, your code should run without throwing the TypeError.
> 
> If you are using Hugging Face's Trainer class, you might have to set the learning rate outside of the training loop. For example, you can use the AdamW optimizer like this:
> 
> ```
> from transformers import AdamW
> 
> optimizer = AdamW(model.parameters(), lr=1e-5)
> ```
> 
> In this case, you need to make sure that the 'lr' argument in AdamW is not being passed to DeepSpeedEngine_backward() function.

Hi bro, thx for comments. but, i use adalomo optimizer... not adamw.

so, If it falls under `adalomo`, the `if` statement should trigger `lomo_backward`, but since it's DeepSpeed, it's hitting the wrong function, which is the issue I pointed out, my friend.

cc @muellerzr 

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.