BART-base flash_attention_2 causes CUDA error
Author: Kripner
Pull request: no

### System Info

- `transformers` version: 4.37.0
- Platform: Linux-6.5.0-15-generic-x86_64-with-glibc2.31
- Python version: 3.10.13
- Huggingface_hub version: 0.20.3
- Safetensors version: 0.4.1
- Accelerate version: 0.26.1
- Accelerate config: 	not found
- PyTorch version (GPU?): 2.1.2 (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?: Yes
- Using distributed or parallel set-up in script?: No

### Who can help?

@ArthurZucker @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

```
model = AutoModelForSeq2SeqLM.from_pretrained(
    "facebook/bart-base",
    attn_implementation="flash_attention_2",
    torch_dtype=torch.bfloat16,
)
model.to("cuda")
tokenizer = AutoTokenizer.from_pretrained("facebook/bart-base")

def preprocess_function(examples):
    inputs = examples["document"]
    outputs = examples["summary"]
    tokenized_inputs = tokenizer(inputs, max_length=1024, padding="max_length", truncation=True)
    tokenized_outputs = tokenizer(outputs, max_length=64, padding="max_length", truncation=True)
    return {
        "input_ids": tokenized_inputs["input_ids"],
        "attention_mask": tokenized_inputs["attention_mask"],
        "labels": tokenized_outputs["input_ids"],
    }

train_data = datasets.load_dataset("xsum", split="train[:100]", trust_remote_code=True)
train_data = train_data.map(preprocess_function, batched=True, remove_columns=["document", "summary"])

training_args = Seq2SeqTrainingArguments(
    output_dir="output",
)
trainer = Seq2SeqTrainer(
    model=model,
    args=training_args,
    tokenizer=tokenizer,
    train_dataset=train_data,
)
trainer.train()
```

### Expected behavior

Expected behavior as per https://huggingface.co/docs/transformers/en/perf_train_gpu_one#flash-attention-2:

> You can speedup the training throughput by using Flash Attention 2 integration in transformers. Check out the appropriate section in the [single GPU section](https://huggingface.co/docs/transformers/en/perf_infer_gpu_one#Flash-Attention-2) to learn more about how to load a model with Flash Attention 2 modules.

With Bart being listed as supported in the quoted link.

However, the script triggers CUDA error. The full output is (with `CUDA_LAUNCH_BLOCKING=1`):

```
You are attempting to use Flash Attention 2.0 with a model not initialized on GPU. Make sure to move the model to GPU after initializing it on CPU with `model.to('cuda')`.
WARNING:dvclive:Can't save experiment without a Git Repo.
Create a Git repo (`git init`) and commit (`git commit`).
  3%|█████                                                                                                                                                                                                 | 1/39 [00:00<00:13,  2.80it/s]Traceback (most recent call last):
  File "/app/pt/experiments/playground.py", line 100, in <module>
    trainer.train()
  File "/opt/conda/lib/python3.10/site-packages/transformers/trainer.py", line 1539, in train
    return inner_training_loop(
  File "/opt/conda/lib/python3.10/site-packages/transformers/trainer.py", line 1869, in _inner_training_loop
    tr_loss_step = self.training_step(model, inputs)
  File "/opt/conda/lib/python3.10/site-packages/transformers/trainer.py", line 2768, in training_step
    loss = self.compute_loss(model, inputs)
  File "/opt/conda/lib/python3.10/site-packages/transformers/trainer.py", line 2791, in compute_loss
    outputs = model(**inputs)
  File "/opt/conda/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/opt/conda/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "/opt/conda/lib/python3.10/site-packages/transformers/models/bart/modeling_bart.py", line 1731, in forward
    outputs = self.model(
  File "/opt/conda/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/opt/conda/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "/opt/conda/lib/python3.10/site-packages/transformers/models/bart/modeling_bart.py", line 1617, in forward
    decoder_outputs = self.decoder(
  File "/opt/conda/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/opt/conda/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "/opt/conda/lib/python3.10/site-packages/transformers/models/bart/modeling_bart.py", line 1470, in forward
    layer_outputs = decoder_layer(
  File "/opt/conda/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/opt/conda/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "/opt/conda/lib/python3.10/site-packages/transformers/models/bart/modeling_bart.py", line 779, in forward
    hidden_states, cross_attn_weights, cross_attn_present_key_value = self.encoder_attn(
  File "/opt/conda/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/opt/conda/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "/opt/conda/lib/python3.10/site-packages/transformers/models/bart/modeling_bart.py", line 403, in forward
    attn_output = self._flash_attention_forward(
  File "/opt/conda/lib/python3.10/site-packages/transformers/models/bart/modeling_bart.py", line 454, in _flash_attention_forward
    attn_output_unpad = flash_attn_varlen_func(
  File "/opt/conda/lib/python3.10/site-packages/flash_attn/flash_attn_interface.py", line 1059, in flash_attn_varlen_func
    return FlashAttnVarlenFunc.apply(
  File "/opt/conda/lib/python3.10/site-packages/torch/autograd/function.py", line 539, in apply
    return super().apply(*args, **kwargs)  # type: ignore[misc]
  File "/opt/conda/lib/python3.10/site-packages/flash_attn/flash_attn_interface.py", line 576, in forward
    out, q, k, v, out_padded, softmax_lse, S_dmask, rng_state = _flash_attn_varlen_forward(
  File "/opt/conda/lib/python3.10/site-packages/flash_attn/flash_attn_interface.py", line 85, in _flash_attn_varlen_forward
    out, q, k, v, out_padded, softmax_lse, S_dmask, rng_state = flash_attn_cuda.varlen_fwd(
RuntimeError: CUDA error: invalid configuration argument
Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.

  3%|▎         | 1/39 [00:00<00:20,  1.82it/s]
```

Comments:
Hey! Thanks for raising the issue, this seems to be a hardware related issue see example issues: https://github.com/nerfstudio-project/nerfacc/issues/207 nothing much we can do on our side I believe. FYI @younesbelkada 

Hmm not sure what is wrong here, can you try to train the model without padd tokens by packing examples together?

I had the same problem. Did you solve it

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.