Add fp16 support of Qwen1.5MoE models (A2.7B) to DeepSpeed-FastGen

[ZonePG]

This PR adds support for Qwen1.5MoE-A2.7B models.

support for deepspeedai/DeepSpeed-MII#457

Test Code

for mii pipeline:

import mii

pipe = mii.pipeline("/data/zonepg/models/Qwen/Qwen1.5-MoE-A2.7B")
responses = pipe("DeepSpeed is", max_new_tokens=128, do_sample=False)
if pipe.is_rank_0:
    print(responses[0])

for huggingface:

import mii

from transformers import AutoModelForCausalLM, AutoTokenizer
from transformers.generation import GenerationConfig
import torch
tokenizer = AutoTokenizer.from_pretrained("/data/zonepg/models/Qwen/Qwen1.5-MoE-A2.7B")
model = AutoModelForCausalLM.from_pretrained("/data/zonepg/models/Qwen/Qwen1.5-MoE-A2.7B", device_map="auto", torch_dtype=torch.float16, trust_remote_code=True).eval()
print(model)
inputs = tokenizer('DeepSpeed is', return_tensors='pt')
inputs = inputs.to(model.device)
pred = model.generate(**inputs, max_new_tokens=128, do_sample=False, repetition_penalty=1.0)
test = tokenizer.decode(pred.cpu()[0], skip_special_tokens=False)
print(test)

Qwen1.5-MoE-A2.7B

Huggingface output with prompt "DeepSpeed is":

 a deep learning framework that is designed to accelerate the training of large-scale neural networks. It is built on top of PyTorch and provides a set of tools and techniques for optimizing the performance of deep learning models.

DeepSpeed supports a variety of hardware accelerators, including GPUs, TPUs, and FPGAs, and can be used to train models on distributed systems, such as clusters of GPUs or TPUs.

One of the key features of DeepSpeed is its ability to automatically parallelize the training of deep learning models across multiple GPUs or TPUs. This can significantly reduce the time required to train large models, as it allows the

DeepSpeed-FastGen output with prompt "DeepSpeed is":

 a deep learning framework that is designed to accelerate the training of large-scale neural networks. It is built on top of PyTorch and provides a set of tools and techniques for optimizing the performance of deep learning models.

DeepSpeed supports a variety of hardware accelerators, including GPUs, TPUs, and FPGAs, and can be used to train models on distributed systems, such as clusters of GPUs or TPUs.

One of the key features of DeepSpeed is its ability to automatically parallelize the training of deep learning models across multiple GPUs or TPUs. This can significantly reduce the time required to train large models, as it allows the

DeepSpeed-FastGen output with prompt "DeepSpeed is" with 8-way sharding:

 a deep learning framework that is designed to accelerate the training of large-scale neural networks. It is built on top of PyTorch and provides a set of tools and techniques for optimizing the performance of deep learning models.

DeepSpeed supports a variety of hardware accelerators, including GPUs, TPUs, and FPGAs, and can be used to train models on distributed systems, such as clusters of GPUs or TPUs.

One of the key features of DeepSpeed is its ability to automatically parallelize the training of deep learning models across multiple GPUs or TPUs. This can significantly reduce the time required to train large models, as it allows the

[ZonePG]

I am not sure if using torch.sigmoid directly will affect performance?

[heiseon]

When I use your source code to build DeepSpeed，and run “for mii pipeline” code， the process is blocked and no error。How should I identify the problem? I use 4090 GPU and transformer is 4.41.0.dev0, torch is 2.2.1, cuda version is 11.8.
BTW the transformer code runs well , just runs very slow.

[ZonePG]

Hi, @heiseon I just created a new conda environment and built it from my deepspeed code and deepspeed-mii offical source code, and it’s ok without any issues.

maybe you can delete ~/.cache/torch_extensions/pyxxx_cuxxx and try it again.

my path is /data/zonepg/.cache/torch_extensions/py311_cu121.

[heiseon]

Hi, @heiseon I just created a new conda environment and built it from my deepspeed code and deepspeed-mii offical source code, and it’s ok without any issues.

maybe you can delete ~/.cache/torch_extensions/pyxxx_cuxxx and try it again.

my path is /data/zonepg/.cache/torch_extensions/py311_cu121.

delete ~/.cache/torch_extensions/pyxxx_cuxxx is worked for me.
I have another question. When using the Qwen1.5-MoE-A2.7B-Chat-GPTQ-Int4 quantized version, an error occurs: 'Could not find a mapping for dependency "mlp.experts.18.gate_proj.bias"'. Does it mean that the GPTQ quantized version of the model is not supported?"
Would you like assistance with addressing this issue or understanding more about GPTQ quantization compatibility?

[ZonePG]

Hi, @heiseon It does not support quantized Qwen models currently. Supporting this would likely require a big effort, so it might not be considered in the short term.