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 によると、Inferentia は、同等の現行世代の GPU ベースの Amazon EC2 インスタンスと比較して、推論あたりのコストを最大 80% 低減し、スループットを最大 2.3 倍向上させることができます。

そのアーキテクチャは、チップ内のカスタムアクセラレータである "Neuron Cores" を利用しています。各 Inferentia チップには 4 つの Neuron Cores が含まれており、ユーザーは以下のいずれかを選択できます：

スループットを最大化するために、コアごとに 1 つのモデルをロードする。
レイテンシを最小限に抑えるために、全コアにまたがって 1 つのモデルをロードする。

BERT デプロイのための技術的ワークフロー

BERT のようなモデル（例：distilbert-base-uncased-finetuned-sst-2-english）を AWS Inferentia にデプロイするには、多段階のコンパイルおよびデプロイプロセスが必要です。

1. AWS Neuron SDK によるモデル変換

AWS Neuron SDK は、PyTorch および TensorFlow モデルを EC2 Inf1 インスタンス用の Neuron 互換フォーマットに変換するためのディープラーニングコンパイラとランタイムを提供します。

AWS Neuron SDK は動的形状（dynamic shapes）をサポートしていないため、コンパイルおよび推論中の入力サイズは静的である必要があります。例えば、モデルがバッチサイズ 1、シーケンス長 128 でコンパイルされている場合、その正確な形状の入力のみを処理できます。

2. カスタム推論スクリプトの作成

Hugging Face Inference Toolkit は多くのモデルに対してコードなしのデプロイをサポートしていますが、AWS Inferentia は現在、カスタムの inference.py スクリプトを必要とします。このスクリプトでは、以下を定義する必要があります：

model_fn: トークナイザー、Neuron モデル、およびモデル構成をロードするため。
predict_fn: 入力埋め込みを処理し、静的なシーケンス長に一致するように（パディングと切り捨てによって）確実にし、予測を実行するため。

スループットを最大化するために、環境変数 NEURON_RT_NUM_CORES=1 を使用して、各 HTTP ワーカーが単一の Neuron コアを利用するようにします。

3. SageMaker デプロイ

モデルがコンパイルされ、推論スクリプトが作成されたら、アーティファクトは model.tar.gz ファイルにアーカイブされ、Amazon S3 にアップロードされます。モデルは、Amazon SageMaker の HuggingFaceModel クラスを使用して、特に ml.inf1.xlarge インスタンスをターゲットとして、リアルタイム推論エンドポイントとしてデプロイされます。

パフォーマンス結果

10,000 件の同期リクエストによる負荷テストにおいて、BERT モデルはシーケンス長 128 の場合に平均レイテンシ 5-6ms を達成しました。

Hugging Face は、このセットアップが CPU ベースの推論よりも高速であり、Neuron Cores にわたって 4 つのモデルを並列に実行することで GPU よりも高いスループットを提供すると結論付けています。これにより、AWS Inferentia が、テキスト分類、トークン分類、および質問回答などのエンコーダータスクに BERT のような Transformers を使用している企業にとって、実行可能な選択肢となります。