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

notebook：sagemaker/18_inferentia_inference

BERT 和 Transformers 的采用持续增长。基于 Transformer 的模型现在不仅在自然语言处理领域取得了最先进的性能，在计算机视觉、语音和时间序列领域也同样如此。💬 🖼 🎤 ⏳

公司正逐步从实验和研究阶段转向生产阶段，以便将 Transformer 模型用于大规模工作负载。但默认情况下，与传统的机器学习算法相比，BERT 及其衍生模型相对较慢、较大且复杂。加速 Transformers 和 BERT 是且将成为未来一个有趣的挑战。

AWS 应对这一挑战的方法是设计一款专为优化推理工作负载而设计的定制机器学习芯片，名为 AWS Inferentia。AWS 表示，AWS Inferentia “与同代基于 GPU 的 Amazon EC2 实例相比，每次推理的成本降低了高达 80%，吞吐量提高了高达 2.3 倍。”

与 GPU 相比，AWS Inferentia 实例的真正价值在于每个设备上可用的多个 Neuron Core。Neuron Core 是 AWS Inferentia 内部的定制加速器。每个 Inferentia 芯片都带有 4 个 Neuron Core。这使您可以在每个核心上加载 1 个模型（以实现高吞吐量），或将 1 个模型分布在所有核心上（以实现更低延迟）。

教程

在这个端到端的教程中，您将学习如何使用 Hugging Face Transformers、Amazon SageMaker 和 AWS Inferentia 加速 BERT 文本分类的推理速度。

您可以在这里找到 notebook：sagemaker/18_inferentia_inference

您将学习如何

• 1. 将您的 Hugging Face Transformer 转换为 AWS Neuron
• 2. 为 text-classification 创建自定义 inference.py 脚本
• 3. 创建 Neuron 模型和推理脚本并将其上传到 Amazon S3
• 4. 在 Amazon SageMaker 上部署实时推理端点
• 5. 在 Inferentia 上运行并评估 BERT 的推理性能

让我们开始吧！🚀

---

如果您打算在本地环境（而非 SageMaker Studio 或 Notebook 实例）中使用 Sagemaker，您需要一个具有 Sagemaker 所需权限的 IAM 角色。您可以在这里找到更多相关信息。

1. 将您的 Hugging Face Transformer 转换为 AWS Neuron

我们将使用 适用于 AWS Inferentia 的 AWS Neuron SDK。Neuron SDK 包含一个深度学习编译器、运行时和工具，用于将 PyTorch 和 TensorFlow 模型转换和编译为与 Neuron 兼容的模型，这些模型可以在 EC2 Inf1 实例上运行。

作为第一步，我们需要安装 Neuron SDK 和所需的软件包。

提示：如果您正在使用 Amazon SageMaker Notebook 实例或 Studio，您可以使用 conda_python3 conda 内核。

# Set Pip repository to point to the Neuron repository
!pip config set global.extra-index-url https://pip.repos.neuron.amazonaws.com

# Install Neuron PyTorch
!pip install torch-neuron==1.9.1.* neuron-cc[tensorflow] sagemaker>=2.79.0 transformers==4.12.3 --upgrade

安装完 Neuron SDK 后，我们可以加载并转换我们的模型。Neuron 模型使用 torch_neuron 及其 trace 方法进行转换，类似于 torchscript。您可以在我们的文档中找到更多信息。

为了能够转换我们的模型，我们首先需要从 hf.co/models 中选择我们想要用于文本分类管道的模型。对于本例，我们选择 distilbert-base-uncased-finetuned-sst-2-english，但可以轻松地换成其他类似 BERT 的模型。

model_id = "distilbert-base-uncased-finetuned-sst-2-english"

在撰写本文时，AWS Neuron SDK 不支持动态形状，这意味着输入大小在编译和推理时需要是静态的。

简单来说，这意味着如果模型是使用例如批量大小为 1 和序列长度为 16 的输入进行编译的，那么该模型只能对具有相同形状的输入运行推理。

使用 t2.medium 实例时，编译大约需要 3 分钟。

import os
import tensorflow  # to workaround a protobuf version conflict issue
import torch
import torch.neuron
from transformers import AutoTokenizer, AutoModelForSequenceClassification

# load tokenizer and model
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForSequenceClassification.from_pretrained(model_id, torchscript=True)

