使用 Hugging Face Transformers 和 AWS Inferentia 加速 BERT 推理

TL;DR

Hugging Face 展示了如何通过使用 AWS Neuron SDK 编译 Transformers 模型，并将其部署在 Amazon SageMaker 上的 AWS Inferentia 芯片上，从而加速 BERT 推理。这种方法将序列长度为 128 时的延迟降低至 5-6ms，并且与传统的基于 GPU 的 EC2 实例相比，提供了更高的吞吐量和更低的成本。

用于 Transformer 加速的 AWS Inferentia

AWS Inferentia 是一款专门为优化推理工作负载而设计的定制机器学习芯片。根据 AWS 的说法，与当前世代的同类基于 GPU 的 Amazon EC2 实例相比，Inferentia 可以提供高达 80% 的更低推理成本和高达 2.3 倍的更高吞吐量。

该架构利用了“Neuron Cores”，即芯片内的定制加速器。每个 Inferentia 芯片包含四个 Neuron Cores，允许用户：

• 每个核心加载一个模型以最大化吞吐量。
• 在所有核心上加载一个模型以最小化延迟。

BERT 部署的技术工作流

在 AWS Inferentia 上部署类似 BERT 的模型（例如 distilbert-base-uncased-finetuned-sst-2-english）涉及一个多步骤的编译和部署过程。

1. 通过 AWS Neuron SDK 进行模型转换

AWS Neuron SDK 提供了一个深度学习编译器和运行时，用于将 PyTorch 和 TensorFlow 模型转换为适用于 EC2 Inf1 实例的 neuron-compatible 格式。

由于 AWS Neuron SDK 不支持动态形状（dynamic shapes），因此输入尺寸在编译和推理期间必须是静态的。例如，如果一个模型是以 batch size 为 1 且序列长度为 128 进行编译的，那么它只能处理该精确形状的输入。

2. 自定义推理脚本编写

虽然 Hugging Face Inference Toolkit 支持许多模型的零代码部署，但 AWS Inferentia 目前需要一个自定义的 inference.py 脚本。该脚本必须定义：

• model_fn: 用于加载 tokenizer、neuron 模型和模型配置。
• predict_fn: 用于处理输入 embeddings、确保它们与静态序列长度匹配（通过 padding 和 truncation），并执行预测。

为了最大化吞吐量，使用了环境变量 NEURON_RT_NUM_CORES=1 来确保每个 HTTP worker 利用单个 Neuron core。

3. SageMaker 部署

一旦模型被编译且推理脚本被创建，这些产物将被归档到 model.tar.gz 文件中并上传到 Amazon S3。然后，使用 Amazon SageMaker 中的 HuggingFaceModel 类，专门针对 ml.inf1.xlarge 实例进行部署，将其作为实时推理端点。

性能结果

在包含 10,000 次同步请求的负载测试中，BERT 模型在序列长度为 128 时实现了 5-6ms 的平均延迟。

Hugging Face 得出结论，通过在 Neuron Cores 上并行运行四个模型，此设置比基于 CPU 的推理更快，并且比 GPU 提供了更高的吞吐量。这使得 AWS Inferentia 对于使用类似 BERT 的 Transformers 进行编码器任务（如文本分类、token 分类和问答）的公司来说是一个可行的选择。