【AI 充电】揭秘大语言模型实践：分布式推理的工程化落地才是关键！

标签：deepspeed AI 模型 rank 工程化 model 推理分布式

分布式推理成为大模型落地的首选方案

随着 3 月 15 日 OpenAI 重磅发布了 GPT4，其在司法考试、程序编程上的惊艳表现，将大家对大模型的热情推向了顶点，人们纷纷讨论是否我们已经进入到通用人工智能的时代。与此同时，基于大语言模型的应用也如雨后春笋出现在大家面前，其在协同办公、客服对话、语言翻译、内容生成等方面的使用均来带了前所未有的畅快体验。

在我们享受大语言模型带来的普惠 AI 能力时，它也给开发者带来了前所未有的挑战。GPT3 模型具有 1750 亿参数量，即使是针对学术界和初级用户的 Alpaca 也具有 70 亿的参数量，因此单机多卡的分布式推理便成为了大模型落地方案的不二选择。

本文将以 Bloom7B1 模型为样例，分享在阿里云容器服务 ACK 上，进行大语言模型分布式推理的具体实践。

工程化落地是大模型分布式推理的关键

随着越来越多的大语言模型发布，其中也有很多表现优秀的开源大语言模型能让大家体验，人们通过已有的大语言模型构建自己的应用也不再遥不可及。然而，与以往的模型不同，单张 GPU 卡的显存可能不足以支撑大语言模型。因此，需要使用模型并行技术，将大语言模型进行切分后，在多张 GPU 卡上进行推理。在本文中，我们使用 DeepSpeed Inference 来部署大语言模型分布式推理服务。

DeepSpeed Inference 是 Microsoft 提供的分布式推理解决方案，能够很好的支持 transformer 类型的大语言模型。DeepSpeed Inference 提供了模型并行能力，在多 GPU 上对大模型并行推理。通过张量并行技术同时利用多个 GPU，提高推理性能。DeepSpeed 还提供了优化过的推理定制内核来提高 GPU 资源利用率，降低推理延迟。详细信息可参考DeepSpeed Inference [3 ] 。

有了大模型分布式推理方案，然而想要在 Kubernetes 集群中高效部署大模型推理服务，还存在很多工程化挑战，比如大规模的 GPU 等异构资源如何高效地管理运维和自动调度？如何快速部署推理服务，服务上线后如何保证资源能够应对波动的访问量？以及没有适合的工具进行推理服务时延、吞吐、GPU 利用率、显存占用等关键指标监控，没有合理的模型切分方案，模型版本管理等。

本文使用阿里云容器服务 ACK 云原生 AI 套件进行 DeepSpeed 分布式推理的实践，可以轻松管理大规模异构资源，精细化的 GPU 调度策略和丰富的 GPU 监控告警能力，使用 Arena 快速提交和管理可弹性伸缩的推理服务，以及服务化运维等。

实践示例概述

本例中会使用以下组件：

Arena：Arena 是基于 Kubernetes 的机器学习轻量级解决方案，支持数据准备、模型开发，模型训练、模型预测的完整生命周期，提升数据科学家工作效率。同时和阿里云的基础云服务深度集成，支持 GPU 共享、CPFS 等服务，可以运行阿里云优化的深度学习框架，最大化使用阿里云异构设备的性能和成本的效益。更多 arena 信息，可以参考云原生 AI 套件开发者使用指南 [ 1] 。
Ingress：在 Kubernetes 集群中，Ingress 作为集群内服务对外暴露的访问接入点，其几乎承载着集群内服务访问的所有流量。Ingress 是 Kubernetes 中的一个资源对象，用来管理集群外部访问集群内部服务的方式。您可以通过 Ingress 资源来配置不同的转发规则，从而达到根据不同的规则设置访问集群内不同的 Service 所对应的后端 Pod。更多 Ingress 信息，可以参考 Ingress 概述 [ 2] 。
DeepSpeed Inference：是 Microsoft 提供的分布式推理解决方案，提供了对 GPT、BLOOM 等 LLM 模型的分布式推理优化，具体可参考 DeepSpeed Inference [ 3] 。

下列示例中，我们通过 Arena 在 Kubernetes 集群中部署了基于 Bloom 7B1 模型的单机多卡分布式推理服务，使用 DJLServing 作为模型服务框架。DJLServing 是由 Deep Java Library (DJL) 提供支持的高性能通用模型服务解决方案，能直接支持 DeepSpeed Inference，通过 HTTP 提供大模型推理服务，详细信息可参考 DJLServing [ 4] 。使用 Arena 提交推理任务，在 Kubernetes 中使用 Deployment 部署推理服务，从共享存储 OSS 中加载模型和配置文件，通过 Service 暴露服务，为推理服务提供弹性伸缩、GPU 共享调度、性能监控、成本分析与优化等功能，降低您的运维成本。

实践示例步骤

环境准备

