meta-llama/Llama-3.2-11B-Vision-Instruct, device_map = 'auto', padding ruins _prepare_4d_causal_attention_mask_with_cache_position
Author: AndriiZelenko
Pull request: no
Labels: name: 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


```

Comments:
Same as https://github.com/huggingface/transformers/pull/35763

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.