PEFT+gradient checkpointing causes attention mask shape mismatch during backward pass
Author: geoffreyangus
Pull request: no

### System Info

- `transformers` version: 4.36.0
- Platform: Linux-5.4.0-152-generic-x86_64-with-glibc2.17
- Python version: 3.8.13
- Huggingface_hub version: 0.19.4
- Safetensors version: 0.4.1
- Accelerate version: 0.24.1
- Accelerate config: 	not found
- PyTorch version (GPU?): 2.1.1+cu121 (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?: x1 80GB A100
- Using distributed or parallel set-up in script?: No

### Who can help?

@ArthurZucker and @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

This should reproduce the error. This was originally run on x1 80GB A100. Note that both Llama-2-7B and Mixtral-8x7B are affected by this change (mixtral is testable– commented out– in the repro script below).

```python
import torch
from torch.optim import Adam
from transformers import BitsAndBytesConfig
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import get_peft_config, get_peft_model, LoraConfig, TaskType

MODEL_ID = "meta-llama/Llama-2-7b-hf"
# MODEL_ID = "mistralai/Mixtral-8x7B-v0.1"  # this is broken too

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
inputs = tokenizer("hello world what's up", return_tensors="pt")
inputs = {k: v.to("cuda") for k, v in inputs.items()}
print(inputs)

model = AutoModelForCausalLM.from_pretrained(MODEL_ID, device_map="auto", attn_implementation="eager", torch_dtype=torch.float16)
peft_config = LoraConfig(task_type=TaskType.CAUSAL_LM, target_modules=['q_proj', 'v_proj'], inference_mode=False, r=8, lora_alpha=32, lora_dropout=0.1)
model = get_peft_model(model, peft_config)
model.print_trainable_parameters()
model.gradient_checkpointing_enable()
model.enable_input_require_grads()

optimizer = Adam(model.parameters(), lr=1e-5)
model.train()

for i in range(10):
  outputs = model(labels=inputs['input_ids'], **inputs)
  loss = outputs.loss
  print(loss)
  loss.backward()
  optimizer.step()
  optimizer.zero_grad()
```

The culprit in the above script is the `model.train()` call after LoRA is configured for `model`. One can workaround it by (1) avoiding calling `model.train()` and (2) if you have to call `model.eval()`, be sure to save and reuse the `module.training` values from the model's initial state when reverting back to train mode.

The above will throw the following error:
```
  File "mixtral_train_debug.py", line 47, in <module>
    loss.backward()
  File "/home/ray/anaconda3/lib/python3.8/site-packages/torch/_tensor.py", line 492, in backward
    torch.autograd.backward(
  File "/home/ray/anaconda3/lib/python3.8/site-packages/torch/autograd/__init__.py", line 251, in backward
    Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
  File "/home/ray/anaconda3/lib/python3.8/site-packages/torch/autograd/function.py", line 288, in apply
    return user_fn(self, *args)
  File "/home/ray/anaconda3/lib/python3.8/site-packages/torch/utils/checkpoint.py", line 271, in backward
    outputs = ctx.run_function(*detached_inputs)
  File "/home/ray/anaconda3/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/ray/anaconda3/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/ray/anaconda3/lib/python3.8/site-packages/transformers/models/llama/modeling_llama.py", line 789, in forward
    hidden_states, self_attn_weights, present_key_value = self.self_attn(
  File "/home/ray/anaconda3/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/ray/anaconda3/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/ray/anaconda3/lib/python3.8/site-packages/transformers/models/llama/modeling_llama.py", line 423, in forward
    raise ValueError(
ValueError: Attention mask should be of size (1, 1, 7, 14), but is torch.Size([1, 1, 7, 7])
```

### Expected behavior

I would expect calling `model.train()` and `model.eval()` to be callable despite the presence of PEFT modules and/or gradient checkpointing.

Comments:
hi @geoffreyangus 
thanks for the clean reproducer, I can confirm #28031 fixes the bug

Ah, @younesbelkada were you able to try with Mixtral as well (the commented out MODEL_ID in the reproducer)? It still seems to be broken on my side.

Hi @geoffreyangus 
Thanks, indeed, just made #28061