multi-gpu: fix tensor device placements for various models
Author: dvrogozh
Pull request: yes

Fixes: https://github.com/huggingface/transformers/issues/35762
CC: @SunMarc @ydshieh @faaany 

Comments:
@SunMarc, that's a big change since models reuse the code from each other and that's enforced on ci level with `make fix-copies`. To simplify the review I have structured commits as follows:
1. First and second commits are making changed to the "primary" sources (those from which other models copy code)
2. Third commit is automatically generated with `make fix-copies` from the first 2

So, I think it might be easier to review first 2 commits separately.

@SunMarc : can you, please, help to review again? The issue being fixed here is being noticed by other folks as well, see https://github.com/huggingface/transformers/issues/35918 and https://github.com/huggingface/transformers/pull/35920.

ping @ArthurZucker 

The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_35763). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

Hi, @SunMarc and @zucchini-nlp, I think the issue is still present for some models.
With Qwen-2.5-VL-7B for example on an 8-gpu node, using the code from the docs:
```
import torch
from transformers import Qwen2_5_VLForConditionalGeneration, AutoTokenizer, AutoProcessor

# Load the model in half-precision on the available device(s)
model = Qwen2_5_VLForConditionalGeneration.from_pretrained("Qwen/Qwen2.5-VL-7B-Instruct", device_map="auto")
processor = AutoProcessor.from_pretrained("Qwen/Qwen2.5-VL-7B-Instruct")
```
results in this device_map:
```
{'visual': 0,
 'model.embed_tokens': 1,
 'model.layers.0': 1,
 'model.layers.1': 1,
 'model.layers.2': 1,
 'model.layers.3': 2,
 'model.layers.4': 2,
 'model.layers.5': 2,
 'model.layers.6': 2,
 'model.layers.7': 2,
 'model.layers.8': 3,
 'model.layers.9': 3,
 'model.layers.10': 3,
 'model.layers.11': 3,
 'model.layers.12': 3,
 'model.layers.13': 4,
 'model.layers.14': 4,
 'model.layers.15': 4,
 'model.layers.16': 4,
 'model.layers.17': 4,
 'model.layers.18': 5,
 'model.layers.19': 5,
 'model.layers.20': 5,
 'model.layers.21': 5,
 'model.layers.22': 5,
 'model.layers.23': 6,
 'model.layers.24': 6,
 'model.layers.25': 6,
 'model.layers.26': 6,
 'model.layers.27': 6,
 'model.norm': 6,
 'model.rotary_emb': 6,
 'lm_head': 7}
 ```
 Which in turn gives the following when running generate:
 ```
 File /fsx/guilherme/miniconda3/envs/vllm/lib/python3.12/site-packages/transformers/models/qwen2_5_vl/modeling_qwen2_5_vl.py:1371, in Qwen2_5_VLModel._prepare_4d_causal_attention_mask_with_cache_position(attention_mask, sequence_length, target_length, dtype, device, cache_position, batch_size, config, past_key_values)
   1367 min_dtype = torch.finfo(dtype).min
   1368 causal_mask = torch.full(
   1369     (sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=device
   1370 )
-> 1371 diagonal_attend_mask = torch.arange(target_length, device=device) > cache_position.reshape(-1, 1)
   1372 if config.sliding_window is not None:
   1373     # if we have sliding window, we should not attend to tokens beyond sliding window length, so we mask them out also
   1374     # the check is needed to verify is current checkpoint was trained with sliding window or not
   1375     if not isinstance(past_key_values, SlidingWindowCache) or sequence_length > target_length:

RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:1 and cuda:0!
```

Manually changing the map to place 'visual' and 'model.embed_tokens' together fixes the issue, so I'm guessing they shouldn't be split
This in transformers==4.49.0 which in theory includes this fix

cc @gante 

@guipenedo : the fix has addressed issues in the reported scope, i.e. https://github.com/huggingface/transformers/issues/35762. Note 7d4b3ddde4da3cba17403072cfdb8b8c76ca1c7c as a commit on which issue was filed. I did verify the PR against reported scope on this commit and on a number of later commits, eventually not all. One of other commits to Transformers has introduced additional issue and this was overlooked likely due to parallel work/merge of the changes.

