[BUG] Mixtral inference OOM

Describe the bug
I'm not sure if DeepSpeed needs to be adapted for Mixtral. When I tried using DeepSpeed inference for model inference, it didn't properly implement model parallelism. Instead, it attempted to load the complete model parameters on each GPU, which ultimately led to Out Of Memory (OOM) errors.
However, when I use llama2 for inference, it does indeed implement model parallelism. So I'm wondering if the Mixtral model requires official adaptation? Additionally, when I deploy Mixtral using the MII library, model parallelism is also successfully implemented, with the model parameters correctly split across different GPUs. But when I directly use DeepSpeed inference, it fails. I'm not sure if it's because it requires official adaptation or if there's a problem with how I'm using it. Is there anyone who can provide some guidance?

To Reproduce

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
import deepspeed

tokenizer = AutoTokenizer.from_pretrained("/workspace/shuaiqi/Model/Mixtral",use_fast=True, add_prefix_space=True)


model = AutoModelForCausalLM.from_pretrained("/workspace/shuaiqi/Model/Mixtral",torch_dtype=torch.float16)
model.eval()

model_engine = deepspeed.init_inference(model,
                                        mp_size=4, # Number of GPU
                                        dtype=torch.float16, # dtype of the weights (fp16)
                                        replace_method="auto", # Lets DS autmatically identify the layer to replace
                                        replace_with_kernel_inject=True, # replace the model with the kernel injector
)
print(model_engine)


inputs = tokenizer("Hello, my dog is cute", return_tensors="pt").to('cuda')
outputs=model_engine.generate(**inputs,max_length=128,do_sample=True,top_k=50,top_p=0.95,temperature=0.9,use_cache=True)
output = tokenizer.decode(outputs[0], skip_special_tokens=True)
print(output)

The launch command is

CUDA_VISIBLE_DEVICES=0,1,3,4 deepspeed --num_nodes=1 --num_gpus=4 example.py

Expected behavior
A clear and concise description of what you expected to happen.

ds_report output

DeepSpeed C++/CUDA extension op report
--------------------------------------------------
NOTE: Ops not installed will be just-in-time (JIT) compiled at
      runtime if needed. Op compatibility means that your system
      meet the required dependencies to JIT install the op.
--------------------------------------------------
JIT compiled ops requires ninja
ninja .................. [OKAY]
--------------------------------------------------
op name ................ installed .. compatible
--------------------------------------------------
 [WARNING]  async_io requires the dev libaio .so object and headers but these were not found.
 [WARNING]  async_io: please install the libaio-dev package with apt
 [WARNING]  If libaio is already installed (perhaps from source), try setting the CFLAGS and LDFLAGS environment variables to where it can be found.
async_io ............... [NO] ....... [NO]
fused_adam ............. [NO] ....... [OKAY]
cpu_adam ............... [NO] ....... [OKAY]
cpu_adagrad ............ [NO] ....... [OKAY]
cpu_lion ............... [NO] ....... [OKAY]
 [WARNING]  Please specify the CUTLASS repo directory as environment variable $CUTLASS_PATH
evoformer_attn ......... [NO] ....... [NO]
fused_lamb ............. [NO] ....... [OKAY]
fused_lion ............. [NO] ....... [OKAY]
inference_core_ops ..... [NO] ....... [OKAY]
cutlass_ops ............ [NO] ....... [OKAY]
quantizer .............. [NO] ....... [OKAY]
ragged_device_ops ...... [NO] ....... [OKAY]
ragged_ops ............. [NO] ....... [OKAY]
random_ltd ............. [NO] ....... [OKAY]
 [WARNING]  sparse_attn requires a torch version >= 1.5 and < 2.0 but detected 2.1
 [WARNING]  using untested triton version (2.1.0), only 1.0.0 is known to be compatible
sparse_attn ............ [NO] ....... [NO]
spatial_inference ...... [NO] ....... [OKAY]
transformer ............ [NO] ....... [OKAY]
stochastic_transformer . [NO] ....... [OKAY]
transformer_inference .. [NO] ....... [OKAY]
--------------------------------------------------
DeepSpeed general environment info:
torch install path ............... ['/workspace/shuaiqi/miniconda3/envs/Shay/lib/python3.11/site-packages/torch']
torch version .................... 2.1.2+cu121
deepspeed install path ........... ['/workspace/shuaiqi/miniconda3/envs/Shay/lib/python3.11/site-packages/deepspeed']
deepspeed info ................... 0.12.6, unknown, unknown
torch cuda version ............... 12.1
torch hip version ................ None
nvcc version ..................... 12.3
deepspeed wheel compiled w. ...... torch 2.1, cuda 12.1
shared memory (/dev/shm) size .... 1.97 TB

Screenshots
If applicable, add screenshots to help explain your problem.

System info (please complete the following information):
transformers 4.36.2
cuda 12.3
pytorch 2.1.2
deepspeed 0.12.6

Docker context
Are you using a specific docker image that you can share?

Additional context
Add any other context about the problem here.

Start with the provided reproduction script and the deepspeed.init_inference entry point, then compare its Mixtral behavior with the reported llama2 and MII behavior. Use the supplied ds_report and launch command while tracing whether model parameters are split across the four GPUs; done means Mixtral inference uses model parallelism without loading the complete model on every GPU or causing OOM.