使用 TorchServe 在 AWS Inferentia2 上部署高性能 Llama 2

作者：Mike Zhang, Li Ning, Sergey Ivanov, Naman Nandan, Hamid Shojanazeri, Geeta Chauhan, Abhi Shivaditya, Michael Nguyen, Pinak Panigrahi

最近，Llama 2 发布，并引起了机器学习社区的广泛关注。Amazon EC2 Inf2 实例，由 AWS Inferentia2 提供支持，现已支持 Llama 2 模型的训练和推理。在本文中，我们将展示如何使用最新的 AWS Neuron SDK 版本，在 Amazon EC2 Inf2 实例上实现 Llama-2 模型的低延迟和高成本效益推理。我们首先介绍如何创建、编译和部署 Llama-2 模型，并解释 AWS Neuron SDK 引入的优化技术，以实现低成本下的高性能。接下来，我们将展示我们的基准测试结果。最后，我们将展示如何通过 Amazon SageMaker 使用 TorchServe 在 Inf2 实例上部署 Llama-2 模型。

什么是 Llama 2

Llama 2 是一种自回归语言模型，使用优化的 Transformer 架构。Llama 2 旨在用于英语的商业和研究用途。它提供多种大小——70 亿、130 亿和 700 亿参数——以及预训练和微调版本。根据 Meta 的说法，微调版本使用监督微调 (SFT) 和带有人类反馈的强化学习 (RLHF) 来与人类在有用性和安全性方面的偏好对齐。Llama 2 在来自公开来源的 2 万亿个 token 数据上进行了预训练。微调模型旨在用于助手式聊天，而预训练模型可以适应各种自然语言生成任务。无论开发者使用哪个版本的模型，Meta 的负责任使用指南都可以帮助指导可能需要的额外微调，以定制和优化模型并采取适当的安全措施。

Amazon EC2 Inf2 实例概述

Amazon EC2 Inf2 实例采用 Inferentia2，与第一代 Inf1 实例相比，提供 3 倍的计算能力、4 倍的加速器内存，从而实现高达 4 倍的吞吐量和高达 10 倍的更低延迟。

大型语言模型 (LLM) 推理是内存密集型工作负载，性能随加速器内存带宽的增加而提升。Inf2 实例是 Amazon EC2 中唯一提供高速加速器互连（NeuronLink）的推理优化实例，可实现高性能大型 LLM 模型部署以及经济高效的分布式推理。现在，您可以在 Inf2 实例上的多个加速器上高效且经济地部署数十亿参数的 LLM。

Inferentia2 支持 FP32、TF32、BF16、FP16、UINT8 以及新的可配置 FP8 (cFP8) 数据类型。AWS Neuron 可以接受高精度 FP32 和 FP16 模型，并将它们自动转换为较低精度的数据类型，同时优化精度和性能。自动类型转换消除了对低精度再训练的需求，并可以使用更小的数据类型实现更高性能的推理，从而缩短上市时间。

为了灵活且可扩展地部署不断发展的深度学习模型，Inf2 实例具有硬件优化和软件支持，可通过标准的 PyTorch 自定义算子编程接口支持动态输入形状以及用 C++ 编写的自定义算子。

Transformers Neuron (transformers-neuronx)

Transformers Neuron 是一个软件包，使 PyTorch 用户能够部署性能优化的 LLM 推理。它包含使用 XLA 高级算子 (HLO) 实现的优化版 Transformer 模型，这使得在多个 NeuronCores 之间分片张量（即张量并行性）以及针对 Neuron 硬件进行性能优化（如并行上下文编码和 KV 缓存）成为可能。Llama 2 在 XLA HLOs 中的源代码可以在 此处找到。

Llama 2 在 Transformers Neuron 中通过 LlamaForSampling 类提供支持。Transformers Neuron 为 Hugging Face 模型提供了无缝的用户体验，可在 Inf2 实例上提供优化的推理。更多详情可在 Transformers Neuron 开发者指南中找到。在下一节中，我们将解释如何使用 Transformers Neuron 部署 Llama-2 13B 模型。并且，此示例也适用于其他基于 Llama 的模型。

使用 Transformers Neuron 进行 Llama 2 模型推理

创建模型、编译和部署

我们在这里有三个简单的步骤来在 Inf2 实例上创建、编译和部署模型。

1. 创建一个 CPU 模型，使用此脚本或以下代码片段将检查点序列化并保存到本地目录中。

