SwinLayer / DonutSwinLayer / ClapAudioLayer attention mask creation always happens on CPU
Author: gorodnitskiy
Pull request: no

### System Info

- `transformers` version: 4.42.0.dev0
- Platform: Linux-5.15.146+-x86_64-with-glibc2.35
- Python version: 3.10.14
- Huggingface_hub version: 0.23.3
- Safetensors version: 0.4.3
- Accelerate version: not installed
- Accelerate config: not found
- PyTorch version (GPU?): 2.3.1+cu121 (False)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using distributed or parallel set-up in script?: no
- GPU: A100 40GB
- CUDA: 12.1

### Who can help?

@amyeroberts

### 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

Hi!

I hope I tagged the right person. If not, I apologize.

I recently used DONUT (OCR-free Document Understanding Transformer) from HuggingFace and noticed that in the DonutSwinLayer (or SwinLayer / ClapAudioLayer), the attention mask is initially created on the CPU, processed, and then moved to the device.


Original code:
```python
class SwinLayer(nn.Module):

    ...

    def get_attn_mask(self, height, width, dtype):
        if self.shift_size > 0:
            # calculate attention mask for SW-MSA
            img_mask = torch.zeros((1, height, width, 1), dtype=dtype)  # <----------------------
            height_slices = (
                slice(0, -self.window_size),
                slice(-self.window_size, -self.shift_size),
                slice(-self.shift_size, None),
            )
            width_slices = (
                slice(0, -self.window_size),
                slice(-self.window_size, -self.shift_size),
                slice(-self.shift_size, None),
            )
            count = 0
            for height_slice in height_slices:
                for width_slice in width_slices:
                    img_mask[:, height_slice, width_slice, :] = count
                    count += 1

            mask_windows = window_partition(img_mask, self.window_size)
            mask_windows = mask_windows.view(-1, self.window_size * self.window_size)
            attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2)
            attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill(attn_mask == 0, float(0.0))
        else:
            attn_mask = None
        return attn_mask

    ...

    def forward(
        self,
        hidden_states: torch.Tensor,
        input_dimensions: Tuple[int, int],
        head_mask: Optional[torch.FloatTensor] = None,
        output_attentions: Optional[bool] = False,
        always_partition: Optional[bool] = False,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        
        ...
        
        attn_mask = self.get_attn_mask(height_pad, width_pad, dtype=hidden_states.dtype)  # <----------------------
        if attn_mask is not None:                                                         # <----------------------
            attn_mask = attn_mask.to(hidden_states_windows.device)                        # <----------------------

        ...

```


Fixed code:
```python
class SwinLayer(nn.Module):

    ...

    def get_attn_mask(self, height, width, dtype, device):  # <-----------------
        if self.shift_size > 0:
            # calculate attention mask for SW-MSA
            img_mask = torch.zeros((1, height, width, 1), dtype=dtype, device=device)  # <-----------------
            height_slices = (
                slice(0, -self.window_size),
                slice(-self.window_size, -self.shift_size),
                slice(-self.shift_size, None),
            )
            width_slices = (
                slice(0, -self.window_size),
                slice(-self.window_size, -self.shift_size),
                slice(-self.shift_size, None),
            )
            count = 0
            for height_slice in height_slices:
                for width_slice in width_slices:
                    img_mask[:, height_slice, width_slice, :] = count
                    count += 1

            mask_windows = window_partition(img_mask, self.window_size)
            mask_windows = mask_windows.view(-1, self.window_size * self.window_size)
            attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2)
            attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill(attn_mask == 0, float(0.0))
        else:
            attn_mask = None
        return attn_mask

    ...

    def forward(
        self,
        hidden_states: torch.Tensor,
        input_dimensions: Tuple[int, int],
        head_mask: Optional[torch.FloatTensor] = None,
        output_attentions: Optional[bool] = False,
        always_partition: Optional[bool] = False,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        
        ...
        
        attn_mask = self.get_attn_mask(                                                            # <-----------------
            height_pad, width_pad, dtype=hidden_states.dtype, device=hidden_states_windows.device  # <-----------------
        )                                                                                          # <-----------------

        ...

```


