Repository: huggingface/transformers
Pull request: fix(models): Fix ViVit when dealing with multiple GPUs

# What does this PR do?

`test_modeling_vivit.py::VivitModelTest::test_model_parallelism` fails on multiple gpus with:


```
  /usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py:1775: in _wrapped_call_impl
  return self._call_impl(*args, **kwargs)
 [...]
  src/transformers/models/vivit/modeling_vivit.py:222: in forward
  context_layer, attention_probs = attention_interface(
  _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
  _ _ _

  module = VivitSelfAttention(
  (query): Linear(in_features=32, out_features=32, bias=True)
  (key): Linear(in_features=32, out_features=32, bias=True)
  (value): Linear(in_features=32, out_features=32, bias=True)
  )
  query = tensor([[[[-0.1032, -0.0349, -0.0562,  ...,  0.1166,  0.1275,  0.0801],
  [ 0.2075, -0.1922,  0.0339,  ..., -0... [-0.0146, -0.0876, -0.0722,  ...,  0.0379,  0.2044, -0.1474]]]],
  device='cuda:0', grad_fn=<TransposeBackward0>)
  key = tensor([[[[-1.1990e-01, -1.7907e-01,  1.0957e-01,  ...,  3.8359e-02,
  -8.9726e-02, -1.6479e-01],
  [...e-01,  ..., -1.0423e-01,
  -8.2719e-02,  1.5556e-01]]]], device='cuda:1',
  grad_fn=<TransposeBackward0>)
  value = tensor([[[[-5.6872e-02,  4.6242e-02,  3.0870e-02,  ..., -9.6589e-02,
  2.0439e-02,  3.4408e-02],
  [...e-02,  ...,  1.2462e-01,
  -7.5010e-02,  6.4968e-02]]]], device='cuda:1',
  grad_fn=<TransposeBackward0>)
  attention_mask = None, dropout = 0.0, scaling = 0.3535533905932738, is_causal = False, kwargs = {}, sdpa_kwargs = {}

[...]

  >       attn_output = torch.nn.functional.scaled_dot_product_attention(
  query,
  key,
  value,
  attn_mask=attention_mask,
  dropout_p=dropout,
  scale=scaling,
  is_causal=is_causal,
  **sdpa_kwargs,
  )
  E       RuntimeError: Expected query, key, and value to have the same device type, but got query.device: cuda:0 key.device: cuda:1 and value.device: cuda:1 instead.

  src/transformers/integrations/sdpa_attention.py:92: RuntimeError
```

This is caused by tensors living in different GPUs when doing operations.

I've added `.to()` calls to make sure they are on the same one.

Fixed verified via a GH worker running multiple GPUs

I don't know if this is the best approach though, or if there are better ways to address this.



## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [x ] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#create-a-pull-request),
      Pull Request section?
- [ ] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [ ] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [x ] Did you write any new necessary tests?


## Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

<!-- Your PR will be replied to more quickly if you can figure out the right person to tag with @

 If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.
 Please tag fewer than 3 people.

Models:

- text models: @ArthurZucker @Cyrilvallez
- vision models: @yonigozlan @molbap
- audio models: @eustlb @ebezzam @vasqu
- multimodal models: @zucchini-nlp
- graph models: @clefourrier

Library:

- generate: @zucchini-nlp (visual-language models) or @gante (all others)
- continuous batching: @remi-or @ArthurZucker @McPatate
- pipelines: @Rocketknight1
- tokenizers: @ArthurZucker and @itazap
- trainer: @SunMarc
- attention: @vasqu @ArthurZucker @CyrilVallez
- model loading (from pretrained, etc): @CyrilVallez
- distributed: @3outeille @ArthurZucker
- CIs: @ydshieh

Integrations:

- ray/raytune: @richardliaw, @amogkam
- Big Model Inference: @SunMarc
- quantization: @SunMarc @MekkCyber
- kernels: @MekkCyber @drbh
- peft: @BenjaminBossan @githubnemo

Devices/Backends:

- AMD ROCm: @ivarflakstad
- Intel XPU: @IlyasMoutawwakil
- Ascend NPU: @ivarflakstad 

Documentation: @stevhliu

Research projects are not maintained and should be taken as is.

 -->


Changed files:
filename: src/transformers/models/vivit/modeling_vivit.py
patch: @@ -364,7 +364,7 @@ class VivitPreTrainedModel(PreTrainedModel):
     main_input_name = "pixel_values"
     input_modalities = "video"
     supports_gradient_checkpointing = True
-    _no_split_modules = []
+    _no_split_modules = ["VivitLayer"]
     _supports_sdpa = True
     _supports_flash_attn = True
     _supports_flex_attn = True