At the moment (https://github.com/huggingface/transformers/commit/31bb662db19355e56b3125ab5c48b780ec13111c) I can reproduce the issue on one of the tests from original scope of 35762:
```
tests/models/aria/test_modeling_aria.py::AriaForConditionalGenerationModelTest::test_model_parallel_beam_search
>           causal_mask *= torch.arange(target_length, device=device) > cache_position.reshape(-1, 1)
E           RuntimeError: Expected all tensors to be on the same device, but found at least two devices, xpu:1 and xpu:0!
```

Then, bisecting and applying this PR retrospectively I get ece8c42488a20ada0dcac87bd5e09f25f503ccf3 from @gante as the commit which introduced this additional issue. Overall, I think this needs additional fix. @guipenedo : I suggest to raise separate issue on that.

```
git bisect start
# bad: [24405127233e1d09b69bea1321c4db920e87bd7a] multi-gpu: fix tensor device placements for various models (#35763)
git bisect bad 24405127233e1d09b69bea1321c4db920e87bd7a
# good: [7d4b3ddde4da3cba17403072cfdb8b8c76ca1c7c] ci: fix xpu skip condition for test_model_parallel_beam_search (#35742)
git bisect good 7d4b3ddde4da3cba17403072cfdb8b8c76ca1c7c
# bad: [4d1d489617f195bae9e1cad89a3fa357c57e4648] Update README.md (#35958)
git bisect bad 4d1d489617f195bae9e1cad89a3fa357c57e4648
# bad: [4d1d489617f195bae9e1cad89a3fa357c57e4648] Update README.md (#35958)
git bisect bad 4d1d489617f195bae9e1cad89a3fa357c57e4648
# good: [62bd83947a95ee209600c4972e2bd11926a4f47a] [chat] docs fix (#35840)
git bisect good 62bd83947a95ee209600c4972e2bd11926a4f47a
# good: [e57b459997344c8231f9195af5143f1e737bdca9] Split and clean up GGUF quantization tests (#35502)
git bisect good e57b459997344c8231f9195af5143f1e737bdca9
# good: [f48ecd7608a2f80765208de7ce05e7984df396d4] Fix TP initialization (#35860)
git bisect good f48ecd7608a2f80765208de7ce05e7984df396d4
# bad: [c600e89f5c989959dea7701489d93a6dddeca3bb] Update `unwrap_and_save_reload_schedule` to use `weights_only=False` (#35952)
git bisect bad c600e89f5c989959dea7701489d93a6dddeca3bb
# bad: [61cbb723fc6760a82c1482e4d3bc76929ef2348d] Remove INC notebook reference in documentation (#35936)
git bisect bad 61cbb723fc6760a82c1482e4d3bc76929ef2348d
# bad: [478c4f2d0d54938af8e2642b034d0199e9803a47] fix(FA): QKV not being casted to target_dtype for FA with dpo lora (#35834)
git bisect bad 478c4f2d0d54938af8e2642b034d0199e9803a47
# bad: [ece8c42488a20ada0dcac87bd5e09f25f503ccf3] Test: generate with `torch.compile(model.forward)` as a fast test (#34544)
git bisect bad ece8c42488a20ada0dcac87bd5e09f25f503ccf3
# first bad commit: [ece8c42488a20ada0dcac87bd5e09f25f503ccf3] Test: generate with `torch.compile(model.forward)` as a fast test (#34544)
```

I used the following script but can't reproduce with transformers 4.49.0:

import torch
from transformers import Qwen2_5_VLForConditionalGeneration, AutoProcessor

```python
# Load the model in half-precision on the available device(s)
device_map = {'visual': 0,
 'model.embed_tokens': 1,
 'model.layers.0': 1,
 'model.layers.1': 1,
 'model.layers.2': 1,
 'model.layers.3': 2,
 'model.layers.4': 2,
 'model.layers.5': 2,
 'model.layers.6': 2,
 'model.layers.7': 2,
 'model.layers.8': 3,
 'model.layers.9': 3,
 'model.layers.10': 3,
 'model.layers.11': 3,
 'model.layers.12': 3,
 'model.layers.13': 3,
 'model.layers.14': 3,
 'model.layers.15': 3,
 'model.layers.16': 3,
 'model.layers.17': 3,
 'model.layers.18': 3,
 'model.layers.19': 3,
 'model.layers.20': 3,
 'model.layers.21': 3,
 'model.layers.22': 3,
 'model.layers.23': 3,
 'model.layers.24': 3,
 'model.layers.25': 3,
 'model.layers.26': 3,
 'model.layers.27': 3,
 'model.norm': 3,
 'model.rotary_emb': 3,
 'lm_head': 3}

model = Qwen2_5_VLForConditionalGeneration.from_pretrained("Qwen/Qwen2.5-VL-7B-Instruct", device_map=device_map)
processor = AutoProcessor.from_pretrained("Qwen/Qwen2.5-VL-7B-Instruct")


conversation = [
    {
        "role":"user",
        "content":[
            {
                "type":"image",
                "url": "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpeg"
            },
            {
                "type":"text",
                "text":"Describe this image."
            }
        ]
    }
]

inputs = processor.apply_chat_template(
    conversation,
    add_generation_prompt=True,
    tokenize=True,
    return_dict=True,
    return_tensors="pt"
).to(model.device)


# Inference: Generation of the output
output_ids = model.generate(**inputs, max_new_tokens=128)
generated_ids = [output_ids[len(input_ids):] for input_ids, output_ids in zip(inputs.input_ids, output_ids)]
output_text = processor.batch_decode(generated_ids, skip_special_tokens=True, clean_up_tokenization_spaces=True)
print(output_text)

```

[dvrogozh](https://github.com/dvrogozh), the commit you shared only do the following modification in the model : `_supports_static_cache = False`. Maybe the new cache used creates this specific issue ? 

> Maybe the new cache used creates this specific issue ?

@SunMarc : I double checked - reverting https://github.com/huggingface/transformers/commit/ece8c42488a20ada0dcac87bd5e09f25f503ccf3 as a whole most definitely helps to fix issue w/ aria model. But I think it's not issue with the new cache since changing back `_supports_static_cache = False` to `True` does not help. But changing back this part helps:
```
--- a/src/transformers/models/aria/modeling_aria.py
+++ b/src/transformers/models/aria/modeling_aria.py
@@ -1564,7 +1564,7 @@ class AriaForConditionalGeneration(AriaPreTrainedModel, GenerationMixin):
             output_hidden_states=output_hidden_states,
             return_dict=return_dict,
             logits_to_keep=logits_to_keep,
-            cache_position=cache_position,
+            #cache_position=cache_position,
         )
```

I.e. it seems that passing `cache_position` to `self.language_model()` modifies `cache_position` and causes the issue later on.

https://github.com/huggingface/transformers/blob/27d1707586c38895ef3b64cb0163645f5b8fd688/src/transformers/models/aria/modeling_aria.py#L1557-L1564



@dvrogozh the issue was not introduced in https://github.com/huggingface/transformers/commit/ece8c42488a20ada0dcac87bd5e09f25f503ccf3 -- if we revert to the commit before it and cherry pick this commit, we get the same failure as on `main`

```
git checkout f48ecd7608a2f80765208de7ce05e7984df396d4
git cherry-pick 24405127233e1d09b69bea1321c4db920e87bd7a
py.test tests/models/aria/test_modeling_aria.py::AriaForConditionalGenerationModelTest::test_model_parallel_beam_search
```
(note that the test will only fail with 4+ GPUs)


This PR did not fully fix the test either:
```
git checkout 24405127233e1d09b69bea1321c4db920e87bd7a
py.test tests/models/aria/test_modeling_aria.py::AriaForConditionalGenerationModelTest::test_model_parallel_beam_search
```
(fails with 4+ GPUs)

```
(...)
            pos_ids = (bucket_coords_h[:, None] * self.num_patches_per_side + bucket_coords_w).flatten()
            position_ids[batch_idx][p_attn_mask.view(-1).cpu()] = pos_ids

        position_ids = position_ids.to(self.position_embedding.weight.device)
>       embeddings = embeddings + self.position_embedding(position_ids)
E       RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cuda:1!

src/transformers/models/idefics3/modeling_idefics3.py:184: RuntimeError
```

___________________________________

The `cache_position` Aria fix in https://github.com/huggingface/transformers/commit/ece8c42488a20ada0dcac87bd5e09f25f503ccf3 is also non-negotiable: without it, the model won't have a full support of generative feature (e.g. generating from embeddings)

In a nutshell: this PR was absolutely in the right direction, but there are still a few bugs to squash 🤗 

And pointing fingers certainly won't help 😉 

> This PR did not fully fix the test either

I don't see the issue for idefics on the system with 1 to 4 Intel PVC cards. However, behavior highly depends on number of cards, and amount of free memory on each card. Depending on that you get model layers distributed differently across devices. So if you have different setup you very well might step into the issues I did not notice. Basically, it's wrong to think that this PR will fix all such issues once and for all. There might be other cases requiring fixes. These needs to be identified, root caused and fixed.

> The cache_position Aria fix in https://github.com/huggingface/transformers/commit/ece8c42488a20ada0dcac87bd5e09f25f503ccf3 is also non-negotiable

@gante No one negotiates the fix. I am just highlighting the place in ece8c42488a20ada0dcac87bd5e09f25f503ccf3 which is responsible for the appeared failure in the very particular scenario for `tests/models/aria/test_modeling_aria.py::AriaForConditionalGenerationModelTest::test_model_parallel_beam_search` test. I suggest that we need to think how we can address the issue I highlighted on top of ece8c42488a20ada0dcac87bd5e09f25f503ccf3. As an author of the change, can you please suggest what can be modified? Should `language_model` layer be added to no split modules or we just need to save aria's `cache_position` device and move it back after calling `language_model()`?


At first glance, a scalable solution to fix all related issues would start by being device agnostic: we are clearly seeing that some failures only start showing once we have a certain number of devices to distribute the model. Once we have that solution in place, finding and fixing all issues should become simple.

@SunMarc @dvrogozh do you know if it is possible to simulate devices on a machine? If it is possible, we could craft a test (or a script) with a very large number of devices to find all related issues 🤗 

Otherwise, I'll do a test run on one of our cluster machines, ideally with 16+ devices.