Multi GPU infrerence not supported with Mixtral(moe)!  
Author: DataCTE
Pull request: no

### System Info

(most recent call last):
  File "/deep-pool/inference/text-generation-webui/modules/callbacks.py", line 57, in gentask
    ret = self.mfunc(callback=_callback, args, self.kwargs)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/deep-pool/inference/text-generation-webui/modules/text_generation.py", line 352, in generate_with_callback
    shared.model.generate(kwargs)
  File "/home/alex/miniconda/envs/textgen/lib/python3.11/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/deep-pool/inference/text-generation-webui/transformers/src/transformers/generation/utils.py", line 1764, in generate
    return self.sample(
           ^^^^^^^^^^^^
  File "/deep-pool/inference/text-generation-webui/transformers/src/transformers/generation/utils.py", line 2861, in sample
    outputs = self(
              ^^^^^
  File "/home/alex/miniconda/envs/textgen/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/alex/miniconda/envs/textgen/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/alex/miniconda/envs/textgen/lib/python3.11/site-packages/accelerate/hooks.py", line 165, in new_forward
    output = module._old_forward(args, **kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/deep-pool/inference/text-generation-webui/transformers/src/transformers/models/mixtral/modeling_mixtral.py", line 1244, in forward
    aux_loss = load_balancing_loss_func(
               ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/deep-pool/inference/text-generation-webui/transformers/src/transformers/models/mixtral/modeling_mixtral.py", line 98, in load_balancing_loss_func
    gate_logits = torch.cat(gate_logits, dim=0)
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cuda:1! (when checking argument for argument tensors in method wrapper_CUDA_cat)
Output generated in 2.42 seconds (0.00 tokens/s, 0 tokens, context 65, seed 459973075)


it seems no matter what I try Mixtral models explicitly do not support multi-GPU inference. No other model on via transformers has this from what I know and this seems to be a bug of some kind. 

thank you so much for your time. 

### Who can help?

_No response_

### Information

- [X] 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

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model_id = "mistralai/Mixtral-8x7B-Instruct-v0.1"
tokenizer = AutoTokenizer.from_pretrained(model_id)

model = AutoModelForCausalLM.from_pretrained(model_id, load_in_4bit=True, device_map="auto" ,use_flash_attention_2=True)

text = "Hello my name is"
inputs = tokenizer(text, return_tensors="pt").to(0)

outputs = model.generate(**inputs, max_new_tokens=20)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))

### Expected behavior

model output (but getting multi gpu inference not supported) 

Comments:
cc @SunMarc? 

There's a related PR with the fix here: #27948 

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.