use_reentrant=False can't be set properly
Author: getao
Pull request: no
Labels: name: trainer

name: DeepSpeed

### System Info

transformers==4.40.1
deepspeed==0.14.2
torch==2.2.1

### Who can help?

@ArthurZucker Hello, I used the tranformers' trainer with deepspeed to train the decoder-only model. As a common solution to reduce memory, I enabled gradient checkpointing and set use_reentrance to False in my code:
```
    parser = transformers.HfArgumentParser((ModelArguments, DataArguments, TrainingArguments))
    model_args, data_args, training_args = parser.parse_args_into_dataclasses()
    training_args.gradient_checkpointing_kwargs = {"use_reentrant": False}
 ```
 When I printed the training_args, it shows propoerly:
```
...
gradient_checkpointing=True,
gradient_checkpointing_kwargs={'use_reentrant': False},
...
```

The training_args is properly passed to Trainer with:
```
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset,
        eval_dataset=eval_dataset,
        data_collator=data_collator
    )
```

However, when training starts, I was still always warned with the message:
```
warnings.warn(
/opt/conda/lib/python3.10/site-packages/torch/utils/checkpoint.py:460: UserWarning: torch.utils.checkpoint: please pass in use_reentrant=True or use_reentrant=False explicitly. The default value of use_reentrant will be updated to be False in the future. To maintain current behavior, pass use_reentrant=True. It is recommended that you use use_reentrant=False. Refer to docs for more details on the differences between the two variants.
```

I don't understand the reason leading to the message but it seems that ```use_reentrant``` is not properly set to take effect.

Could anyone please help me take a look at the problem?

### 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

```
def train_model(model, train_dataset, eval_dataset, training_args, data_collator=None):

    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset,
        eval_dataset=eval_dataset,
        data_collator=data_collator
    )

def main():
    parser = transformers.HfArgumentParser((ModelArguments, DataArguments, TrainingArguments))
    model_args, data_args, training_args = parser.parse_args_into_dataclasses()
    training_args.gradient_checkpointing_kwargs = {"use_reentrant": False}

    tokenizer = AutoTokenizer.from_pretrained(model_args.model_name_or_path, trust_remote_code=True)

    train_dataset, eval_dataset = load_data(...)

    train_dataset = train_dataset.map(
        group_texts,
        batched=True,
        fn_kwargs={"block_size": data_args.max_seq_length},
        load_from_cache_file=True
    )
    eval_dataset = eval_dataset.map(
        group_texts,
        batched=True,
        fn_kwargs={"block_size": data_args.max_seq_length},
        load_from_cache_file=True
    )

    model = AutoModelForCausalLM.from_pretrained(model_args.model_name_or_path, torch_dtype=torch.float16 if not training_args.bf16 else torch.bfloat16, resume_download=True, trust_remote_code=True, attn_implementation="flash_attention_2")

    train_model(model, train_dataset, eval_dataset, training_args)
```

The some configs of deepspeed are as follows:

```
    "zero_optimization": {
        "stage": 3,
        "allgather_partitions": true,
        "allgather_bucket_size": 5e7,
        "overlap_comm": true,
        "reduce_scatter": true,
        "reduce_bucket_size": 5e7,
        "contiguous_gradients" : true,
        "stage3_max_live_parameters" : 1e8,
        "stage3_max_reuse_distance" : 1e8,
        "stage3_prefetch_bucket_size" : 5e7,
        "stage3_param_persistence_threshold" : 1e7,
        "sub_group_size" : 1e9,
        "stage3_gather_16bit_weights_on_model_save": true
    },

    "activation_checkpointing": {
        "partition_activations": true,
        "contiguous_memory_optimization": false,
        "cpu_checkpointing": false,
        "number_checkpoints": null,
        "synchronize_checkpoint_boundary": false,
        "profile": false
    },
```

### Expected behavior

No warning message

Comments:
Hi there! It seems like you are encountering a warning related to the `use_reentrant` parameter when training your model with deepspeed and transformers. The warning is advising you to explicitly pass in `use_reentrant=True` or `use_reentrant=False` to the torch checkpoint function.

In your code snippet, you have set `use_reentrant` to `False` in the `gradient_checkpointing_kwargs`, but the warning indicates that it needs to be explicitly passed when using torch's checkpointing mechanism.

To address this warning and ensure that `use_reentrant` is properly set, you can explicitly pass this argument to the `Trainer` constructor where you instantiate the trainer, like so:

```python
checkpointing_args = {"use_reentrant": False}

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    data_collator=data_collator,
    gradient_checkpointing=True,
    gradient_checkpointing_kwargs=checkpointing_args
)
```

By explicitly passing `gradient_checkpointing=True` and `gradient_checkpointing_kwargs=checkpointing_args`, you are ensuring that the `use_reentrant` argument is set correctly and should prevent the warning message from appearing during training.

If you have any further questions or if the issue persists, feel free to ask for more assistance!

> Hi there! It seems like you are encountering a warning related to the `use_reentrant` parameter when training your model with deepspeed and transformers. The warning is advising you to explicitly pass in `use_reentrant=True` or `use_reentrant=False` to the torch checkpoint function.
> 
> In your code snippet, you have set `use_reentrant` to `False` in the `gradient_checkpointing_kwargs`, but the warning indicates that it needs to be explicitly passed when using torch's checkpointing mechanism.
> 
> To address this warning and ensure that `use_reentrant` is properly set, you can explicitly pass this argument to the `Trainer` constructor where you instantiate the trainer, like so:
> 
> ```python
> checkpointing_args = {"use_reentrant": False}
> 
> trainer = Trainer(
>     model=model,
>     args=training_args,
>     train_dataset=train_dataset,
>     eval_dataset=eval_dataset,
>     data_collator=data_collator,
>     gradient_checkpointing=True,
>     gradient_checkpointing_kwargs=checkpointing_args
> )
> ```
> 
> By explicitly passing `gradient_checkpointing=True` and `gradient_checkpointing_kwargs=checkpointing_args`, you are ensuring that the `use_reentrant` argument is set correctly and should prevent the warning message from appearing during training.
> 
> If you have any further questions or if the issue persists, feel free to ask for more assistance!

TypeError: Trainer.__init__() got an unexpected keyword argument 'gradient_checkpointing_kwargs'

It doesn't work.

BTW, are you a model?

@cw235 There are a series of issues you've commented on e.g. [here](https://github.com/huggingface/transformers/issues/30731#issuecomment-2105375405), [here](https://github.com/huggingface/transformers/issues/30750#issuecomment-2105373137) and [here](https://github.com/huggingface/transformers/issues/30725#issuecomment-2105373583) in which the issue has clearly been fed through to a chat model e.g. ChatGPT to produce an output which does not answer or address the question. Please refrain from doing this, it is both unhelpful to the person reporting the issue and isn't scalable behaviour: what would happen if everyone started doing this? 

All of these will be marked as spam so as to hide them and make the issues navigable. If you keep doing this, the account @cw235 will be reported. 

cc @pacman100 @muellerzr 

I couldn't recreate this warning. Can you try either:

1. Using the dev version of transformers `pip install git+https://github.com/huggingface/transformers`
2. Upgrading your PyTorch version `pip install torch -U` 

And let us know if you get this warning? 

Thanks!

> I couldn't recreate this warning. Can you try either:
> 
> 1. Using the dev version of transformers `pip install git+https://github.com/huggingface/transformers`
> 2. Upgrading your PyTorch version `pip install torch -U`
> 
> And let us know if you get this warning?
> 
> Thanks!

Thank you. Is it possible to be related to specific models with the argument ```trust_remote_code=True```?

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.