Hugging Face TransformersとAWS Inferentiaを使用して、BERT推論を高速化する

使用するのはHugging Face TransformersとAWS Inferentiaで、BERT推論を高速化します

ノートブック：sagemaker/18_inferentia_inference

BERTとTransformersの採用はますます広がっています。Transformerベースのモデルは、自然言語処理だけでなく、コンピュータビジョン、音声、時系列でも最先端のパフォーマンスを達成しています。💬 🖼 🎤 ⏳

企業は、大規模なワークロードのためにトランスフォーマーモデルを使用するため、実験および研究フェーズから本番フェーズにゆっくりと移行しています。ただし、デフォルトでは、BERTとその仲間は、従来の機械学習アルゴリズムと比較して、比較的遅く、大きく、複雑なモデルです。TransformersとBERTの高速化は、将来的に解決すべき興味深い課題となるでしょう。

AWSはこの課題を解決するために、最適化された推論ワークロード向けに設計されたカスタムマシンラーニングチップであるAWS Inferentiaを開発しました。AWSは、AWS Inferentiaが「現行世代のGPUベースのAmazon EC2インスタンスと比較して、推論ごとのコストを最大80%低減し、スループットを最大2.3倍高める」と述べています。

• Hugging FaceでのDecision Transformersの紹介 🤗
• ~自分自身を~ 繰り返さない
• Habana LabsとHugging Faceが提携し、Transformerモデルのトレーニングを加速化する

AWS Inferentiaインスタンスの真の価値は、各デバイスに搭載された複数のNeuronコアを通じて実現されます。Neuronコアは、AWS Inferentia内部のカスタムアクセラレータです。各Inferentiaチップには4つのNeuronコアが搭載されています。これにより、高スループットのために各コアに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またはノートブックインスタンスではない場合）、Sagemakerに必要な権限を持つIAMロールへのアクセスが必要です。詳細については、こちらをご覧ください。

1. Hugging Face TransformerをAWS Neuronに変換する

AWS InferentiaのためにAWS Neuron SDKを使用します。Neuron SDKには、PyTorchおよびTensorFlowモデルをニューロン互換モデルに変換およびコンパイルするためのディープラーニングコンパイラ、ランタイム、ツールが含まれています。これらのモデルはEC2 Inf1インスタンスで実行できます。

最初のステップとして、Neuron SDKと必要なパッケージをインストールする必要があります。

ヒント：Amazon SageMakerのノートブックインスタンスまたは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をインストールした後、モデルをロードして変換することができます。Neuronモデルは、torchscriptに似たtorch_neuronのtraceメソッドを使用して変換されます。詳細については、ドキュメントを参照してください。

モデルを変換するためには、まずテキスト分類パイプラインに使用するモデルを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  # プロトコルバージョンの競合問題を回避するために使用
import torch
import torch.neuron
from transformers import AutoTokenizer, AutoModelForSequenceClassification

# トークナイザとモデルをロードする
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForSequenceClassification.from_pretrained(model_id, torchscript=True)

# マックス長さ128のダミー入力を作成
dummy_input = "後でパディングされるダミー入力"
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_neuron = torch.neuron.trace(model, neuron_inputs)
model.config.update({"traced_sequence_length": max_length})

# トークナイザ、ニューロンモデル、構成を保存して後で使用する
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 Inference Toolkitは、🤗 Transformers のパイプライン機能の上にゼロコードのデプロイメントをサポートしています。これにより、推論スクリプトを使用せずに Hugging Face transformers をデプロイできます [ 例を参照 ]。

現在、この機能はAWS Inferentiaではサポートされていませんので、推論の実行には inference.py スクリプトを提供する必要があります。

Inferentiaのゼロコードデプロイメントのサポートに興味がある場合は、フォーラムにご連絡ください。

---

推論スクリプトを使用するために、inference.py スクリプトを作成する必要があります。この例では、model_fn を上書きしてニューロンモデルをロードし、predict_fn を作成してテキスト分類パイプラインを作成します。

inference.py スクリプトについて詳しく知りたい場合は、この例を参照してください。その中で、model_fn と predict_fn について説明しています。

!mkdir code

各HTTPワーカーがスループットを最大化するために、NEURON_RT_NUM_CORES=1 を使用しています。

%%writefile code/inference.py

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

# ワーカーごとに1つのニューロンコアを使用するため
os.environ["NEURON_RT_NUM_CORES"] = "1"

# 保存されたウェイト名
AWS_NEURON_TRACED_WEIGHTS_NAME = "neuron_model.pt"