from transformers import AutoModelForCausalLM
from transformers_neuronx.module import save_pretrained_split
model_cpu = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-13b-hf", low_cpu_mem_usage=True)
model_dir = "./llama-2-13b-split"
save_pretrained_split(model_cpu, model_dir)

2. 使用以下方法从保存序列化检查点的本地目录加载并编译模型。为了加载 Llama 2 模型，我们使用 Transformers Neuron 中的 LlamaForSampling。注意，环境变量 NEURON_RT_NUM_CORES 指定了运行时使用的 NeuronCores 数量，应与模型指定的张量并行度 (TP) 匹配。此外，NEURON_CC_FLAGS 可用于启用针对仅解码器 LLM 模型的编译器优化。

from transformers_neuronx.llama.model import LlamaForSampling
os.environ['NEURON_RT_NUM_CORES'] = '24'
os.environ['NEURON_CC_FLAGS'] = '--model-type=transformer'
model = LlamaForSampling.from_pretrained(
        model_dir,
        batch_size=1,
        tp_degree=24,
        amp='bf16',
        n_positions=16,
        context_length_estimate=[8]
    )

现在，让我们使用一行代码 API 编译模型并将模型权重加载到设备内存中。

model.to_neuron()

3. 最后，让我们在编译后的模型上运行推理。请注意，sample 函数的输入和输出都是 token 序列。

inputs = torch.tensor([[1, 16644, 31844, 312, 31876, 31836, 260, 3067, 2228, 31844]])
seq_len = 16
outputs = model.sample(inputs, seq_len, top_k=1)

Transformers Neuron 中的推理优化

张量并行性

Transformer Neuron 在多个 NeuronCores 上实现并行张量操作。我们将用于推理的核心数量称为 TP 度。较大的 TP 度提供了更高的内存带宽，从而导致更低的延迟，因为 LLM token 生成是内存 I/O 密集型工作负载。随着 TP 度的增加，推理延迟显著降低，我们的结果显示，TP 度从 2 增加到 24 时，整体速度提升了约 4 倍。对于 Llama-2 7B 模型，使用 2 个核心时的延迟从 30.1 ms/token 降低到使用 24 个核心时的 7.9 ms/token；类似地，对于 Llama-2 13B 模型，延迟从 57.3 ms/token 降低到 11.1 ms/token。

并行上下文编码

在 Transformer 架构中，token 是通过称为自回归采样（autoregressive sampling）的顺序过程生成的，而输入提示 token 则可以通过并行上下文编码并行处理。这可以显著降低自回归采样生成 token 之前的输入提示上下文编码延迟。默认情况下，参数 context_length_estimate 将设置为一系列 2 的幂次方数字，旨在覆盖各种上下文长度。根据用例，可以将其设置为自定义数字。这可以在使用 LlamaForSampling.from_pretrained 创建 Llama 2 模型时完成。我们分析了输入 token 长度对端到端 (E2E) 延迟的影响。如图所示，由于并行上下文编码，使用 Llama-2 7B 模型进行文本生成的延迟仅随输入提示变大而略有增加。

KV 缓存

自注意力模块使用 KV 向量执行自注意力操作。并且，KV 向量使用 token 嵌入和 KV 权重计算，因此与 token 相关联。在朴素实现中，对于每个生成的 token，都会重新计算整个 KV 缓存，但这会降低性能。因此，Transformers Neuron 库重用先前计算的 KV 向量以避免不必要的计算，这也被称为 KV 缓存，以降低自回归采样阶段的延迟。

基准测试结果

我们在不同条件下（即输出 token 数量、实例类型）对 Llama-2 7B 和 13B 模型的延迟和成本进行了基准测试。除非另有说明，我们使用数据类型 ‘bf16’ 和批量大小 1，因为这是聊天机器人和代码助手等实时应用程序的常见配置。

延迟

下图显示了在 TP 度为 24 的 inf2.48xlarge 实例上的每个 token 延迟。这里，每个输出 token 的延迟计算为端到端延迟除以输出 token 的数量。我们的实验表明，与 Llama-2 7B 在其他可比的推理优化 EC2 实例上生成 256 个 token 的端到端延迟相比，速度提高了 2 倍。

吞吐量

我们现在展示 inf2.48xlarge 实例每秒可以处理的 Llama-2 7B 和 13B 模型生成的 token 数量。在 TP 度为 24，充分利用所有 24 个 NeuronCores 的情况下，Llama-2 7B 和 13B 模型可以分别达到 130 tokens/秒 和 90 tokens/秒 的吞吐量。

成本

