[Bug] Fix device mismatch in Ministral3 RoPE forward pass
Author: dysby
Pull request: no
Labels: name: bug

### System Info

I encountered a `RuntimeError` when running `Ministral3` where the tensors in the Rotary Positional Embedding (RoPE) calculation were on different devices (CPU vs CUDA).

While `inv_freq_expanded` is correctly moved to `x.device`, the `position_ids_expanded` tensor is created without an explicit device cast. If the input `position_ids` are not manually moved to the GPU before the forward pass, the matrix multiplication fails.

**Error Traceback:**

```text
File ".../transformers/models/ministral3/modeling_ministral3.py", line 336, in forward
    freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
RuntimeError: Expected all tensors to be on the same device, but got mat2 is on cpu, different from other tensors on cuda:0

```

### **Location:**

* **File:** `src/transformers/models/ministral3/modular_ministral3.py`
* **Method:** `forward` (inside the RoPE embedding class)

### **The Fix:**

I resolved this by ensuring `position_ids_expanded` is explicitly moved to the same device as the input `x`, mirroring the behavior of `inv_freq_expanded`.

**Before:**

```python
inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device)
position_ids_expanded = position_ids[:, None, :].float() 

```

**After:**

```python
inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device)
position_ids_expanded = position_ids[:, None, :].float().to(x.device)  # <-- Added .to(x.device)

```


@zucchini-nlp 

Name: transformers Version: 5.0.0rc3


### Who can help?

_No response_

### Information

- [ ] The official example scripts
- [ ] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

Demo from the ministral3 model docs

### Expected behavior

Fix

Comments:
Spent some time on this.

In `Ministral3RotaryEmbedding.forward()` the `inv_freq_expanded` gets `.to(x.device)` but `position_ids_expanded` doesn't — it stays on whatever device `position_ids` came from. Then the matmul blows up if they're on different devices.

This happens when you pass `position_ids` on CPU while the model is on CUDA, or in multi-GPU setups with custom `device_map`. Works fine normally because when `position_ids=None` it's created from `cache_position` which inherits `inputs_embeds.device`.

btw this isn't Ministral3-specific — same pattern in ~107 models (Llama, Mistral, Qwen2, Gemma2, DeepSeek-V2, etc). They all do `position_ids_expanded = position_ids[:, None, :].float()` without `.to(x.device)`.

The fix is just adding `.to(x.device)` to `position_ids_expanded`. Safe because `tensor.to(same_device)` is a no-op in PyTorch.

One thing — `modeling_ministral3.py` is auto-generated from modular which inherits from Mistral. So probably makes more sense to fix it in `MistralRotaryEmbedding` or `LlamaRotaryEmbedding` to cascade down.

Hi @raimbekovm, this makes sense but the conditions to trigger it seem a bit weird. In particular, all models will fail if you pass in a CPU tensor input while they're on CUDA, so that's not really a 'bug'. So I'd guess this only really applies with very specific custom device maps, right?

Fair point. The inv_freq needs .to(x.device) because it's an internal buffer that can end up on different device with device_map. But position_ids is user input like input_ids/attention_mask - those don't get auto-moved either. So yeah, this is user error, not a bug.

Hey, could you clarify if you are trying to pass `position_ids` created manually to model's forward but those are kept on CPU? Like this:

```python
model  = model.to("cuda")
input_ids = tokenizer('hellow world', return_tensors="pt").to(model.device)
position_ids = torch.arange(input_ids.shape[-1])
logits = model(input_ids, position_ids=position_ids)
```


If I got it right, I think you just have to move positions to model's device manually, doesn't look like a bug. In multigpu it is usually handled by `accelerate` which moves all inputs to the same device as the layer

Hi, 
i just follow Ministral3 transformers code demo in the model hf page.

```python
tokenized = tokenizer.apply_chat_template(messages, return_tensors="pt", return_dict=True)

tokenized["input_ids"] = tokenized["input_ids"].to(device="cuda")
tokenized["pixel_values"] = tokenized["pixel_values"].to(dtype=torch.bfloat16, device="cuda")
image_sizes = [tokenized["pixel_values"].shape[-2:]]
```

Is it that apply_chat_template creates position_ids and they ignore?

Nit - ideally `tokenized.to(device="cuda")` will move all fields in `tokenized` to CUDA, so we don't move each one separately.

> Is it that apply_chat_template creates position_ids and they ignore?

AFAIK it doesn't create any position ids. Can you link to the demo snippet so I could reproduce it?

Here, 
from https://huggingface.co/mistralai/Ministral-3-8B-Instruct-2512

```python
import torch
from transformers import Mistral3ForConditionalGeneration, MistralCommonBackend

model_id = "mistralai/Ministral-3-8B-Instruct-2512"

tokenizer = MistralCommonBackend.from_pretrained(model_id)
model = Mistral3ForConditionalGeneration.from_pretrained(model_id, device_map="auto")

image_url = "https://static.wikia.nocookie.net/essentialsdocs/images/7/70/Battle.png/revision/latest?cb=20220523172438"

messages = [
    {
        "role": "user",
        "content": [
            {
                "type": "text",
                "text": "What action do you think I should take in this situation? List all the possible actions and explain why you think they are good or bad.",
            },
            {"type": "image_url", "image_url": {"url": image_url}},
        ],
    },
]

tokenized = tokenizer.apply_chat_template(messages, return_tensors="pt", return_dict=True)

tokenized["input_ids"] = tokenized["input_ids"].to(device="cuda")
tokenized["pixel_values"] = tokenized["pixel_values"].to(dtype=torch.bfloat16, device="cuda")
image_sizes = [tokenized["pixel_values"].shape[-2:]]

output = model.generate(
    **tokenized,
    image_sizes=image_sizes,
    max_new_tokens=512,
)[0]

decoded_output = tokenizer.decode(output[len(tokenized["input_ids"][0]):])
print(decoded_output)
```

Thanks for the pointer! I see that there is also an attn mask returned which is on CPU, therefore the position ids are internally created on CPU.

Adding `tokenized["attention_mask"] = tokenized["attention_mask"].to(device="cuda")` will solve it, but as mentioned above the recommended way is to move everything like `tokenized = tokenized.to(device="cuda")`. 

I will close the issue as resolved then. Feel free to open a PR on the ministral model page fixing the demo snippet :)

@zucchini-nlp ,  just a quick additional note 

```
tokenized = tokenized.to(device="cuda")
```

will not work directly, we get an exception

```
lib/python3.12/site-packages/torch/nn/modules/conv.py", line 543, in _conv_forward
    return F.conv2d(
           ^^^^^^^^^
RuntimeError: Input type (torch.cuda.DoubleTensor) and weight type (CUDABFloat16Type) should be the same
```

unless we specifically convert pixel_values to the model dtype

```
tokenized = tokenized.to(device="cuda")
tokenized["pixel_values"] = tokenized["pixel_values"].to(dtype=torch.bfloat16, device="cuda")
```

What is `transformers` recommended way to deal with VLMs inputs?

Oh it has pixels, so it's a BatchFeature I suppose. We should be able to do `tokenized = tokenized.to(device="cuda", dtype=torch.bfloat16)` as well on multimodal input dict