摆脱 GPU 资源束缚，Amazon Inf2 让大模型推理更高效

引言

大语言模型（LLM）的兴起改变了自然语言处理（NLP），在文本分类、内容创作和问答等任务中展现出了卓越的能力。然而 GPU 资源的紧缺，为生产环境中部署这些庞大的模型带来了挑战。这些问题阻碍了 LLM 的广泛采用，特别是对于计算能力有限的企业。借助亚马逊云科技的自研芯片，可以有效的缓解资源不足的问题。此外 vLLM 作为主流的推理框架，也将 Inf2 添加到了功能清单，我们将借助 vLLM 强大的推理能力，在 Inf2 上搭建推理服务。

关于 Amazon Inf2

Amazon EC2 Inf2 实例由最多 12 个 Inferentia2 芯片提供支持，为部署大型语言模型和其他生成式 AI 工作负载提供了经济高效、高性能的解决方案。

借助高达 2.3 PetaFLOPS 的计算能力和 384GB 的总加速器内存，以及 9.8TB/秒的总内存带宽，Inf2 实例比其他可比的 Amazon EC2 实例的性价比高出 40%。为了实现加速器之间的快速通信，Inf2 实例支持 192GB/秒的 NeuronLink，可实现高性能的集群通信操作。这允许跨多个 Inferentia2 设备分片大型模型（例如通过张量并行），从而优化延迟和吞吐量。

目前，Amazon EC2 Inf2 提供四种实例类型，其中最大的实例规格（Inf2.48xlarge）能够部署高达 1750 亿参数的模型。

Instance size	# of Inferentia2 devices	vCPUs	Host Memory (GiB)	FP8/FP16/BF16/TF32 TFLOPS	FP32 TFLOPS	Device Memory (GiB)	Instance Memory Bandwidth (GiB/sec)	NeuronLink-v2 device-to-device (GiB/sec/device)
Inf2.xlarge	1	4	16	190	47.5	32	820	N/A
Inf2.8xlarge	1	32	128	190	47.5	32	820	N/A
Inf2.24xlarge	6	96	384	1140	285	192	4920	192
Inf2.48xlarge	12	192	768	2280	570	384	9840	192

（https://awsdocs-neuron.readthedocs-hosted.com/en/latest/general/arch/neuron-hardware/inf2-arch.html）

AWS Neuron

AWS Neuron 是一个用于在 AWS Inferentia 和 AWS Trainium 基础实例上运行深度学习工作负载的开发工具包（SDK）。AWS Neuron 原生集成了 TensorFlow、PyTorch 和 Apache MXNet 等主流深度学习框架。提供了一套完整的工具链，包括深度学习编译器、运行时和相关工具。并为 Hugging Face 等热门模型库中的模型实现开箱即用的性能优化。通过使用 AWS Neuron，可以在 AWS 平台上快速构建、训练和部署高性能的深度学习模型。借助 AWS Neuron SDK，可以在 AWS Inf2 实例上部署模型。

vLLM

vLLM 是一个由加州大学伯克利分校的研究人员开发的开源库，旨在加速大语言模型（LLM）的推理和服务。vLLM 引入了创新的算法和优化，显著提高了 LLM 推理和服务的速度和效率。Hugging Face Transformers（HF）和 Hugging Face Text Generation Inference（TGI）的基准测试显示，vLLM 实现了高达 24 倍于 HF 和 3.5 倍于 TGI 的吞吐量。vLLM 从版本 0.3.3 开始，vL 加了对 AWS Inferentia2 的实验性支持。

使用 vLLM 在 Inf2 进行 Mistral 模型的推理

1. 环境说明

我们将使用以下步骤在 Inf2 上安装 vLLM v0.4.8 和 Neuron Compiler [2.14.213.0] ，在此基础上运行 Mistral 7B v0.3 的大模型。

2. 配置 Inf2 实例

启动 Inf2实例，选择 DeepLearning AMI Neuron (Ubuntu 22.04) 20240708

在这里我们将使用预装了 Neuron 和 Ubuntu 22.04 的 DeepLearning AMI 作为运行环境。Neuron Deep Learning AMI 是一个预构建的 Amazon 镜像，预装了经过优化的机器学习框架和库，用于在 AWS Neuron 实例上运行深度学习工作负载。极大地简化了环境设置。启动实例后，我们将激活 Neuron 环境：

source /opt/aws_neuronx_venv_pytorch_2_1/bin/activate

您可以通过运行以下命令检查 Neuronx 的相关版本：

neuronx-cc --version

3. 安装 vllm

下载 vLLM 并安装所需的依赖项：

git clone -b v0.4.3 https://github.com/vllm-project/vllm.git
pip install -U -r requirements-neuron.txt
pip install .

在我们安装完成后，可以看到对应的 vLLM 版本和 Neuron 版本。

4. 调整 vllm 配置

打开文件 ~/vllm/vllm/engine/arg_utils.py，找到第 264 行：

parser.add_argument('--block-size',
                            type=int,
                            default=EngineArgs.block_size,
                            choices=[8, 16, 32,128],
                            help='Token block size for contiguous chunks of '
                            'tokens.')

将 block_size 设置为等于 max_model_len，这里我们设置为 128。

5. 启动模型 vLLM 通过 FastAPI 进行推理

vLLM 的模型服务可以通过 Python 脚本在本机或 Docker 环境中启动。这里我们将展示使用 Python 命令的方式加载 Mistral-7B 模型，并同步启动内置 FastAPI 的 vLLM 服务。

python3 -m vllm.entrypoints.api_server \
 --model  mistralai/Mistral-7B-v0.3 \
 --max-num-seqs 8 \ 
 --max-model-len 128 \ 
 --block-size 128 \
 --device neuron \
 --tensor-parallel-size 2

启动后会监听 8000 端口，接收相关的推理请求。

6. 调用推理接口

我们将通过 HTTP POST 请求的方式调用 vLLM 进行推理。

curl -X POST -H "Content-Type: application/json" \
  localhost:8000/generate \
  -d '{"prompt": "[INST]  Hello, my name is [/INST]"}'

{"text":["[INST] Hello my name is[/INST] Dillon and I'm here to share with you what I've learned"]}

至此，我们已将 Mistral 7B v0.3 模型通过 vLLM 框架部署在 Inf2 实例上，并提供 API 服务。

总结

AWS Inferentia2 加速芯片为大型语言模型的高效部署带来了新的可能性。本文以开源的 Mistral 7B 大型语言模型为例,详细介绍了如何借助 vLLM 在 Inferentia2 上进行大模型的推理。使用 Inferentia2 将极大地解决 GPU 资源紧缺问题，为大模型推理提供更加多样化的选择。