# create dummy input for max length 128
dummy_input = "dummy input which will be padded later"
max_length = 128
embeddings = tokenizer(dummy_input, max_length=max_length, padding="max_length",return_tensors="pt")
neuron_inputs = tuple(embeddings.values())

# compile model with torch.neuron.trace and update config
model_neuron = torch.neuron.trace(model, neuron_inputs)
model.config.update({"traced_sequence_length": max_length})

# save tokenizer, neuron model and config for later use
save_dir="tmp"
os.makedirs("tmp",exist_ok=True)
model_neuron.save(os.path.join(save_dir,"neuron_model.pt"))
tokenizer.save_pretrained(save_dir)
model.config.save_pretrained(save_dir)

2. 为 text-classification 创建自定义 inference.py 脚本

Hugging Face 推理工具包支持基于 🤗 Transformers 的 pipeline 功能实现零代码部署。这允许用户无需编写推理脚本即可部署 Hugging Face Transformers 模型 [示例]。

目前，AWS Inferentia 尚不支持此功能，这意味着我们需要提供一个 inference.py 脚本来运行推理。

如果您对 Inferentia 的零代码部署支持感兴趣，请在论坛上告诉我们。

---

要使用推理脚本，我们需要创建一个 inference.py 脚本。在我们的示例中，我们将重写 model_fn 来加载我们的 Neuron 模型，并重写 predict_fn 来创建一个文本分类管道。

如果您想了解更多关于 inference.py 脚本的信息，请查看这个示例。它解释了 model_fn 和 predict_fn 等内容。

!mkdir code

我们使用 NEURON_RT_NUM_CORES=1 来确保每个 HTTP 工作进程使用 1 个 Neuron Core，以最大化吞吐量。

%%writefile code/inference.py

import os
from transformers import AutoConfig, AutoTokenizer
import torch
import torch.neuron

# To use one neuron core per worker
os.environ["NEURON_RT_NUM_CORES"] = "1"

# saved weights name
AWS_NEURON_TRACED_WEIGHTS_NAME = "neuron_model.pt"

def model_fn(model_dir):
    # load tokenizer and neuron model from model_dir
    tokenizer = AutoTokenizer.from_pretrained(model_dir)
    model = torch.jit.load(os.path.join(model_dir, AWS_NEURON_TRACED_WEIGHTS_NAME))
    model_config = AutoConfig.from_pretrained(model_dir)

    return model, tokenizer, model_config

def predict_fn(data, model_tokenizer_model_config):
    # destruct model, tokenizer and model config
    model, tokenizer, model_config = model_tokenizer_model_config

    # create embeddings for inputs
    inputs = data.pop("inputs", data)
    embeddings = tokenizer(
        inputs,
        return_tensors="pt",
        max_length=model_config.traced_sequence_length,
        padding="max_length",
        truncation=True,
    )
    # convert to tuple for neuron model
    neuron_inputs = tuple(embeddings.values())

    # run prediciton
    with torch.no_grad():
        predictions = model(*neuron_inputs)[0]
        scores = torch.nn.Softmax(dim=1)(predictions)

    # return dictonary, which will be json serializable
    return [{"label": model_config.id2label[item.argmax().item()], "score": item.max().item()} for item in scores]

3. 创建 Neuron 模型和推理脚本并将其上传到 Amazon S3

在我们能将 Neuron 模型部署到 Amazon SageMaker 之前，我们需要创建一个 model.tar.gz 归档文件，其中包含我们保存在 tmp/ 目录下的所有模型构件（例如 neuron_model.pt），并将其上传到 Amazon S3。

为此，我们需要设置我们的权限。

import sagemaker
import boto3
sess = sagemaker.Session()
# sagemaker session bucket -> used for uploading data, models and logs
# sagemaker will automatically create this bucket if it not exists
sagemaker_session_bucket=None
if sagemaker_session_bucket is None and sess is not None:
    # set to default bucket if a bucket name is not given
    sagemaker_session_bucket = sess.default_bucket()

try:
    role = sagemaker.get_execution_role()
except ValueError:
    iam = boto3.client('iam')
    role = iam.get_role(RoleName='sagemaker_execution_role')['Role']['Arn']

sess = sagemaker.Session(default_bucket=sagemaker_session_bucket)

print(f"sagemaker role arn: {role}")
print(f"sagemaker bucket: {sess.default_bucket()}")
print(f"sagemaker session region: {sess.boto_region_name}")

接下来，我们创建我们的 model.tar.gz。inference.py 脚本将被放置在一个 code/ 文件夹中。