Benchmark code:
```python
import torch
from time import perf_counter
import numpy as np

from transformers import DonutProcessor, VisionEncoderDecoderModel
from datasets import load_dataset

_TOTAL_ITERATION = 500
_WARMUP = 20


def main() -> None:
    # load processor and model
    processor = DonutProcessor.from_pretrained(
        "naver-clova-ix/donut-base-finetuned-rvlcdip"
    )
    model = VisionEncoderDecoderModel.from_pretrained(
        "naver-clova-ix/donut-base-finetuned-rvlcdip"
    )
    
    model.half()
    model.eval()
    model.to(torch.device("cuda"))
    
    # load test image
    dataset = load_dataset(
        "hf-internal-testing/example-documents", split="test"
    )
    images = dataset[1]["image"]
    
    # first run
    with torch.no_grad():
        task_prompt = "<s_rvlcdip>"
        decoder_input_ids = processor.tokenizer(
            task_prompt, add_special_tokens=False, return_tensors="pt"
        ).input_ids.to(model.decoder.device)
        pixel_values = processor(images, return_tensors="pt").pixel_values.to(
            model.encoder.device, model.encoder.dtype
        )
        _ = model.generate(
            pixel_values,
            decoder_input_ids=decoder_input_ids,
            max_length=model.decoder.config.max_position_embeddings,
            pad_token_id=processor.tokenizer.pad_token_id,
            eos_token_id=processor.tokenizer.eos_token_id,
            use_cache=True,
            bad_words_ids=[[processor.tokenizer.unk_token_id]],
            return_dict_in_generate=True,
        )
    torch.cuda.synchronize()
    
    # main run
    times = []
    times_gen = []
    times_pr = []
    times_tok = []
    for _ in range(_TOTAL_ITERATION):
        torch.cuda.synchronize()
        t1 = perf_counter()

        with torch.no_grad():
            task_prompt = "<s_rvlcdip>"

            torch.cuda.synchronize()
            t1_tok = perf_counter()
            decoder_input_ids = processor.tokenizer(
                task_prompt, add_special_tokens=False, return_tensors="pt"
            ).input_ids.to(model.decoder.device)
            torch.cuda.synchronize()
            t2_tok = perf_counter()
            times_tok.append(t2_tok - t1_tok)
            
            torch.cuda.synchronize()
            t1_pr = perf_counter()
            pixel_values = processor(images, return_tensors="pt").pixel_values.to(
                model.encoder.device, model.encoder.dtype
            )
            torch.cuda.synchronize()
            t2_pr = perf_counter()
            times_pr.append(t2_pr - t1_pr)
            
            torch.cuda.synchronize()
            t1_gen = perf_counter()
            _ = model.generate(
                pixel_values,
                decoder_input_ids=decoder_input_ids,
                max_length=model.decoder.config.max_position_embeddings,
                pad_token_id=processor.tokenizer.pad_token_id,
                eos_token_id=processor.tokenizer.eos_token_id,
                use_cache=True,
                bad_words_ids=[[processor.tokenizer.unk_token_id]],
                return_dict_in_generate=True,
            )
            torch.cuda.synchronize()
            t2_gen = perf_counter()
            times_gen.append(t2_gen - t1_gen)
        
        torch.cuda.synchronize()
        t2 = perf_counter()
        times.append(t2 - t1)
    
    print("processor.tokenizer", np.mean(times_tok[_WARMUP:]), np.std(times_tok[_WARMUP:]))
    print("processor.__call__", np.mean(times_pr[_WARMUP:]), np.std(times_pr[_WARMUP:]))
    print("model.generate",np.mean(times_gen[_WARMUP:]), np.std(times_gen[_WARMUP:]))
    print("Total", np.mean(times[_WARMUP:]), np.std(times[_WARMUP:]))


if __name__ == "__main__":
    main()
```

### Expected behavior

In the case of using a GPU, this issue leads to a significant slowdown in inference speed. By adding a device argument to the first tensor initialization for the attention mask, I achieved a 39% speed improvement for `model.generate` on A100:
- Original:  0.15148±0.00961 s/iter
- Fixed:  0.09294±0.00118 s/iter

I have created this [PR](https://github.com/huggingface/transformers/pull/31295) to fix this issue.

Comments:
Closing as #31295 has been merged in. Thanks again for reporting and for opening a PR to address this so quickly!