对于延迟优先的应用程序，我们展示了在 inf2.48xlarge 实例上托管 Llama-2 模型的成本，7B 和 13B 模型分别为每 1000 token $0.011 和 $0.016，与其他可比的推理优化 EC2 实例相比，成本节省了 3 倍。请注意，我们报告的成本是基于3 年预留实例价格计算的，这是客户用于大型生产部署的方式。

我们还比较了在 inf2.xlarge 和 inf2.48xlarge 实例上托管 Llama-2 7B 模型的成本。我们可以看到，inf2.xlarge 比 inf2.48xlarge 便宜 4 倍以上，但由于 TP 度较小，代价是延迟更长。例如，在 inf2.48xlarge 上，模型生成 256 个输出 token（输入 token 为 256 个）需要 7.9 ms，而在 Inf2.xlarge 上则需要 30.1 ms。

使用 TorchServe 在 EC2 Inf2 实例上服务 Llama2

现在，我们进入模型部署部分。在本节中，我们将向您展示如何通过 SageMaker 使用 TorchServe 部署 Llama-2 13B 模型，TorchServe 是 PyTorch 推荐的模型服务器，预装在 AWS PyTorch 深度学习容器 (DLC) 中。

本节描述了使用 TorchServe 所需的准备工作，特别是如何配置 model_config.yaml 和 inf2_handler.py，以及如何生成模型工件并预编译模型以供后续模型部署使用。提前准备模型工件可以避免在模型部署过程中进行模型编译，从而减少模型加载时间。

模型配置 model-config.yaml

在 handler 和 micro_batching 部分定义的参数用于自定义处理程序 inf2_handler.py。更多关于 model_config.yaml 的细节在这里。TorchServe 微批量处理是一种并行预处理和后处理一批推理请求的机制。当后端持续接收到传入数据时，通过更好地利用可用加速器，可以实现更高的吞吐量，更多细节请参阅这里。对于 Inf2 上的模型推理，micro_batch_size, amp, tp_degree 和 max_length 分别指定批量大小、数据类型、张量并行度和最大序列长度。

# TorchServe Frontend Parameters
minWorkers: 1
maxWorkers: 1
maxBatchDelay: 100
responseTimeout: 10800
batchSize: 16

# TorchServe Backend Custom Handler Parameters
handler:
    model_checkpoint_dir: "llama-2-13b-split"
    amp: "bf16"
    tp_degree: 12
    max_length: 100

micro_batching:
    # Used by batch_size in function LlamaForSampling.from_pretrained
    micro_batch_size: 1  
    parallelism:
        preprocess: 2
        inference: 1
        postprocess: 2

自定义处理程序 inf2_handler.py

TorchServe 中的自定义处理程序是一个简单的 Python 脚本，允许您将模型初始化、预处理、推理和后处理逻辑定义为函数。在这里，我们创建了我们的 Inf2 自定义处理程序。

1. The initialize 函数用于加载模型。在这里，Neuron SDK 将首次编译模型，并将预编译的模型保存在由 NEURONX_DUMP_TO 指定的目录中，该功能由 NEURONX_CACHE 启用。第一次之后，后续运行将检查是否已存在预编译的模型工件。如果存在，则会跳过模型编译。模型加载后，我们将发起预热推理请求，以便编译版本被缓存。利用 Neuron 持久缓存可以显著减少模型加载延迟，确保后续推理运行迅速。

os.environ["NEURONX_CACHE"] = "on"
os.environ["NEURONX_DUMP_TO"] = f"{model_dir}/neuron_cache"

TorchServe TextIteratorStreamerBatch 扩展了 Hugging Face transformers BaseStreamer，以支持 batchSize 大于 1 时的响应流式传输。

self.output_streamer = TextIteratorStreamerBatch(
    self.tokenizer,
    batch_size=self.handle.micro_batch_size,
    skip_special_tokens=True,
)

2. The inference 函数调用 send_intermediate_predict_response 来发送流式响应。

for new_text in self.output_streamer:
    logger.debug("send response stream")
    send_intermediate_predict_response(
        new_text[: len(micro_batch_req_id_map)],
        micro_batch_req_id_map,
        "Intermediate Prediction success",
        200,
        self.context,
    )

打包模型工件

使用 torch-model-archiver 将所有模型工件打包到一个文件夹 llama-2-13b-neuronx-b1 中。

torch-model-archiver --model-name llama-2-13b-neuronx-b1 --version 1.0 --handler inf2_handler.py -r requirements.txt --config-file model-config.yaml --archive-format no-archive

服务模型