def model_fn(model_dir):
    # トークナイザとニューロンモデルを 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):
    # モデル、トークナイザ、モデル構成を分解する
    model, tokenizer, model_config = model_tokenizer_model_config

    # 入力の埋め込みを作成する
    inputs = data.pop("inputs", data)
    embeddings = tokenizer(
        inputs,
        return_tensors="pt",
        max_length=model_config.traced_sequence_length,
        padding="max_length",
        truncation=True,
    )
    # ニューロンモデル用にタプルに変換する
    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. ニューロンモデルと推論スクリプトを Amazon S3 に作成してアップロードする

Amazon SageMakerにニューロンモデルをデプロイする前に、すべてのモデルアーティファクト（neuron_model.pt など）が tmp/ に保存された model.tar.gz アーカイブを作成し、これを Amazon S3 にアップロードする必要があります。

これを行うには、アクセス許可を設定する必要があります。

import sagemaker
import boto3
sess = sagemaker.Session()
# sagemaker session bucket -> データ、モデル、ログをアップロードするために使用されます
# 存在しない場合、sagemakerは自動的にこのバケットを作成します
sagemaker_session_bucket=None
if sagemaker_session_bucket is None and sess is not None:
    # バケット名が指定されていない場合はデフォルトのバケットに設定します
    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/フォルダに配置されます。

# inference.pyをモデルディレクトリのcode/ディレクトリにコピーします。
!cp -r code/ tmp/code/
# 全てのモデルアーティファクトとinference.pyスクリプトを含むmodel.tar.gzアーカイブを作成します。
%cd tmp
!tar zcvf model.tar.gz *
%cd ..

これで、model.tar.gzをsagemakerを使用してセッションの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"model artifcats uploaded to {s3_model_uri}")

4. Amazon SageMakerでリアルタイム推論エンドポイントをデプロイする

model.tar.gzをAmazon S3にアップロードした後、カスタムのHuggingfaceModelを作成します。このクラスは、Amazon SageMaker上でリアルタイム推論エンドポイントを作成およびデプロイするために使用されます。

from sagemaker.huggingface.model import HuggingFaceModel

# Hugging Face Modelクラスを作成します
huggingface_model = HuggingFaceModel(
   model_data=s3_model_uri,       # モデルとスクリプトへのパス
   role=role,                    # Endpointを作成するためのIAMロール
   transformers_version="4.12",  # 使用されるtransformersのバージョン
   pytorch_version="1.9",        # 使用されるpytorchのバージョン
   py_version='py37',            # 使用されるPythonのバージョン
)

# SageMakerに既にモデルがコンパイルされていることを知らせます
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": "the mesmerizing performances of the leads keep the film grounded and keep the audience riveted .",
}

res = predictor.predict(data=data)
res

モデルをAWS Inferentiaにデプロイしました。パフォーマンスをテストしてみましょう。ダミーロードテストとして、10,000の同期リクエストをエンドポイントに送信します。

# 10000のリクエストを送信します
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}")

私たちのBERTモデルの平均レイテンシは、シーケンス長が128の場合には5-6msです。

図1. モデルのレイテンシ

モデルとエンドポイントの削除

クリーンアップのため、モデルとエンドポイントを削除することができます。

predictor.delete_model()
predictor.delete_endpoint()

結論

私たちは、バニラのHugging Face TransformersモデルをAWS Inferentia互換のNeuronモデルに正常にコンパイルしました。その後、新しいHugging Face Inference DLCを使用してNeuronモデルをAmazon SageMakerにデプロイしました。ニューロンコアごとに5-6msのレイテンシを達成しました。これはレイテンシの点でCPUよりも高速であり、4つのモデルを並列に実行したため、GPUよりもスループットが高くなります。

もしも、お客様またはお客様の企業がエンコーダータスク（テキスト分類、トークン分類、質問応答など）にBERTのようなTransformerを使用しており、レイテンシが要件を満たしている場合は、AWS Inferentiaに切り替えることをお勧めします。これにより、コストを節約するだけでなく、モデルの効率とパフォーマンスも向上することができます。

将来的には、Transformerのコストパフォーマンスに関するより詳細な事例研究を行う予定ですので、お楽しみに！

また、Transformerの高速化について詳しく学びたい場合は、Hugging Face Optimumもチェックしてみてください。

---

読んでいただきありがとうございます！ご質問がある場合は、Githubやフォーラムを通じて、またはTwitterやLinkedInで私に連絡することもできます。

We will continue to update VoAGI; if you have any questions or suggestions, please contact us!