# copy inference.py into the code/ directory of the model directory.
!cp -r code/ tmp/code/
# create a model.tar.gz archive with all the model artifacts and the inference.py script.
%cd tmp
!tar zcvf model.tar.gz *
%cd ..

现在我们可以使用 sagemaker 将我们的 model.tar.gz 上传到我们的会话 S3 存储桶中。

from sagemaker.s3 import S3Uploader

# create s3 uri
s3_model_path = f"s3://{sess.default_bucket()}/{model_id}"

# upload model.tar.gz
s3_model_uri = S3Uploader.upload(local_path="tmp/model.tar.gz",desired_s3_uri=s3_model_path)
print(f"model artifcats uploaded to {s3_model_uri}")

4. 在 Amazon SageMaker 上部署实时推理端点

在我们将 model.tar.gz 上传到 Amazon S3 之后，我们可以创建一个自定义的 HuggingfaceModel。这个类将用于在 Amazon SageMaker 上创建和部署我们的实时推理端点。

from sagemaker.huggingface.model import HuggingFaceModel

# create Hugging Face Model Class
huggingface_model = HuggingFaceModel(
   model_data=s3_model_uri,       # path to your model and script
   role=role,                    # iam role with permissions to create an Endpoint
   transformers_version="4.12",  # transformers version used
   pytorch_version="1.9",        # pytorch version used
   py_version='py37',            # python version used
)

# Let SageMaker know that we've already compiled the model via neuron-cc
huggingface_model._is_compiled_model = True

# deploy the endpoint endpoint
predictor = huggingface_model.deploy(
    initial_instance_count=1,      # number of instances
    instance_type="ml.inf1.xlarge" # AWS Inferentia Instance
)

5. 在 Inferentia 上运行并评估 BERT 的推理性能

.deploy() 方法返回一个 HuggingFacePredictor 对象，可用于请求推理。

data = {
  "inputs": "the mesmerizing performances of the leads keep the film grounded and keep the audience riveted .",
}

res = predictor.predict(data=data)
res

我们成功地将我们编译为 Neuron 的 BERT 模型部署到了 Amazon SageMaker 上的 AWS Inferentia。现在，让我们来测试它的性能。作为一个简单的负载测试，我们将循环发送 10,000 个同步请求到我们的端点。

# send 10000 requests
for i in range(10000):
    resp = predictor.predict(
        data={"inputs": "it 's a charming and often affecting journey ."}
    )

让我们在 CloudWatch 中检查性能。

print(f"https://console.aws.amazon.com/cloudwatch/home?region={sess.boto_region_name}#metricsV2:graph=~(metrics~(~(~'AWS*2fSageMaker~'ModelLatency~'EndpointName~'{predictor.endpoint_name}~'VariantName~'AllTraffic))~view~'timeSeries~stacked~false~region~'{sess.boto_region_name}~start~'-PT5M~end~'P0D~stat~'Average~period~30);query=~'*7bAWS*2fSageMaker*2cEndpointName*2cVariantName*7d*20{predictor.endpoint_name}")

对于序列长度为 128，我们的 BERT 模型的平均延迟为 5-6ms。

删除模型和端点

为了清理，我们可以删除模型和端点。

predictor.delete_model()
predictor.delete_endpoint()

结论

我们成功地将一个原生的 Hugging Face Transformers 模型编译成与 AWS Inferentia 兼容的 Neuron 模型。之后，我们使用新的 Hugging Face Inference DLC 将我们的 Neuron 模型部署到 Amazon SageMaker。我们成功地在每个 Neuron Core 上实现了 5-6ms 的延迟，这在延迟方面比 CPU 更快，并且由于我们并行运行了 4 个模型，实现了比 GPU 更高的吞吐量。

如果您或您的公司目前正在使用类似 BERT 的 Transformer 模型进行编码器任务（如文本分类、token 分类、问答等），并且延迟满足您的要求，您应该切换到 AWS Inferentia。这不仅可以节省成本，还可以提高模型的效率和性能。

我们计划在未来对 Transformers 的成本性能进行更详细的案例研究，敬请期待！

此外，如果您想了解更多关于加速 Transformers 的信息，您也应该查看 Hugging Face 的 optimum。

---

感谢阅读！如果您有任何问题，请随时通过 Github 或论坛与我联系。您也可以在 Twitter 或 LinkedIn 上与我联系。