Peft + gradient checkpointing crashes
Author: snailrowen1337
Pull request: no

### System Info

>>> transformers.__version__
'4.37.0.dev0'

>>> peft.__version__
'0.7.2.dev0'

>>> torch.__version__
'1.13.0'

### Who can help?

_No response_

### Information

- [ ] The official example scripts
- [X] 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

```
from functools import partial

import torch

from datasets import load_dataset
import transformers
from peft import LoraConfig, get_peft_model, TaskType
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    DataCollatorForSeq2Seq,
    TrainingArguments,
    set_seed,
    Trainer,
)



def prepare_datasets(raw_datasets, train_key, tokenizer, max_seq):
    # Preprocessing the datasets.
    if "messages" in raw_datasets[train_key].column_names:
        encode_function = partial(
            encode_with_messages_format,
            tokenizer=tokenizer,
            max_seq_length=max_seq,
        )
    else:
        raise ValueError("You need to have either 'prompt'&'completion' or 'messages' in your column names.")
    # To speed up this part, we use multiprocessing.
    lm_datasets = raw_datasets.map(encode_function, batched=False, num_proc=32)
    lm_datasets.set_format(type="pt")
    lm_datasets = lm_datasets.filter(lambda example: (example['labels'] != -100).any())
    return lm_datasets

def encode_with_messages_format(example, tokenizer, max_seq_length):
    '''
    Here we assume each example has a 'messages' field Each message is a dict with 'role' and 'content' fields.
    We concatenate all messages with the roles as delimiters and tokenize them together.
    '''
    messages = example['messages']
    if len(messages) == 0:
        raise ValueError('messages field is empty.')
    
    def _concat_messages(messages):
        message_text = ""
        for message in messages:
            if message["role"] == "system":
                message_text += "
\n" + message["content"].strip() + "\n"
            elif message["role"] == "user":
                message_text += "
\n" + message["content"].strip() + "\n"
            elif message["role"] == "assistant":
                message_text += "
\n" + message["content"].strip() + tokenizer.eos_token + "\n"
            else:
                raise ValueError("Invalid role: {}".format(message["role"]))
        return message_text
        
    example_text = _concat_messages(messages).strip()
    tokenized_example = tokenizer(example_text, return_tensors='pt', max_length=max_seq_length, truncation=True)
    input_ids = tokenized_example.input_ids
    labels = input_ids.clone()

    # mask the non-assistant part for avoiding loss
    for message_idx, message in enumerate(messages):
        if message["role"] != "assistant":
            if message_idx == 0:
                message_start_idx = 0
            else:
                message_start_idx = tokenizer(
                    _concat_messages(messages[:message_idx]), return_tensors='pt', max_length=max_seq_length, truncation=True
                ).input_ids.shape[1]
            if message_idx < len(messages) - 1 and messages[message_idx+1]["role"] == "assistant":
                # here we also ignore the role of the assistant
                messages_so_far = _concat_messages(messages[:message_idx+1]) + "
\n"
            else:
                messages_so_far = _concat_messages(messages[:message_idx+1])
            message_end_idx = tokenizer(
                messages_so_far,
                return_tensors='pt', 
                max_length=max_seq_length, 
                truncation=True
            ).input_ids.shape[1]
            labels[:, message_start_idx:message_end_idx] = -100
            
            if message_end_idx >= max_seq_length:
                break

    attention_mask = torch.ones_like(input_ids)
    return {
        'input_ids': input_ids.flatten(),
        'labels': labels.flatten(),
        'attention_mask': attention_mask.flatten(),
    }

def get_dataset(dataset_name, train_key, tokenizer, max_seq):
    raw_datasets = load_dataset(dataset_name)
    lm_datasets = prepare_datasets(raw_datasets, train_key, tokenizer, max_seq)
    train_dataset = lm_datasets[train_key]
    eval_dataset = lm_datasets['test_sft']
    return train_dataset, eval_dataset

def main():
    dataset_name = "HuggingFaceH4/ultrachat_200k"
    model_name = "NousResearch/Yarn-Mistral-7b-128k"
    train_key = 'train_sft'
    max_seq = 512
    set_seed(123)
    
    model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True, use_cache=False)
    model.enable_input_require_grads()
    peft_config = LoraConfig(inference_mode=False, r=8, lora_alpha=32, lora_dropout=0.1, task_type=TaskType.CAUSAL_LM)
    model = get_peft_model(model, peft_config)
    
    # tokenizer
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    tokenizer.add_special_tokens({'pad_token': '[PAD]'})
    embedding_size = model.get_input_embeddings().weight.shape[0]
    if len(tokenizer) > embedding_size:
        model.resize_token_embeddings(len(tokenizer))

    train_dataset, eval_dataset = get_dataset(dataset_name, train_key, tokenizer, max_seq)

    training_args = TrainingArguments(
        output_dir="output",
        report_to="tensorboard",
        per_device_train_batch_size=1,
        gradient_accumulation_steps=4,
        learning_rate=7e-5,
        logging_steps=1,
        num_train_epochs=5,
        max_steps=-1,
        save_steps=100,
        save_total_limit=10,
        warmup_ratio=0.05,
        lr_scheduler_type='cosine',
        evaluation_strategy ='steps',
        eval_steps=250,
        per_device_eval_batch_size=2,
        gradient_checkpointing=True,
    )
        
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset,
        eval_dataset=eval_dataset,
        tokenizer=tokenizer,
        data_collator=DataCollatorForSeq2Seq(tokenizer=tokenizer, model=model),
    )
    train_result = trainer.train()
    trainer.save_model()  # Saves the tokenizer too for easy upload


if __name__ == "__main__":
    main()
```