创建包含 GPU 的 Kubernetes 集群 [ 5]
安装云原生 AI 套件 [ 6]

大模型推理实践

接下来演示如何使用 Arena 命令行工具，在 ACK 容器服务中提交一个 Bloom7B1 模型的单机多卡分布式推理任务，并配置 Ingress 来进行服务访问。

1. 模型配置编写

模型配置中包括了两个方面的内容：

配置文件，对应本例中的 serving.properties 文件，里面描述了模型配置的相关信息。这里重点关注两个参数：

- tensor_parallel_degree：用于指定 tensor parallel 的 size，本例中设置为 2，也就是使用 2 张 GPU 卡进行分布式推理；
- model_id：为模型的名称，huggingface 中 model 的名称，也可以是 download 后的模型地址；本例样例中，会将 bloom7B1 模型下载到 OSS 中，并通过 PVC 的形式挂载到容器内，因此这里会指定 OSS 的地址。
推理逻辑文件，用于完成模型的加载和 request 的处理，具体如下：

- get_model 函数：先进行 model 和分词器的加载，然后将 model 通过 deepspeed.init_inference 转换为具有分布式推理能力的 model，最后通过新生成的 model 来构建推理 pipeline；
- handle 函数：通过调用 get_model 函数中生成的 pipeline 来完成 tokenize，forward 和 detokenize 流程。

serving.properties 内容如下：

这里的 model_id 指定为 pvc 挂载后的容器内地址；如果没有提前 download 模型到本地，可以指定为 bigscience/bloom-7b1，程序会执行自动下载（模型文件一共 15G 作业）

engine=DeepSpeed
option.parallel_loading=true
option.tensor_parallel_degree=2
option.model_loading_timeout=600
option.model_id=model/LLM/bloom-7b1/deepspeed/bloom-7b1
option.data_type=fp16
option.max_new_tokens=100

model.py 内容如下：

mport os
import torch
from typing import Optional

import deepspeed
import logging
logging.basicConfig(format='[%(asctime)s] %(filename)s %(funcName)s():%(lineno)i [%(levelname)s] %(message)s', level=logging.DEBUG)
from djl_python.inputs import Input
from djl_python.outputs import Output
from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer

predictor = None


def get_model(properties: dict):
    model_dir = properties.get("model_dir")
    model_id = properties.get("model_id")
    mp_size = int(properties.get("tensor_parallel_degree", "2"))
    local_rank = int(os.getenv('OMPI_COMM_WORLD_LOCAL_RANK', '0'))
    logging.info(f"process [{os.getpid()}  rank is [{local_rank}]]")
    if not model_id:
        model_id = model_dir
    logging.info(f"rank[{local_rank}] start load model")
    model = AutoModelForCausalLM.from_pretrained(model_id)
    tokenizer = AutoTokenizer.from_pretrained(model_id)
    logging.info(f"rank[{local_rank}] success load model")

    model = deepspeed.init_inference(model,
                                     mp_size=mp_size,
                                     dtype=torch.float16,
                                     replace_method='auto',
                                     replace_with_kernel_inject=True)
    logging.info(f"rank[{local_rank}] success to convert model to deepspeed kernel")

    return pipeline(task='text-generation',
                    model=model,
                    tokenizer=tokenizer,
                    device=local_rank)


def handle(inputs: Input) -> Optional[Output]:
    global predictor
    if not predictor:
        predictor = get_model(inputs.get_properties())

    if inputs.is_empty():
        # Model server makes an empty call to warmup the model on startup
        return None

    data = inputs.get_as_string()
    output = Output()
    output.add_property("content-type", "application/json")
    result = predictor(data, do_sample=True, max_new_tokens=50)
    return output.add(result)

分别将 serving.properties、model.py 和模型文件（可选）上传到 OSS 上。具体操作，请参见控制台上传文件 [ 7] 。

上传到 OSS 之后，分别创建名称为 bloom7b1-pv 和 bloom7b1-pvc 的 PV 和 PVC，以用于推理服务的容器挂载。具体操作，请参见使用 OSS 静态存储卷 [ 8] 。

2. 启动服务

将配置文件信息放入 PVC 中，可通过下列 arena 命令启动推理服务。

--gpus：设置为 2，表示需要使用 2 张 GPU 卡进行分布式推理
--data：bloom7b1-pvc 为上一步创建的 pvc，/model 为 pvc 挂载到容器中的路径

arena serve custom \
    --name=bloom7b1-deepspeed \
    --gpus=2 \
    --version=alpha \
    --replicas=1 \
    --restful-port=8080 \
    --data=bloom7b1-pvc:/model \
    --image=ai-studio-registry-vpc.cn-beijing.cr.aliyuncs.com/kube-ai/djl-serving:2023-05-19 \
    "djl-serving -m "

查看任务运行情况。

$ kubectl get pod | grep bloom7b1-deepspeed-alpha-custom-serving
bloom7b1-deepspeed-alpha-custom-serving-766467967d-j8l2l    1/1     Running     0          8s

