使用 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 上與我聯絡。