multi-gpu: test_model_parallel_beam_search tests fail with "RuntimeError: Expected all tensors to be on the same device"
Author: dvrogozh
Pull request: no

With:
* Transformers: 7d4b3ddde
* https://github.com/pytorch/pytorch/commit/4e4b8592a32f701b4974679ab1381ba7cccd4844 (torch 2.7 candidate)

On:
* 2 card Intel(R) Data Center GPU Max 1550 (aka PVC), note: each card has 2 tiles, in total there are 4 torch devices available

`test_model_parallel_beam_search` tests for a number of models fail with "RuntimeError: Expected all tensors to be on the same device":
```
# TRANSFORMERS_TEST_DEVICE_SPEC=spec.py python3 -m pytest -k test_model_parallel_beam_search \
  tests/models/aria \
  tests/models/falcon_mamba \
  tests/models/gpt2 \
  tests/models/gpt_bigcode \
  tests/models/idefics2 \
  tests/models/imagegpt \
  tests/models/instructblip \
  tests/models/instructblipvideo \
  tests/models/mamba \
  tests/models/mbart \
  tests/models/opt \
  tests/models/qwen2_vl \
  tests/models/xglm
...
FAILED tests/models/aria/test_modeling_aria.py::AriaForConditionalGenerationModelTest::test_model_parallel_beam_search - RuntimeError: Expected all tensors to be on the same device, but found at least two devices, xpu:1 and xpu:0!
FAILED tests/models/falcon_mamba/test_modeling_falcon_mamba.py::FalconMambaModelTest::test_model_parallel_beam_search - RuntimeError: Expected all tensors to be on the same device, but found at least two devices, xpu:0 and xpu:1!
FAILED tests/models/gpt2/test_modeling_gpt2.py::GPT2ModelTest::test_model_parallel_beam_search - RuntimeError: Expected all tensors to be on the same device, but found at least two devices, xpu:0 and xpu:1!
FAILED tests/models/gpt_bigcode/test_modeling_gpt_bigcode.py::GPTBigCodeModelTest::test_model_parallel_beam_search - RuntimeError: Expected all tensors to be on the same device, but found at least two devices, xpu:1 and xpu:2!
FAILED tests/models/gpt_bigcode/test_modeling_gpt_bigcode.py::GPTBigCodeMHAModelTest::test_model_parallel_beam_search - RuntimeError: Expected all tensors to be on the same device, but found at least two devices, xpu:0 and xpu:1!
FAILED tests/models/idefics2/test_modeling_idefics2.py::Idefics2ForConditionalGenerationModelTest::test_model_parallel_beam_search - RuntimeError: Expected all tensors to be on the same device, but found at least two devices, xpu:1 and xpu:0!
FAILED tests/models/imagegpt/test_modeling_imagegpt.py::ImageGPTModelTest::test_model_parallel_beam_search - RuntimeError: Expected all tensors to be on the same device, but found at least two devices, xpu:0 and xpu:1!
FAILED tests/models/instructblip/test_modeling_instructblip.py::InstructBlipForConditionalGenerationDecoderOnlyTest::test_model_parallel_beam_search - RuntimeError: Expected all tensors to be on the same device, but found at least two devices, xpu:2 and xpu:1!
FAILED tests/models/instructblipvideo/test_modeling_instructblipvideo.py::InstructBlipVideoForConditionalGenerationDecoderOnlyTest::test_model_parallel_beam_search - RuntimeError: Expected all tensors to be on the same device, but found at least two devices, xpu:2 and xpu:1!
FAILED tests/models/mamba/test_modeling_mamba.py::MambaModelTest::test_model_parallel_beam_search - RuntimeError: Expected all tensors to be on the same device, but found at least two devices, xpu:0 and xpu:1!
FAILED tests/models/mbart/test_modeling_mbart.py::MBartModelTest::test_model_parallel_beam_search - RuntimeError: Expected query, key, and value to have the same device type, but got query.device: xpu:1 key.device: xp...
FAILED tests/models/mllama/test_modeling_mllama.py::MllamaForConditionalGenerationModelTest::test_model_parallel_beam_search - RuntimeError: Expected all tensors to be on the same device, but found at least two devices, xpu:1 and xpu:3! (when c...
FAILED tests/models/opt/test_modeling_opt.py::OPTModelTest::test_model_parallel_beam_search - RuntimeError: Expected all tensors to be on the same device, but found at least two devices, xpu:0 and xpu:1!
FAILED tests/models/qwen2_vl/test_modeling_qwen2_vl.py::Qwen2VLModelTest::test_model_parallel_beam_search - RuntimeError: indices should be either on cpu or on the same device as the indexed tensor (xpu:1)
FAILED tests/models/xglm/test_modeling_xglm.py::XGLMModelTest::test_model_parallel_beam_search - RuntimeError: Expected all tensors to be on the same device, but found at least two devices, xpu:1 and xpu:2!
=============================== 15 failed, 2 passed, 4399 deselected, 2 warnings in 13.95s ===============================
```

