Repository: transformers
Issue: meta-llama/Llama-3.2-11B-Vision-Instruct, device_map = 'auto', padding ruins _prepare_4d_causal_attention_mask_with_cache_position
Labels: bug

### System Info

### OS info 
- `transformers` version: 4.48.1
- Platform: Linux-5.4.0-192-generic-x86_64-with-glibc2.31
- Python version: 3.12.0
- Huggingface_hub version: 0.26.5
- Safetensors version: 0.4.5
- Accelerate version: 1.2.0
- Accelerate config:    not found
- PyTorch version (GPU?): 2.4.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 parallel device setup
- Using GPU in script (8 gpu's in my case)
- GPU type: NVIDIA GeForce RTX 3090

### Who can help?

@amyeroberts , @qubvel @ArthurZucker 

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

## Steps to replicate

```python
import requests
from PIL import Image
from transformers import MllamaForConditionalGeneration, AutoProcessor


url = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/0052a70beed5bf71b92610a43a52df6d286cd5f3/diffusers/rabbit.jpg"
image = Image.open(requests.get(url, stream=True).raw)
texts = ["

What do you see here?", "

What do you see here but longer?"]
images = [image, image]


repo_id = "meta-llama/Llama-3.2-11B-Vision-Instruct"

processor = AutoProcessor.from_pretrained(repo_id)
model = MllamaForConditionalGeneration.from_pretrained(repo_id, device_map = 'auto')


batch = processor(text=texts, images=images, return_tensors="pt", padding=True)

with torch.no_grad():
    model_output = model(
        input_ids = batch['input_ids'],
        attention_mask = batch['attention_mask'], 
        pixel_values = batch['pixel_values'], 
        aspect_ratio_ids = batch['aspect_ratio_ids'], 
        aspect_ratio_mask = batch['aspect_ratio_mask'], 
        cross_attention_mask = batch['cross_attention_mask'],
        )
```


### Full traceback: 
```
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
Cell In[8], line 2
      1 with torch.no_grad():
----> 2     model_output = model(
      3         input_ids = batch['input_ids'],
      4         attention_mask = batch['attention_mask'], 
      5         pixel_values = batch['pixel_values'], 
      6         aspect_ratio_ids = batch['aspect_ratio_ids'], 
      7         aspect_ratio_mask = batch['aspect_ratio_mask'], 
      8         cross_attention_mask = batch['cross_attention_mask'],
      9         # labels = batch['labels']
     10         )
     12     # model_output = model.generate(**batch, max_new_tokens = 64) 

File ~/mtrix2/mt-rix/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py:1553, in Module._wrapped_call_impl(self, *args, **kwargs)
   1551     return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   1552 else:
-> 1553     return self._call_impl(*args, **kwargs)

File ~/mtrix2/mt-rix/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py:1562, in Module._call_impl(self, *args, **kwargs)
   1557 # If we don't have any hooks, we want to skip the rest of the logic in
   1558 # this function, and just call forward.
   1559 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
   1560         or _global_backward_pre_hooks or _global_backward_hooks
   1561         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1562     return forward_call(*args, **kwargs)
   1564 try:
   1565     result = None

File ~/mtrix2/mt-rix/.venv/lib/python3.12/site-packages/accelerate/hooks.py:170, in add_hook_to_module.<locals>.new_forward(module, *args, **kwargs)
    168         output = module._old_forward(*args, **kwargs)
    169 else:
--> 170     output = module._old_forward(*args, **kwargs)
    171 return module._hf_hook.post_forward(module, output)

File ~/mtrix2/mt-rix/.venv/lib/python3.12/site-packages/transformers/models/mllama/modeling_mllama.py:2131, in MllamaForConditionalGeneration.forward(self, input_ids, pixel_values, aspect_ratio_mask, aspect_ratio_ids, attention_mask, cross_attention_mask, cross_attention_states, position_ids, past_key_values, inputs_embeds, labels, use_cache, output_attentions, output_hidden_states, return_dict, cache_position, num_logits_to_keep)
   2128     cross_attention_mask = cross_attention_mask[:, :, cache_position]
   2129     full_text_row_masked_out_mask = full_text_row_masked_out_mask[:, :, cache_position]
-> 2131 outputs = self.language_model(
   2132     input_ids=input_ids,
   2133     attention_mask=attention_mask,
   2134     position_ids=position_ids,
   2135     cross_attention_states=cross_attention_states,
   2136     cross_attention_mask=cross_attention_mask,
   2137     full_text_row_masked_out_mask=full_text_row_masked_out_mask,
   2138     past_key_values=past_key_values,
   2139     use_cache=use_cache,
   2140     inputs_embeds=inputs_embeds,
   2141     labels=labels,
   2142     output_hidden_states=output_hidden_states,
   2143     output_attentions=output_attentions,
   2144     return_dict=return_dict,
   2145     cache_position=cache_position,
   2146     num_logits_to_keep=num_logits_to_keep,
   2147 )
   2149 return outputs

File ~/mtrix2/mt-rix/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py:1553, in Module._wrapped_call_impl(self, *args, **kwargs)
   1551     return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   1552 else:
-> 1553     return self._call_impl(*args, **kwargs)

File ~/mtrix2/mt-rix/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py:1562, in Module._call_impl(self, *args, **kwargs)
   1557 # If we don't have any hooks, we want to skip the rest of the logic in
   1558 # this function, and just call forward.
   1559 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
   1560         or _global_backward_pre_hooks or _global_backward_hooks
   1561         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1562     return forward_call(*args, **kwargs)
   1564 try:
   1565     result = None

File ~/mtrix2/mt-rix/.venv/lib/python3.12/site-packages/transformers/models/mllama/modeling_mllama.py:1939, in MllamaForCausalLM.forward(self, input_ids, attention_mask, position_ids, cross_attention_states, cross_attention_mask, full_text_row_masked_out_mask, past_key_values, inputs_embeds, labels, use_cache, output_attentions, output_hidden_states, return_dict, cache_position, num_logits_to_keep, **loss_kwargs)
   1936 return_dict = return_dict if return_dict is not None else self.config.use_return_dict
   1938 # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
-> 1939 outputs = self.model(
   1940     input_ids=input_ids,
   1941     cross_attention_states=cross_attention_states,
   1942     attention_mask=attention_mask,
   1943     position_ids=position_ids,
   1944     cross_attention_mask=cross_attention_mask,
   1945     full_text_row_masked_out_mask=full_text_row_masked_out_mask,
   1946     past_key_values=past_key_values,
   1947     inputs_embeds=inputs_embeds,
   1948     use_cache=use_cache,
   1949     output_attentions=output_attentions,
   1950     output_hidden_states=output_hidden_states,
   1951     return_dict=return_dict,
   1952     cache_position=cache_position,
   1953 )
   1955 hidden_states = outputs[0]
   1956 logits = self.lm_head(hidden_states[:, -num_logits_to_keep:, :]).float()

File ~/mtrix2/mt-rix/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py:1553, in Module._wrapped_call_impl(self, *args, **kwargs)
   1551     return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   1552 else:
-> 1553     return self._call_impl(*args, **kwargs)

File ~/mtrix2/mt-rix/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py:1562, in Module._call_impl(self, *args, **kwargs)
   1557 # If we don't have any hooks, we want to skip the rest of the logic in
   1558 # this function, and just call forward.
   1559 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
   1560         or _global_backward_pre_hooks or _global_backward_hooks
   1561         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1562     return forward_call(*args, **kwargs)
   1564 try:
   1565     result = None

File ~/mtrix2/mt-rix/.venv/lib/python3.12/site-packages/transformers/models/mllama/modeling_mllama.py:1758, in MllamaTextModel.forward(self, input_ids, attention_mask, position_ids, cross_attention_states, cross_attention_mask, full_text_row_masked_out_mask, past_key_values, inputs_embeds, use_cache, output_attentions, output_hidden_states, return_dict, cache_position)
   1755 if position_ids is None:
   1756     position_ids = cache_position.unsqueeze(0)
-> 1758 causal_mask = self._update_causal_mask(
   1759     attention_mask, inputs_embeds, cache_position, past_key_values, output_attentions
   1760 )
   1762 # create position embeddings to be shared across the decoder layers
   1763 position_embeddings = self.rotary_emb(hidden_states, position_ids)

File ~/mtrix2/mt-rix/.venv/lib/python3.12/site-packages/transformers/models/mllama/modeling_mllama.py:1111, in MllamaPreTrainedModel._update_causal_mask(self, attention_mask, input_tensor, cache_position, past_key_values, output_attentions)
   1104     target_length = (
   1105         attention_mask.shape[-1]
   1106         if isinstance(attention_mask, torch.Tensor)
   1107         else past_seen_tokens + sequence_length + 1
   1108     )
   1110 # In case the provided `attention` mask is 2D, we generate a causal mask here (4D).
-> 1111 causal_mask = self._prepare_4d_causal_attention_mask_with_cache_position(
   1112     attention_mask,
   1113     sequence_length=sequence_length,
   1114     target_length=target_length,
   1115     dtype=dtype,
   1116     device=device,
   1117     cache_position=cache_position,
   1118     batch_size=input_tensor.shape[0],
   1119 )
   1121 if (
   1122     self.config._attn_implementation == "sdpa"
   1123     and attention_mask is not None
   (...)
   1128     # using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path.
   1129     # Details: https://github.com/pytorch/pytorch/issues/110213
   1130     min_dtype = torch.finfo(dtype).min

File ~/mtrix2/mt-rix/.venv/lib/python3.12/site-packages/transformers/models/mllama/modeling_mllama.py:1187, in MllamaPreTrainedModel._prepare_4d_causal_attention_mask_with_cache_position(attention_mask, sequence_length, target_length, dtype, device, cache_position, batch_size, **kwargs)
   1185 causal_mask = causal_mask.clone()  # copy to contiguous memory for in-place edit
   1186 mask_length = attention_mask.shape[-1]
-> 1187 padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :]
   1188 padding_mask = padding_mask == 0
   1189 causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill(
   1190     padding_mask, min_dtype
   1191 )

RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:1 and cuda:0!
```


### Expected behavior

## Expected Behaviour 
Described in  Potential mitigation BP, basically attention mask(causal_mask should handle device similar to attention mask for multiple gpu setups)

## Notes:
1. I've also tried to use: ```processor.tokenizer.padding_side = 'left' ``` throws same error 
2. Potential mitigation: transfer attention_mask within `MllamaPreTrainedModel._prepare_4d_causal_attention_mask_with_cache_position` to self.device if provided, updated code: 
```python
    def _prepare_4d_causal_attention_mask_with_cache_position(
        attention_mask: torch.Tensor,
        sequence_length: int,
        target_length: int,
        dtype: torch.dtype,
        device: torch.device,
        cache_position: torch.Tensor,
        batch_size: int,
        **kwargs,
    ):
        """
        Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape
        `(batch_size, key_value_length)`, or if the input `attention_mask` is already 4D, do nothing.

        Args:
            attention_mask (`torch.Tensor`):
                A 2D attention mask of shape `(batch_size, key_value_length)` or a 4D attention mask of shape
                `(batch_size, 1, query_length, key_value_length)`.
            sequence_length (`int`):
                The sequence length being processed.
            target_length (`int`):
                The target length: when generating with static cache, the mask should be as long as the static cache,
                to account for the 0 padding, the part of the cache that is not filled yet.
            dtype (`torch.dtype`):
                The dtype to use for the 4D attention mask.
            device (`torch.device`):
                The device to plcae the 4D attention mask on.
            cache_position (`torch.Tensor`):
                Indices depicting the position of the input sequence tokens in the sequence.
            batch_size (`torch.Tensor`):
                Batch size.
        """
        if attention_mask is not None:
            attention_mask = attention_mask.to(device)
        cache_position = cache_position.to(device)

        if attention_mask is not None and attention_mask.dim() == 4:
            # In this case we assume that the mask comes already in inverted form and requires no inversion or slicing.
            causal_mask = attention_mask
        else:
            min_dtype = torch.finfo(dtype).min
            causal_mask = torch.full(
                (sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=device
            )
            if sequence_length != 1:
                causal_mask = torch.triu(causal_mask, diagonal=1)
            causal_mask *= torch.arange(target_length, device=device) > cache_position.reshape(-1, 1)
            causal_mask = causal_mask[None, None, :, :].expand(batch_size, 1, -1, -1)
            if attention_mask is not None:
                causal_mask = causal_mask.clone()  # copy to contiguous memory for in-place edit
                mask_length = attention_mask.shape[-1]
                padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :]
                padding_mask = padding_mask == 0
                causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill(
                    padding_mask, min_dtype
                )

        return causal_mask


```