### Expected behavior

Training should not crash. Instead, I get

```
  warnings.warn("None of the inputs have requires_grad=True. Gradients will be None")
Traceback (most recent call last):
  File "repro.py", line 154, in <module>
    main()
  File "repro.py", line 149, in main
    train_result = trainer.train()
  File "/opt/conda/envs/ptca/lib/python3.8/site-packages/transformers/trainer.py", line 1543, in train
    return inner_training_loop(
  File "/opt/conda/envs/ptca/lib/python3.8/site-packages/transformers/trainer.py", line 1860, in _inner_training_loop
    tr_loss_step = self.training_step(model, inputs)
  File "/opt/conda/envs/ptca/lib/python3.8/site-packages/transformers/trainer.py", line 2746, in training_step
    self.accelerator.backward(loss)
  File "/opt/conda/envs/ptca/lib/python3.8/site-packages/accelerate/accelerator.py", line 1905, in backward
    loss.backward(**kwargs)
  File "/opt/conda/envs/ptca/lib/python3.8/site-packages/torch/_tensor.py", line 487, in backward
    torch.autograd.backward(
  File "/opt/conda/envs/ptca/lib/python3.8/site-packages/torch/autograd/__init__.py", line 197, in backward
    Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn
  0%|           
```

Comments:
hey! Could you try upgrading your torch version? 1.13.0 is a bit old. 
cc @younesbelkada 

@ArthurZucker  -- I've updated to the latest pytorch but still face the issue

Alright, now you are using a remote model so would recommend you to open the issue on `https://huggingface.co/NousResearch/Yarn-Mistral-7b-128k/discussions` unless this also fails with `mistralai/Mistral-7B-v0.1` 🤗 

@ArthurZucker  I still face this issue with the mistral model. Other users have seen similar issues [https://github.com/huggingface/peft/issues/137] but these fixes do not work for me.

CC @younesbelkada 

Hi @snailrowen1337 
The error should disappear if you pass `gradient_checkpointing_kwargs={"use_reentrant":False}` in `TrainingArguments`

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.