From the log, errors occur on the following lines:

https://github.com/huggingface/transformers/blob/7d4b3ddde4da3cba17403072cfdb8b8c76ca1c7c/src/transformers/models/aria/modeling_aria.py#L1116
https://github.com/huggingface/transformers/blob/7d4b3ddde4da3cba17403072cfdb8b8c76ca1c7c/src/transformers/models/falcon_mamba/modeling_falcon_mamba.py#L312
https://github.com/huggingface/transformers/blob/7d4b3ddde4da3cba17403072cfdb8b8c76ca1c7c/src/transformers/models/gpt2/modeling_gpt2.py#L821
https://github.com/huggingface/transformers/blob/7d4b3ddde4da3cba17403072cfdb8b8c76ca1c7c/src/transformers/models/gpt_bigcode/modeling_gpt_bigcode.py#L962
https://github.com/huggingface/transformers/blob/7d4b3ddde4da3cba17403072cfdb8b8c76ca1c7c/src/transformers/models/gpt_bigcode/modeling_gpt_bigcode.py#L962
https://github.com/huggingface/transformers/blob/7d4b3ddde4da3cba17403072cfdb8b8c76ca1c7c/src/transformers/models/idefics2/modeling_idefics2.py#L1307
https://github.com/huggingface/transformers/blob/7d4b3ddde4da3cba17403072cfdb8b8c76ca1c7c/src/transformers/models/imagegpt/modeling_imagegpt.py#L778
https://github.com/huggingface/transformers/blob/7d4b3ddde4da3cba17403072cfdb8b8c76ca1c7c/src/transformers/models/instructblip/modeling_instructblip.py#L1609
https://github.com/huggingface/transformers/blob/7d4b3ddde4da3cba17403072cfdb8b8c76ca1c7c/src/transformers/models/instructblipvideo/modeling_instructblipvideo.py#L1644
https://github.com/huggingface/transformers/blob/7d4b3ddde4da3cba17403072cfdb8b8c76ca1c7c/src/transformers/models/mamba/modeling_mamba.py#L264
https://github.com/huggingface/transformers/blob/7d4b3ddde4da3cba17403072cfdb8b8c76ca1c7c/src/transformers/models/mbart/modeling_mbart.py#L494
https://github.com/huggingface/transformers/blob/7d4b3ddde4da3cba17403072cfdb8b8c76ca1c7c/src/transformers/models/mllama/modeling_mllama.py#L1489
https://github.com/huggingface/transformers/blob/7d4b3ddde4da3cba17403072cfdb8b8c76ca1c7c/src/transformers/models/opt/modeling_opt.py#L885
https://github.com/huggingface/transformers/blob/7d4b3ddde4da3cba17403072cfdb8b8c76ca1c7c/src/transformers/models/qwen2_vl/modeling_qwen2_vl.py#L1489
https://github.com/huggingface/transformers/blob/7d4b3ddde4da3cba17403072cfdb8b8c76ca1c7c/src/transformers/models/xglm/modeling_xglm.py#L597

CC: @SunMarc @ydshieh @faaany 

Comments:
See proposed fix in:
* #35763

I ran the tests with 2 GPUs (CUDA) and only one is failing. So, I don't think this is an issue with the modeling code. Could you explore a bit more why this happens only on xpu ? 

```
FAILED tests/models/qwen2_vl/test_modeling_qwen2_vl.py::Qwen2VLModelTest::test_model_parallel_beam_search - RuntimeError: indices should be either on cpu or on the same device as the indexed tensor (cuda:0)
=========================== 1 failed, 14 passed, 4135 deselected, 1 warning in 64.23s (0:01:04) ============================
```

