DeepSpeed-MII: 高吞吐低延迟的模型推理利器

DeepSpeed-MII

DeepSpeed-MII简介

DeepSpeed-MII(Model Implementations for Inference)是由微软DeepSpeed团队开发的开源Python库,旨在实现高吞吐、低延迟和高性价比的模型推理。它为用户提供了一种简单而强大的方式来部署和优化大型语言模型(LLM)的推理过程。

DeepSpeed-MII的核心目标是"民主化"强大的模型推理能力,让更多开发者和研究人员能够轻松地使用和部署大规模语言模型。通过集成多项先进技术,MII能够显著提升模型推理的性能,同时保持了使用的简便性。

DeepSpeed-MII架构图

关键技术

DeepSpeed-MII通过集成多项关键技术来实现高效的文本生成推理:

块式KV缓存(Blocked KV Caching): 这项技术优化了模型推理过程中的内存使用,提高了大规模语言模型的推理效率。
连续批处理(Continuous Batching): 允许动态添加新的请求到正在处理的批次中,提高了模型的吞吐量。
动态分割融合(Dynamic SplitFuse): 这是一种自适应的优化技术,能够根据输入动态调整计算图,进一步提升性能。
高性能CUDA核心: MII利用了经过优化的CUDA核心,充分发挥GPU的计算能力。

这些技术的结合使得DeepSpeed-MII能够支持如Llama-2-70B、Mixtral(MoE)8x7B和Phi-2等大型语言模型的快速、高吞吐量文本生成。根据官方的性能测试,MII在某些情况下可以提供比其他领先系统(如vLLM)高出2.5倍的有效吞吐量。

支持的模型

DeepSpeed-MII目前支持超过37,000个模型,涵盖了8种流行的模型架构。这些模型包括:

Falcon (7B - 180B)
Llama (7B - 65B)
Llama-2 (7B - 70B)
Llama-3 (8B - 405B)
Mistral (7B)
Mixtral (MoE) (8x7B)
OPT (0.1B - 66B)
Phi-2 (2.7B)
Qwen (7B - 72B)
Qwen2 (0.5B - 72B)

MII利用Hugging Face作为后端来提供模型权重和对应的分词器。这种广泛的模型支持使得DeepSpeed-MII成为一个非常灵活和强大的推理工具。

使用DeepSpeed-MII

DeepSpeed-MII的使用非常简单直观。它提供了两种主要的部署模式:非持久化管道和持久化部署。

非持久化管道

非持久化管道是尝试DeepSpeed-MII的理想方式。它只在Python脚本运行期间存在。以下是一个简单的示例:

import mii
pipe = mii.pipeline("mistralai/Mistral-7B-v0.1")
response = pipe(["DeepSpeed is", "Seattle is"], max_new_tokens=128)
print(response)

这个简单的4行代码就能够加载模型并生成文本。

持久化部署

对于长期运行的生产应用,持久化部署是更好的选择。它使用轻量级的GRPC服务器,可以同时被多个客户端查询:

import mii
client = mii.serve("mistralai/Mistral-7B-v0.1")
response = client.generate(["Deepspeed is", "Seattle is"], max_new_tokens=128)
print(response)

持久化部署还支持模型并行和模型复制,可以充分利用多GPU系统来提高性能:

client = mii.serve("mistralai/Mistral-7B-v0.1", tensor_parallel=2, replica_num=2)

这将在4个GPU上运行2个模型副本,每个副本跨2个GPU拆分。

REST API支持

DeepSpeed-MII还提供了REST API支持,使得通过HTTP请求进行模型推理变得非常简单:

client = mii.serve(
    "mistralai/Mistral-7B-v0.1",
    deployment_name="mistral-deployment",
    enable_restful_api=True,
    restful_api_port=28080,
)

这样设置后,就可以通过HTTP请求来进行模型推理:

curl --header "Content-Type: application/json" --request POST \
     -d '{"prompts": ["DeepSpeed is", "Seattle is"], "max_length": 128}' \
     http://localhost:28080/mii/mistral-deployment