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

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

笔记本：sagemaker/18_inferentia_inference

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

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

解决这个挑战的AWS方法是设计了一款专为优化推理工作负载而设计的自定义机器学习芯片，名为AWS Inferentia。AWS表示AWS Inferentia“比当前世代基于GPU的Amazon EC2实例的每次推理成本低80％，并且吞吐量高出2.3倍”。

AWS Inferentia实例与GPU相比的真正价值在于每个设备上提供的多个Neuron Core。Neuron Core是AWS Inferentia内部的自定义加速器。每个Inferentia芯片配备了4个Neuron Core。这使您可以在每个核心上加载1个模型（以实现高吞吐量）或在所有核心上加载1个模型（以实现较低的延迟）。

教程

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

您可以在这里找到笔记本：sagemaker/18_inferentia_inference

您将学习如何：

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

让我们开始吧！🚀

如果您要在本地环境中使用Sagemaker（而不是SageMaker Studio或Notebook Instances），您需要访问具有Sagemaker所需权限的IAM角色。您可以在此处了解更多信息。

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

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

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

提示：如果您使用的是Amazon SageMaker Notebook Instances或Studio，可以选择conda_python3 conda内核。

# 将Pip仓库设置为指向Neuron仓库
!pip config set global.extra-index-url https://pip.repos.neuron.amazonaws.com

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

安装完Neuron SDK后，我们可以加载和转换我们的模型。使用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  # 解决protobuf版本冲突问题
import torch
import torch.neuron
from transformers import AutoTokenizer, AutoModelForSequenceClassification

# 加载tokenizer和model
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForSequenceClassification.from_pretrained(model_id, torchscript=True)

# 创建最大长度为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())

# 使用torch.neuron.trace编译model并更新配置
model_neuron = torch.neuron.trace(model, neuron_inputs)
model.config.update({"traced_sequence_length": max_length})

# 保存tokenizer、neuron model和配置以供后续使用
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. 创建用于文本分类的自定义inference.py脚本

Hugging Face推理工具包支持在🤗 Transformers的pipeline功能之上进行零代码部署。这使用户可以在没有推理脚本的情况下部署Hugging Face transformers [示例]。

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

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

要使用推理脚本，我们需要创建一个inference.py脚本。在我们的示例中，我们将覆盖model_fn以加载我们的neuron model，并覆盖predict_fn以创建一个文本分类pipeline。

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

!mkdir code

我们使用NEURON_RT_NUM_CORES=1来确保每个HTTP worker使用1个Neuron核心以最大化吞吐量。

%%writefile code/inference.py

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

# 使用每个worker一个neuron核心
os.environ["NEURON_RT_NUM_CORES"] = "1"

# 保存的权重名称
AWS_NEURON_TRACED_WEIGHTS_NAME = "neuron_model.pt"

def model_fn(model_dir):
    # 从model_dir加载tokenizer和neuron model
    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):
    # 解构model、tokenizer和model config
    model, tokenizer, model_config = model_tokenizer_model_config

    # 为输入创建embeddings
    inputs = data.pop("inputs", data)
    embeddings = tokenizer(
        inputs,
        return_tensors="pt",
        max_length=model_config.traced_sequence_length,
        padding="max_length",
        truncation=True,
    )
    # 转换为neuron model的元组
    neuron_inputs = tuple(embeddings.values())

    # 进行预测
    with torch.no_grad():
        predictions = model(*neuron_inputs)[0]
        scores = torch.nn.Softmax(dim=1)(predictions)

    # 返回可json序列化的字典
    return [{"label": model_config.id2label[item.argmax().item()], "score": item.max().item()} for item in scores]

3. 创建并上传neuron model和推理脚本到Amazon S3

在将neuron model部署到Amazon SageMaker之前，我们需要创建一个包含所有模型工件的model.tar.gz存档，保存在tmp/中，例如neuron_model.pt，并将其上传到Amazon S3。

为了实现这一点，我们需要设置权限。

import sagemaker
import boto3
sess = sagemaker.Session()
# sagemaker 会话 bucket -> 用于上传数据、模型和日志
# 如果不存在，sagemaker 会自动创建该 bucket
sagemaker_session_bucket=None
if sagemaker_session_bucket is None and sess is not None:
    # 如果没有给定 bucket 名称，则设置为默认 bucket
    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 角色 arn: {role}")
print(f"sagemaker bucket: {sess.default_bucket()}")
print(f"sagemaker 会话区域: {sess.boto_region_name}")

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

# 将 inference.py 复制到模型目录下的 code/ 目录中。
!cp -r code/ tmp/code/
# 使用所有模型工件和 inference.py 脚本创建一个 model.tar.gz 归档文件。
%cd tmp
!tar zcvf model.tar.gz *
%cd ..

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

from sagemaker.s3 import S3Uploader

# 创建 S3 URI
s3_model_path = f"s3://{sess.default_bucket()}/{model_id}"

# 上传 model.tar.gz
s3_model_uri = S3Uploader.upload(local_path="tmp/model.tar.gz",desired_s3_uri=s3_model_path)
print(f"模型工件已上传至 {s3_model_uri}")

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

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

from sagemaker.huggingface.model import HuggingFaceModel

# 创建 Hugging Face 模型类
huggingface_model = HuggingFaceModel(
   model_data=s3_model_uri,       # 模型和脚本的路径
   role=role,                    # 具有创建端点权限的 IAM 角色
   transformers_version="4.12",  # 使用的 transformers 版本
   pytorch_version="1.9",        # 使用的 pytorch 版本
   py_version='py37',            # 使用的 python 版本
)

# 让 SageMaker 知道我们已经通过 neuron-cc 编译了模型
huggingface_model._is_compiled_model = True

# 部署推理端点
predictor = huggingface_model.deploy(
    initial_instance_count=1,      # 实例数量
    instance_type="ml.inf1.xlarge" # AWS Inferentia 实例
)

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

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

data = {
  "inputs": "主演们令人着迷的表演使电影保持了平衡，吸引了观众的注意力。",
}

res = predictor.predict(data=data)
res

我们成功将编译后的 BERT 部署到了 AWS Inferentia 上的 Amazon SageMaker。现在，让我们测试其性能。作为一个虚拟负载测试，我们将循环发送 10,000 个同步请求到我们的端点。

# 发送 10000 个请求
for i in range(10000):
    resp = predictor.predict(
        data={"inputs": "这是一段迷人且常常令人感动的旅程。"}
    )

让我们在 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}")

我们的BERT模型的平均延迟为5-6毫秒，适用于序列长度为128。

图1. 模型延迟

删除模型和端点

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

predictor.delete_model()
predictor.delete_endpoint()

结论

我们成功地将一个普通的Hugging Face Transformers模型编译为适用于AWS Inferentia的Neuron模型。之后，我们使用新的Hugging Face推理DLC将我们的Neuron模型部署到Amazon SageMaker。我们每个神经元核心实现了5-6毫秒的延迟，这在延迟方面比CPU更快，并且比GPU实现了更高的吞吐量，因为我们同时运行了4个模型。

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

我们计划在未来进行更详细的关于transformers成本和性能的案例研究，敬请关注！

如果您想了解更多关于加速transformers的信息，您还应该查看Hugging Face优化。

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

Habana Labs与Hugging Face合作加速Transforme...
机器学习专家 – Lewis Tunstall
CO2排放和🤗 Hub：引领变革
使用Habana Gaudi开始使用Transformers
机器学习洞察总监
深度强化学习简介
人工智能

人工智能

探索人工智能的世界及AI的未来潜力