@SunMarc : unfortunately I don't have 2+ CUDA system handy and can't compare execution side by side. What is happening with execution on XPU is clear regarding "how" - tensors involved in the same operation (like sum or whatever) are placed on the different devices. Next question is "why" we have this for XPU and you suggest that's not seen on CUDA. With the later I disagree - I found at least 2 changes in the existing Transformers sources which were fixing similar issues for CUDA. I left respective comments in PR-35763, reposting them here:
https://github.com/huggingface/transformers/blob/5fa35344755d8d9c29610b57d175efd03776ae9e/src/transformers/models/mamba2/modeling_mamba2.py#L482
https://github.com/huggingface/transformers/blob/5fa35344755d8d9c29610b57d175efd03776ae9e/src/transformers/models/whisper/modeling_whisper.py#L1270

I also see this discussion https://github.com/huggingface/transformers/pull/30836#discussion_r1604774353 between you and @ArthurZucker on why this change is being made. It seems that issues similar to the one I filed here were seen for CUDA as well.

With the answering on "why", can you, please, give some guidance to me:
1. I spotted issue for XPU in many models. Do you have concerns with all these changes or a subset of them? For example, there are fixes in my PR similar to the changes done previously for CUDA which I noted above. Should I separate such fixes and stand-alone PR and we will continue discussing the remainder?
2. I think that the reason for "why" is how tensors were placed across multiple devices. And here I seek help and guidance since I am not quite familiar with this logic in Transformers. How Transformers manage tensors placement for `device_map="auto"`? Is there any guarantee that placement will be the same for different kind of devices based on some common characteristic, like for devices with the same amount of memory, etc.? Since if not, then why that's seen for XPU and not for CUDA is understandable - devices are different, so tensors placement is different, hence issues not seen before.



> I ran the tests with 2 GPUs (CUDA) and only one is failing. So, I don't think this is an issue with the modeling code. Could you explore a bit more why this happens only on xpu ?

@SunMarc, as I noted in the previous comment, I think that issue is related with how model was distributed among multiple GPU devices. I had 4 XPU devices in initial filing of the issue with multiple tests failing. But if I will force only 2 GPU devices, then I see the same result as you - only one case fails for qwen. Same as for CUDA:
```
# ZE_AFFINITY_MASK=0,1 <same pytest command as in description>
...
FAILED tests/models/qwen2_vl/test_modeling_qwen2_vl.py::Qwen2VLModelTest::test_model_parallel_beam_search - RuntimeError: indices should be either on cpu or on the same device as the indexed tensor (xpu:0)
===================== 1 failed, 14 passed, 4090 deselected, 1 warning in 9.78s =====================
```

Can someone try to run on the system with 4 CUDA devices? But overall, if I understand correctly, `device_map` distributes tensors declared in the model across multiple devices based on available device memory, distribution algorithm and some hints from model implementation like `no_split_modules`. Further it's needed to make sure that operation's tensor operands will be on the same device, but this highly depends on how tensors were initially distributed and it seems there is no automated logic to guarantee that those operands will be on the same device. If so, this falls to us to run model for multiple distribution cases and identify issues around distribution. Or implement some automation here which is likely challanging.


Thanks for giving more context and reproducing the results for 2 devices. I'll have a look at the PR then.

> Can someone try to run on the system with 4 CUDA devices? But overall, if I understand correctly, device_map distributes tensors declared in the model across multiple devices based on available device memory, distribution algorithm and some hints from model implementation like no_split_modules. Further it's needed to make sure that operation's tensor operands will be on the same device, but this highly depends on how tensors were initially distributed and it seems there is no automated logic to guarantee that those operands will be on the same device. If so, this falls to us to run model for multiple distribution cases and identify issues around distribution. Or implement some automation here which is likely challanging.

That's pretty much it. We put modules that contain residual connection inside `no_split_modules`. However, we can't do that for all modules as they might represent more than 50% of the size of the model. Usually, we put the embedding layer or the decoder layer. If there are other residual connections or big modules that we can't really put inside `no_split_modules`, we need to do a manual dispatch as you saw in the other PRs. For example, we can't really put `WhisperDecoder` inside  `no_split_modules`, so in the forward,  we don't have a choice but to the `positions` to the same device as `inputs_embeds`: 
`hidden_states = inputs_embeds + positions.to(inputs_embeds.device)`. 
Can you check that the modification you did is due to that ? 



> Can you check that the modification you did is due to that ?

@SunMarc : most of the changes in the PR are around `positions` and `inputs_embeds`. Few other changes I've made:
1. Placement of `attention_mask`, `causal_mask`, `padding_mask`
2. Placement of `hidden_states` and `residual`
3. Placement of query, keys and values in `scaled_dot_product_attention` calls