export TS_INSTALL_PY_DEP_PER_MODEL="true"
torchserve --ncs --start --model-store model_store --models llama-2-13b-neuronx-b1

一旦日志显示“WORKER_MODEL_LOADED”，预编译的模型应该会保存在文件夹 llama-2-13b-neuronx-b1/neuron_cache 中，该文件夹与 Neuron SDK 版本紧密关联。然后，将文件夹 llama-2-13b-neuronx-b1 上传到您的 S3 存储桶，以便后续用于生产部署。本博客中的 Llama-2 13B 模型工件可以在这里找到，该工件与 Neuron SDK 2.13.2 相关联，位于 TorchServe 模型库中。

使用 TorchServe 在 SageMaker Inf2 实例上部署 Llama-2 13B 模型

在本节中，我们将使用 PyTorch Neuronx 容器在 SageMaker 端点上部署 Llama-2 13B 模型，使用的托管实例是 ml.inf2.24xlarge，它有 6 个 Inferentia2 加速器，与我们的模型配置 model_config.yaml 中 handler 的设置 - tp_degree: 12 相对应。鉴于我们已经使用 torch-model-archiver 将所有模型工件打包到一个文件夹中并上传到 S3 存储桶，我们现在将使用 SageMaker Python SDK 创建一个 SageMaker 模型，并使用部署未压缩模型的方法将其部署到 SageMaker 实时端点。使用 SageMaker 以这种方式部署的主要优势是速度，您将获得一个功能齐全的生产就绪端点，包含一个安全的 RESTful 端点，而无需在基础设施上花费任何精力。在 SageMaker 上部署模型和运行推理有 3 个步骤。笔记本示例可以在这里找到。

1. 创建一个 SageMaker 模型

from datetime import datetime

instance_type = "ml.inf2.24xlarge"
endpoint_name = sagemaker.utils.name_from_base("ts-inf2-llama2-13b-b1")

model = Model(
    name="torchserve-inf2-llama2-13b" + datetime.now().strftime("%Y-%m-%d-%H-%M-%S"),
    # Enable SageMaker uncompressed model artifacts
    model_data={
        "S3DataSource": {
                "S3Uri": s3_uri,
                "S3DataType": "S3Prefix",
                "CompressionType": "None",
        }
    },
    image_uri=container,
    role=role,
    sagemaker_session=sess,
    env={"TS_INSTALL_PY_DEP_PER_MODEL": "true"},
)

2. 部署一个 SageMaker 模型

model.deploy(
    initial_instance_count=1,
    instance_type=instance_type,
    endpoint_name=endpoint_name,
    volume_size=512, # increase the size to store large model
    model_data_download_timeout=3600, # increase the timeout to download large model
    container_startup_health_check_timeout=600, # increase the timeout to load large model
)

3. 在 SageMaker 上运行流式响应推理 当端点处于服务状态时，您可以使用 invoke_endpoint_with_response_stream API 调用来调用模型。此功能使每个生成的 token 都能返回给用户，从而增强用户体验。当生成整个序列耗时较长时，此功能特别有用。

import json

body = "Today the weather is really nice and I am planning on".encode('utf-8')
resp = smr.invoke_endpoint_with_response_stream(EndpointName=endpoint_name, Body=body, ContentType="application/json")
event_stream = resp['Body']
parser = Parser()
for event in event_stream:
    parser.write(event['PayloadPart']['Bytes'])
    for line in parser.scan_lines():
        print(line.decode("utf-8"), end=' ')

推理示例

输入

“今天天气很好，我正打算”

输出

“今天天气真的很好，我正打算去海边。我打算带上相机，拍一些海边的照片。我打算拍沙滩、海水和人的照片。我还打算拍日落的照片。我非常兴奋能去海边拍照。

海边是拍照的好地方。沙滩、海水和人都是很好的拍照对象。日落也是很好的拍照对象。”

结论

在本文中，我们展示了如何使用 Transformers Neuron 运行 Llama 2 模型推理，以及如何通过 Amazon SageMaker 在 EC2 Inf2 实例上使用 TorchServe 部署 Llama 2 模型服务。我们展示了使用 Inferentia2 的优势——低延迟和低成本——这得益于 AWS Neuron SDK 中的优化，包括张量并行性、并行上下文编码和 KV 缓存，这些优化尤其适用于 LLM 推理。要及时了解最新信息，请关注AWS Neuron 的最新发布版本以获取新功能。

立即开始在 EC2 和通过 SageMaker 使用 Llama 2 示例，并敬请关注如何在 Inf2 上优化 Llama 70B！