最近，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 在来自公开来源的 2000 亿个数据标记上进行了预训练。微调模型旨在用于类似助手的聊天，而预训练模型可以适应各种自然语言生成任务。无论开发者使用哪种版本的模型，Meta 的负责任使用指南都可以帮助指导可能需要的额外微调，以定制和优化模型并采取适当的安全缓解措施。
Amazon EC2 Inf2 实例概述
配备 Inferentia2 的 Amazon EC2 Inf2 实例提供 3 倍的计算能力，4 倍的加速器内存，与第一代 Inf1 实例相比，吞吐量最高提高 4 倍，延迟降低至 10 倍。
大型语言模型（LLM）推理是内存密集型工作负载，性能随着更多加速器内存带宽的增加而提升。Inf2 实例是 Amazon EC2 中唯一针对推理优化的实例，提供高速加速器互连（NeuronLink），从而实现高性能的大LLM模型部署，同时实现成本效益的分布式推理。现在您可以在 Inf2 实例上高效且经济地部署跨多个加速器的亿级LLMs。
Inferentia2 支持 FP32、TF32、BF16、FP16、UINT8 以及新的可配置 FP8（cFP8）数据类型。AWS Neuron 可以将高精度 FP32 和 FP16 模型自动转换为低精度数据类型，同时优化精度和性能。自动转换通过消除对低精度重新训练的需求，并允许使用更小的数据类型进行高性能推理，从而缩短上市时间。
为了使其灵活且可扩展，以便部署不断演变的深度学习模型，Inf2 实例具有硬件优化和软件支持，包括动态输入形状以及通过标准 PyTorch 自定义操作符编程接口编写的自定义操作符。
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 实例上提供优化的推理。更多详情可从 Transformers 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指定了在运行时使用的 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()
- 最后，让我们在编译好的模型上运行推理。注意，
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 中的推理优化
张量并行
Transformer 神经元在多个 NeuronCores 上实现了并行张量操作。我们用 TP 度表示用于推理的核心数量。较大的 TP 度提供更高的内存带宽，从而降低延迟，因为LLM令牌生成是一个内存-IO 密集型的工作负载。随着 TP 度的增加，推理延迟显著降低，我们的结果显示，从 2 度增加到 24 度，整体速度提高了约 4 倍。对于 Llama-2 7B 模型，延迟从 2 个核心的 30.1 ms/令牌降低到 24 个核心的 7.9 ms/令牌；同样地，对于 Llama-2 13B 模型，它从 57.3 ms/令牌降低到 11.1 ms/令牌。
并行上下文编码
在 Transformer 架构中，通过自回归采样以顺序方式生成标记，而输入提示标记可以与并行上下文编码并行处理。这可以显著减少通过自回归采样生成标记之前对输入提示上下文编码的延迟。默认情况下，参数 context_length_estimate 将被设置为一系列 2 的幂次方数，旨在覆盖广泛的上下文长度。根据用例，它可以设置为自定义数字。这可以在创建 Llama 2 模型时使用 LlamaForSampling.from_pretrained 完成。我们分析了输入标记长度对端到端（E2E）延迟的影响。如图所示，由于并行上下文编码，使用 Llama-2 7B 模型进行文本生成的延迟仅随着更大的输入提示略有增加。
KV 缓存
自注意力块使用 KV 向量执行自注意力操作。同时，KV 向量通过标记嵌入和 KV 权重计算得出，与标记相关联。在简单的实现中，对于每个生成的标记，都会重新计算整个 KV 缓存，但这会降低性能。因此，Transformers Neuron 库重用之前计算的 KV 向量以避免不必要的计算，也称为 KV 缓存，以减少自回归采样阶段的延迟。
基准测试结果
我们在不同条件下对 Llama-2 7B 和 13B 模型进行了延迟和成本的基准测试，即输出标记的数量、实例类型。除非特别说明，我们使用数据类型‘bf16’和批处理大小为 1，因为这通常是聊天机器人、代码助手等实时应用的常见配置。
延迟
以下图表显示了在 inf2.48xlarge 实例上每个 token 的延迟，TP 度为 24。在这里，每个输出 token 的延迟计算为端到端延迟除以输出 token 的数量。我们的实验表明，Llama-2 7B 生成 256 个 token 的端到端延迟比其他同类推理优化的 EC2 实例快 2 倍。
吞吐量
现在我们展示了 inf2.48xlarge 实例可以提供的 Llama-2 7B 和 13B 模型每秒生成的 token 数量。在 TP 度为 24 的情况下，充分利用所有 24 个 NeuronCores，Llama-2 7B 和 13B 模型分别可以达到 130 个 token/秒和 90 个 token/秒。
成本
对于以延迟优先的应用程序，我们展示了在 inf2.48xlarge 实例上托管 Llama-2 模型的成本，每 1000 个 token 为 0.011 美元，7B 和 13B 模型分别为 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 提供 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 and 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 自定义处理器。
- 初始化函数用于加载模型。在这里，Neuron SDK 将首次编译模型，并将预编译的模型保存到由
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,
)
- 推理函数调用 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,
)
打包模型工件
将所有模型工件打包到一个文件夹 llama-2-13b-neuronx-b1 中，使用 torch-model-archiver .
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 模型库中。
在 SageMaker Inf2 实例上使用 TorchServe 部署 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 实时端点，使用的是部署未压缩模型的方法。使用 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 调用调用模型。此功能使每个生成的标记都能返回给用户，从而提升用户体验。当生成整个序列耗时较长时，尤其有益。
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 模型推理，并使用 TorchServe 在 Amazon SageMaker 上部署 Llama 2 模型服务，在 EC2 Inf2 实例上。我们展示了使用 Inferentia2 的优势——低延迟和低成本——这是由 AWS Neuron SDK 中的优化实现的，包括张量并行、并行上下文编码和 KV 缓存，特别是对于LLM推理。为了保持最新，请关注 AWS Neuron 的最新发布以获取新功能。
今天开始使用 Llama 2 在 EC2 和 SageMaker 上的示例，并关注如何优化 Inf2 上的 Llama 70B！
