近期，Llama 2 的发布在机器学习社区引发了广泛关注。由 AWS Inferentia2 提供支持的 Amazon EC2 Inf2 实例现已支持 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 实例概述
采用 Inferentia2 的 Amazon EC2 Inf2 实例与第一代 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 模型，并自动将其转换为低精度数据类型，同时优化准确性和性能。自动转换 (Autocasting) 通过消除对低精度重新训练的需求，以及实现基于更小数据类型的高性能推理，从而缩短了上市时间。
为了使其能够灵活扩展以部署不断演进的深度学习模型，Inf2 实例在硬件上进行了优化，并在软件上支持动态输入形状以及通过标准 PyTorch 自定义算子编程接口编写的 C++ 自定义算子。
Transformers Neuron (transformers-neuronx)
Transformers Neuron 是一个软件包，支持 PyTorch 用户部署性能优化的 LLM 推理。它拥有一套使用 XLA 高级算子 (HLO) 实现的优化版 Transformer 模型，支持跨多个 NeuronCore 分片张量（即张量并行），并针对 Neuron 硬件提供了并行上下文编码和 KV 缓存等性能优化。XLA HLO 中的 Llama 2 源代码可点击此处获取。
Llama 2 在 Transformers Neuron 中通过 LlamaForSampling 类获得支持。Transformers Neuron 与 Hugging Face 模型无缝集成，在 Inf2 实例上提供优化后的推理体验。更多详细信息请参阅 Transforms Neuron 开发指南。在接下来的章节中，我们将解释如何使用 Transformers Neuron 部署 Llama-2 13B 模型。此示例同样适用于其他基于 Llama 的模型。
使用 Transformers Neuron 进行 Llama 2 模型推理
创建模型、编译和部署
我们在 Inf2 实例上创建、编译和部署模型只需简单的三个步骤。
- 创建 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)
- 使用以下内容从保存序列化检查点的本地目录加载并编译模型。为了加载 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()
- 最后，我们在编译后的模型上运行推理。请注意，
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 在多个 NeuronCore 之间实现了并行张量运算。我们将用于推理的核心数量称为 TP 程度 (TP degree)。由于 LLM Token 生成是内存 IO 密集型工作负载，更大的 TP 程度提供了更高的内存带宽，从而降低了延迟。随着 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 是通过称为自回归采样的顺序过程产生的，而输入提示词 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 的数量。我们的实验表明，在生成 256 个 Token 时，Llama-2 7B 的端到端延迟比其他可比较的推理优化 EC2 实例快 2 倍。
吞吐量
我们现在展示 inf2.48xlarge 实例每秒可为 Llama-2 7B 和 13B 模型生成的 Token 数量。在 TP 程度为 24、完全利用所有 24 个 NeuronCore 的情况下，Llama-2 7B 和 13B 模型分别可以达到 130 tokens/sec 和 90 tokens/sec。
成本
对于延迟优先的应用程序，我们展示了在 inf2.48xlarge 实例上托管 Llama-2 模型的成本：7B 模型为每 1000 个 Token 0.011 美元，13B 模型为每 1000 个 Token 0.016 美元，这比其他可比较的推理优化 EC2 实例节省了 3 倍的成本。请注意，我们报告的成本基于3年期预留实例价格，这是客户进行大规模生产部署时所使用的价格。
我们还比较了在 inf2.xlarge 和 inf2.48xlarge 实例上托管 Llama-2 7B 模型的成本。我们可以看到，inf2.xlarge 比 inf2.48xlarge 便宜 4 倍以上，但代价是由于 TP 程度较低，延迟会更长。例如，模型在 inf2.48xlarge 上生成 256 个输出 Token 需要 7.9 毫秒，而在 Inf2.xlarge 上则需要 30.1 毫秒。
在 EC2 Inf2 实例上使用 TorchServe 提供 Llama 2 服务
现在，我们转向模型部署。在本节中，我们将向您展示如何通过 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 微批处理 (micro-batching) 是一种并行预处理和后处理推理请求批次的机制。当后端稳定地接收数据时，它能够通过更好地利用可用加速器来实现更高的吞吐量，详见此处。对于 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 自定义处理器。
- initialize 函数用于加载模型。Neuron SDK 将在此首次编译模型，并根据
NEURONX_CACHE的设置，将预编译模型保存在NEURONX_DUMP_TO指定的目录中。此后，后续运行将检查是否存在预编译的模型工件。如果存在，它将跳过模型编译。模型加载完成后，我们启动预热推理请求，以便缓存编译后的版本。当利用 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,
)
- 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”，预编译的模型应已保存在与 Neuron SDK 版本紧密耦合的 llama-2-13b-neuronx-b1/neuron_cache 文件夹中。然后，将 llama-2-13b-neuronx-b1 文件夹上传到您的 S3 存储桶，以便后续在生产部署中使用。本文中的 Llama-2 13B 模型工件可点击此处（位于 TorchServe 模型库中，关联 Neuron SDK 2.13.2 版本）找到。
使用 TorchServe 在 SageMaker Inf2 实例上部署 Llama-2 13B 模型
在本节中，我们将使用 PyTorch Neuronx 容器在 SageMaker 终端节点上部署 Llama-2 13B 模型。我们使用 ml.inf2.24xlarge 托管实例，该实例拥有 6 个 Inferentia2 加速器，对应于我们 model_config.yaml 处理器的设置 – tp_degree: 12。鉴于我们已使用 torch-model-archiver 将所有模型工件打包并上传到 S3 存储桶，现在我们将使用 SageMaker Python SDK 创建一个 SageMaker 模型，并使用部署未压缩模型方法将其部署到 SageMaker 实时终端节点。这种部署方式的关键优势在于速度，您无需花费精力在基础设施上，即可获得功能完备、生产就绪的终端节点，以及安全的 RESTful 接口。在 SageMaker 上部署模型并运行推理有 3 个步骤。笔记本示例可在此处找到。
- 创建 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"},
)
- 部署 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
)
- 在 SageMaker 上运行流式响应推理。当终端节点处于服务状态时，您可以使用
invoke_endpoint_with_response_streamAPI 调用来调用模型。此功能允许将每个生成的 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=' ')
推理示例：
输入
“Today the weather is really nice and I am planning on”
输出
“Today the weather is really nice and I am planning on going to the beach. I am going to take my camera and take some pictures of the beach. I am going to take pictures of the sand, the water, and the people. I am also going to take pictures of the sunset. I am really excited to go to the beach and take pictures.
The beach is a great place to take pictures. The sand, the water, and the people are all great subjects for pictures. The sunset is also a great subject for pictures.”
结论
在本文中，我们展示了如何使用 Transformers Neuron 运行 Llama 2 模型推理，以及如何通过 Amazon SageMaker 在 EC2 Inf2 实例上使用 TorchServe 部署 Llama 2 模型服务。我们演示了 AWS Neuron SDK（包括张量并行、并行上下文编码和 KV 缓存）中的优化如何为 LLM 推理带来低延迟和低成本的益处。为了保持最新状态，请关注 AWS Neuron 的最新发布，了解新功能。
立即在 EC2 和 SageMaker 上开始使用 Llama 2 示例吧，并敬请期待如何在 Inf2 上优化 Llama 70B 模型！