跳转到主要内容
博客

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

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

Llama 2 is an auto-regressive language model that uses an optimized transformer architecture

什么是 Llama 2

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

Amazon EC2 Inf2 实例概览

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

大型语言模型 (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 模型，这使得张量可以在多个 NeuronCore 上进行分片（即张量并行），并进行性能优化，例如 Neuron 硬件的并行上下文编码和 KV 缓存。Llama 2 在 XLA HLO 中的源代码可以在这里找到。

Llama 2 通过 LlamaForSampling 类在 Transformers Neuron 中得到支持。Transformers Neuron 提供了与 Hugging Face 模型无缝的用户体验，以在 Inf2 实例上提供优化的推理。更多详细信息可以在 Transforms 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)
  1. 使用以下方法从保存序列化检查点的本地目录加载和编译模型。要加载 Llama 2 模型，我们使用 Transformers Neuron 中的 LlamaForSampling。请注意，环境变量 NEURON_RT_NUM_CORES 指定运行时要使用的 NeuronCore 数量，它应与为模型指定的张量并行 (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()
  1. 最后，让我们在编译好的模型上运行推理。请注意，`sample` 函数的输入和输出都是令牌序列。
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 中的推理优化

张量并行

Latency with different TP degrees

Transformer Neuron 在多个 NeuronCore 之间实现并行张量操作。我们将用于推理的核心数表示为 TP 度。TP 度越大，内存带宽越高，从而导致延迟越低，因为 LLM 令牌生成是内存-IO 密集型工作负载。随着 TP 度的增加，推理延迟显著降低，我们的结果显示，随着 TP 度从 2 增加到 24，整体速度提高了约 4 倍。对于 Llama-2 7B 模型，延迟从 2 个核心的 30.1 毫秒/令牌降低到 24 个核心的 7.9 毫秒/令牌；类似地，对于 Llama-2 13B 模型，它从 57.3 毫秒/令牌降低到 11.1 毫秒/令牌。

并行上下文编码

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

E2E latency

KV 缓存

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

基准测试结果

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

延迟

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

Latency on inf2

吞吐量

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

E2E throughput

成本

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

Cost on inf2

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

Cost on Llama

在 EC2 Inf2 实例上使用 TorchServe 提供 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. initialize 函数用于加载模型。在这里，Neuron SDK 将首次编译模型，并将预编译的模型保存在由 `NEURONX_CACHE` 和 `NEURONX_DUMP_TO` 指定的目录中。第一次之后，后续运行将检查是否已存在预编译的模型工件。如果存在，它将跳过模型编译。一旦模型加载完成，我们就会启动预热推理请求，以便编译后的版本被缓存。当利用 神经元持久缓存 时，它可以显著减少模型加载延迟，确保后续推理运行迅速。
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,
)
  1. `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 模型

在本节中，我们将在 SageMaker 端点上部署 Llama-2 13B 模型，使用 PyTorch Neuronx 容器 和 ml.inf2.24xlarge 托管实例，该实例有 6 个 Inferentia2 加速器，与我们的模型配置 `model_config.yaml` 处理程序的设置 – `tp_degree: 12` 对应。鉴于我们已经使用 torch-model-archiver 将所有模型工件打包到一个文件夹中并上传到 S3 存储桶，我们现在将使用 SageMaker Python SDK 创建一个 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"},
)
  1. 部署 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
)
  1. 在 SageMaker 上运行流式响应推理 当端点处于服务中时，您可以使用 `invoke_endpoint_with_response_stream` API 调用来调用模型。此功能可以将每个生成的令牌返回给用户，从而增强用户体验。当生成整个序列耗时时，它尤其有益。
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！