云原生 AI 工程化实践之 FasterTransformer 加速 LLM 推理

01 背景
OpenAI 在 3 月 15 日发布了备受瞩目的 GPT4,它在司法考试和程序编程领域的惊人表现让大家对大语言模型的热情达到了顶点。人们纷纷议论我们是否已经跨入通用人工智能的时代。与此同时,基于大语言模型的应用也如雨后春笋般出现,为我们带来了协同办公、客服对话、语言翻译、内容生成等方面前所未有的畅快体验。
然而,当我们享受着大语言模型带来的普惠 AI 能力时,它也给开发者们带来了前所未有的挑战。随着模型不断增大,计算量也达到了空前的高度,直接导致推理时间变长。为了解决大语言模型推理的延迟问题,业界已经提供了一些解决方案,比如 Tensorrt、FasterTransformer 和 vllm。为了帮助用户解决云原生系统中的大语言模型推理加速问题,云原生 AI 套件引入了 FasterTransformer 推理加速方案。
本文将在 ACK 容器服务上,以 Bloom7B1 模型为例展示如何使用 FasterTransformer 进行推理加速。本例中会使用以下组件:
Arena
Arena 是基于 Kubernetes 的机器学习轻量级解决方案,支持数据准备、模型开发,模型训练、模型预测的完整生命周期,提升数据科学家工作效率。同时和阿里云的基础云服务深度集成,支持 GPU 共享、CPFS 等服务,可以运行阿里云优化的深度学习框架,最大化使用阿里云异构设备的性能和成本的效益。更多 arena 信息,可以参考云原生 AI 套件开发者使用指南。
Triton Server
Triton Server为Nvidia 提供了机器学习推理引擎,可以支持 Tensorflow、Pytorch、Tensorrt 和 Fastertransformer 多种 backend。云原生 AI 套件已经将 Triton Server 加入到 Arena 中,用户可以通过简单的命令行或 SDK 来在云原生系统中完成 Triton Server 服务的拉起、运维和监控。更多 AI 套件中使用 Triton Server 信息,可以参考部署 PyTorch 模型推理服务。
FasterTransformer
FasterTransformer 是真对于 Transofrmer 类型模型(也包括 encoder-only、decoder-only)的推理加速方案,其提供了 Kernel Fuse、Memory reuse、kv cache、量化等多种优化方案,同时也提供了 Tensor Parallel 和 Pipeline Parallel 两种分布式推理方案。本文将介绍如何在云原生 AI 套件中使用 FasterTransformer 进行模型的推理加速。
02 环境准备
环境准备分为两个部分,第一个部分是创建包含 GPU 的 Kubernetes 集群和安装云原生 AI 套件,第二个部分是从 huggingface 官网下载 bloom-7b1 模型。
模型的下载命令如下:
通过上面的命令,可以将 huggingface repo 中的文件下载到本地:

下载完成后,我们将 bloom-71 文件夹上传到 OSS 中,作为推理时的共享存储,OSS 的使用可以参考开始使用 OSS。
上传到 OSS 之后,分别创建名称为 bloom7b1-pv 和 bloom7b1-pvc 的 PV 和 PVC,以用于推理服务的容器挂载。具体操作,请参见使用 OSS 静态存储卷。
03 模型转换
FasterTransformer 本质上是对模型的重写,它通过 CUDA、cuDNN 和 cuBLAS 重写了 Transformer 模型结构,因此其具有自己的模型结构和模型参数的描述方式。而我们的模型一般是通过 Pytorch、Tesorflow、Megatron 或 huggingface 这样的训练框架产出,其往往又具有自己单独的一套模型结构和参数的表达,因此在使用FasterTransformer时,就需要将模型原有的 checkpoint 转换为 FasterTransformer 的结构。
FasterTransformer 中已经支持了多种类型的转换脚本,这里我们使用 FasterTransofrmer 提供的 examples/pytorch/gpt/utils/huggingface_bloom_convert.py。
云原生 AI 套件已经接入了上述的转换逻辑,因此,通过如下脚本即可完成一次模型的转换。
通过 arena log 来观察转换的日志:
通过 arena list 命令查看转换是否执行结束:
转换完成后,会在 OSS 上创建一个 model/arena/bloom-7b1-ft-fp16 文件夹,文件中会存储 FasterTransofrmer 所对应的 checkpoint。
04 性能对比
此时,我们的 OSS 上已经有两份 bloom-7b1 checkpoint,一份是 bloom-7b 文件夹存储了 huggingface 原生的 checkpoint,另一份是 bloom-7b-ft-fp16 文件夹存储了转换后的 FasterTransformer 的 checkpoint。我们将使用这两份 checkpoint 进行性能对比,看一下来 FasterTransformer 是否能够带来性能的提升。
性能对比使用 Fastertransformer 提供的 examples/pytorch/gpt/bloom_lambada.py,我们也已经集成到了 AI 套件中。这里我们分别提交两个性能评测命令。
对 Huggingface Bloom-7b1 评测的命令:
查看 HuggingFace 的结果:
对 Fastertransformer Blooom-7b 评测的命令:
查看 FasterTransformer 的结果,可以看见带来了 2.5 倍的性能提升。
通过结果对比可以看见,Fastertransformer 与原生的 Huggingface 相比有比较明显的性能提升。
05 模型部署
在这一小节,我们使用 Triton Server 对 FasterTransformer 进行部署,Triton Server 中原生并不支持 FasterTransformer 的 backend,需要我们配合 Nvidia 提供的 Fastertransformer backend 来使用。通过使用 FasterTransformer backend,Triton Server 不再进行 GPU 资源的分配,FasterTransformer backend 会根据 CUDA_VISIBLE_DEVICES 判断当前可用 GPU 资源,并分配给对应的 RANK 来执行分布式的推理。
FasterTransformer 对应的模型 Repo 目录如下所示:
使用功能 Arena 的如下命令来启动 FasterTransformer:
通过 kubectl logs,我们可以看到 triton server 的部署日志,通过日志可以看到,triton server 启动了两个 gpu 来进行分布式推理。
06 服务请求
启动 forward 进行验证:
这里我们使用 Triton Server 提供的 python SDK 所编写的脚本来向 Triton Server 发起请求。脚本中主要完成三件事情:
发起 client 请求命令如下:
07 总结
本文我们通过 Bloom-7b1 模型展示了如何在云原生 AI 套件中使用 FasterTransformer 对大语言模型进行加速,通过与 HuggingFace 的版本对比可以带来 2.5 倍的性能提升。后续我们会逐步推出更多大模型相关的推理加速方案,以满足不同的业务需求,大家敬请期待。