# 查看启动日志
kubectl logs bloom7b1-deepspeed-alpha-custom-serving-766467967d-j8l2l -f

服务启动日志如下，通过日志我们可以看到：

使用的 tensor parallel size 为 2 的分布式并行进行推理
服务中启动了 process id 为 92 和 93 的两个进程，rank id 分别为 0 和 1
rank0 和 rank0 会同时进行 kernel 的转换和模型的加载，以实现分布式推理的任务

INFO  ModelServer Starting model server ...
INFO  ModelServer Starting djl-serving: 0.23.0-SNAPSHOT ...
INFO  ModelServer
INFO  PyModel Loading model in MPI mode with TP: 2.
INFO  PyProcess [1,0]<stdout>:process [92  rank is [0]]
INFO  PyProcess [1,0]<stdout>:rank[0] start load model
INFO  PyProcess [1,1]<stdout>:process [93  rank is [1]]
INFO  PyProcess [1,1]<stdout>:rank[1] start load model
INFO  PyProcess [1,0]<stdout>:rank[0] success to convert model to deepspeed kernel
INFO  PyProcess [1,1]<stdout>:rank[1] success to convert model to deepspeed kernel
INFO  PyProcess [1,0]<stdout>:rank[0] success load model
INFO  PyProcess [1,1]<stdout>:rank[1] success load model
INFO  PyProcess Model [deepspeed] initialized.
INFO  PyProcess Model [deepspeed] initialized.
INFO  PyModel deepspeed model loaded in 297083 ms.
INFO  ModelServer Initialize BOTH server with: EpollServerSocketChannel.
INFO  ModelServer BOTH API bind to: http://0.0.0.0:8080

3. 服务验证

这里我们启动 port-forward 来进行快速验证

# 使用 kubectl 启动port-forward
kubectl  -n default-group port-forward svc/bloom7b1-deepspeed-alpha 9090:8080

在另一个终端，请求服务

# 打开新的终端，执行下列命令
$ curl -X POST http://127.0.0.1:9090/predictions/deepspeed -H "Content-type: text/plain" -d "I'm very thirsty, I need"
[
  {
    "generated_text":"I'm very thirsty, I need some water.\nWhat are you?\n- I'm a witch.\n- I thought you'd say that.\nI know a great witch.\nShe's right in here.\n- You know where we can go?\n- That's right, in one moment.\n- You want to"
  }
]

4. Ingress 配置

我们可配置 Ingress 来将模型服务对外透出，以用来对外部流量进行管理，保证模型可用性。为上面创建的服务配置 Ingress 流程如下：

登录容器服务管理控制台，在左侧导航栏选择集群。
在集群列表页面，单击目标集群名称，然后在左侧导航栏，选择网络 > 路由。
在路由页面，单击创建 Ingress，在创建 Ingress 对话框配置路由。

更详细的 Ingress 配置策略可以参考：创建 Nginx Ingress [ 9]

填写如下信息

Ingress 创建成功后，可以 Ingress 配置的域名来对 Bloom 模型进行访问。

% curl -X POST http://deepspeed-bloom7b1.c78d407e5fa034a5aa9ab10e577e75ae9.cn-beijing.alicontainer.com/predictions/deepspeed -H "Content-type: text/plain" -d "I'm very thirsty, I need"
[
  {
    "generated_text":"I'm very thirsty, I need to drink!\nI want more water.\nWhere is the water?\nLet me have the water, let me have the water...\nWait!\nYou're the father aren't you?\nDo you have water?\nAre you going to let me have some?\nGive me the"
  }
]

总结和展望

通过上面的例子，我们展示了如何使用 Arena 部署了一个 Bloom7B1 模型的单机多卡推理服务，使用 DeepSpeed-Inference 的模型并行推理技术，在多张 GPU 上进行推理。除了 DeepSpeed-Inference，当前也有一些其他的大模型分布式推理方案，比如 FastTransformer + Triton。后续我们也将不断探索，希望能够通过云原生 AI 套件，结合大模型分布式推理方案，用更低的成本支持高性能、低延迟、可弹性伸缩的大模型推理服务。

相关链接：

[1] 云原生 AI 套件开发者使用指南

https://help.aliyun.com/document_detail/336968.html?spm=a2c4g.212117.0.0.14a47822tIePxy

[2] Ingress 概述

https://help.aliyun.com/document_detail/198892.html?spm=a2c4g.181477.0.0.67d5225chicJHP

[3] DeepSpeed Inference

https://www.deepspeed.ai/tutorials/inference-tutorial/

[4] DJLServing

https://github.com/deepjavalibrary/djl-serving

[5] 创建托管 GPU 集群

https://help.aliyun.com/document_detail/171074.html?spm=a2c4g.171073.0.0.4c78f95a00Mb5P

[6] 安装云原生 AI 套件