AWS Inferentia上のPyTorch自然言語処理アプリケーションにおいて、12倍のスループットと最小のレイテンシーを実現

この記事は、2021年5月4日に Fabio Nonato de Paula、Mahadevan Balasubramaniam によって投稿された Achieve 12x higher throughput and lowest latency for PyTorch Natural Language Processing applications out-of-the-box on AWS Inferentia を翻訳したものです。2021年6月25日に発表された Amazon EC2 Inf1 インスタンスで新機能、パフォーマンスの向上、値下げを実現 の記事も併せてご参照下さい。

Snap、Alexa、Autodesk などの AWS のお客様は、 AWS Inferentia を使用して、さまざまな機械学習 (ML) デプロイで最高のパフォーマンスと最小のコストを達成しています。自然言語処理 (NLP) モデルは、リアルタイムおよびオフラインのバッチ処理のユースケースで人気が高まっています。当社のお客様は、サポートチャットボット、検索、ランキング、ドキュメントの要約、自然言語理解など、多くのアプリケーションにこれらのモデルをデプロイしています。AWS Inferentia を使うことで、オープンソースの NLP モデルをカスタマイズ無しですぐに実行でき、かつ最高のパフォーマンスと最小のコストを実現できます。

この記事では、AWS Inferentia 上で、レイテンシーに対する要求が厳しいリアルタイムアプリケーションと、最大スループットと最小コストが重要なパフォーマンス目標であるバッチ処理の両方で、スループットを最大化する方法について説明します。この記事では、HuggingFace Transformers の事前学習された BERT base モデルを修正せずに使用し、PyTorch フレームワークレベルでコードを1行追加するのみで、NLP ベースのソリューションをデプロイします。このソリューションでは、同じモデルを GPU にデプロイする場合と比較して、AWS Inferentia 上で 70% 低いコストで 12 倍のスループットを実現します。

AWS Inferentia で Hugging Face モデルの推論パフォーマンスを最大化するには、AWS Neuron PyTorch フレームワークインテグレーションを使用します。Neuron は、TensorFlow や PyTorch などの一般的な ML フレームワークと統合するソフトウェア開発キット (SDK) で、フレームワーク API を拡張して Amazon EC2 Inf1 インスタンスで簡単かつコスト効率よく高性能な推論を実行できるようにします。最小限のコード変更で、事前学習済みのモデルをコンパイルし、最適化して AWS Inferentia 上で実行できます。Neuron チームは、新機能やモデルの性能を向上させたアップデートを定期的にリリースしています。v1.13のリリースでは、Transformers ベースのモデルのパフォーマンスがさらに10%～15%向上し、大規模な NLP ワークロードでも、最小のレイテンシーと最大のスループットの境界を押し広げることができました。

Neuron SDK の機能を自分でテストするには、PyTorch チュートリアルの Utilizing Neuron Capabilities をご覧ください。

Neuronコアパイプラインモードの説明

Inf1 インスタンスファミリーで利用可能な各 AWS Inferentia チップ上には、4 つの Neuronコアが搭載されています。インスタンスサイズ毎に 1 ～ 16 チップが提供され、最大インスタンスサイズ inf1.24xlarge では計64 の Neuronコアが搭載されています。Neuronコアは、ニューラルネットワーク (NN) グラフのオペレーションを行うコンピュートユニットです。

パイプラインモードなしでモデルをコンパイルすると、Neuron コンパイラはサポートされる NN オペレーションを1つの Neuronコア上で実行するように最適化します。AWS Inferentiaチップをまたいで、複数の Neuronコアをグループにまとめてコンパイルモデルを実行することもできます。この構成では、AWS Inferentia チップ間で複数の Neuronコアをデータパラレルモードで使用します。つまり、最小のインスタンスサイズでも、4 つのモデルをいつでもアクティブにできます。4 つ（またはそれ以上）のモデルをデータパラレルで実装することで、多くの場合、最高のスループットと低いコストを実現します。AWS Inferentia は小さなバッチサイズでスループットを最大化するように最適化されているため、このパフォーマンス向上はレイテンシーへの影響を最小限に抑えます。

パイプラインモードでは、Neuron コンパイラは、要求された数の Neuronコア全体にわたって、単一の NN グラフの分割と配置を完全自動プロセスで最適化します。パイプライン上の Neuronコアは、モデルの重みを保持するために高速なオンチップキャッシュを使用してストリーミング推論リクエストを実行するので、ハードウェアを効率的に使用できます。パイプライン上のコアの 1 つが最初のリクエストの処理を完了すると、最後のコアが最初のリクエストの処理を完了するのを待たずに、後続のリクエストの処理を開始できます。このストリーミングパイプラインの推論は、バッチサイズ 1 などのリアルタイムアプリケーションで小さなバッチサイズの推論を実行する場合でも、コアあたりのハードウェアの利用率が向上します。

単一の大きなモデルに適合する最適な Neuronコアの数を見つけることは、経験的なプロセスです。良い出発点は、次の近似式を使用することですが、最適解を得るには、複数の構成を試してみることをお勧めします。

neuronCore_pipeline_cores = 4*round(number-of-weights-in-model/(2E7))

コンパイラは neuroncore-pipeline-cores コンパイルフラグの値を直接受け取りますが、それだけです！この機能を有効にするには、目的のフレームワークの通常のコンパイルフローにこの引数を追加します。

TensorFlow Neuron では、以下のコードを使用します。

import numpy as np
import tensorflow.neuron as tfn

example_input = np.zeros([1,224,224,3], dtype='float16')
tfn.saved_model.compile("<Path to your saved model>",
                        "<Path to write compiled model>/1",
                        model_feed_dict={'input_1:0' : example_input },
                        compiler_args = ['--neuroncore-pipeline-cores', '8'])

PyTorch Neuron では、以下のコードを使用します。

import torch
import torch_neuron

model = torch.jit.load(<Path to your traced model>)
inputs = torch.zeros([1, 3, 224, 224], dtype=torch.float32)

model_compiled = torch.neuron.trace(model, 
                           example_inputs=inputs, 
                           compiler_args = ['--neuroncore-pipeline-cores', '8'])

Neuronコアパイプラインおよびその他の Neuron 機能の詳細については、Neuron Features を参照してください。

AWS Inferentia で HuggingFace 質問応答モデルを実行

AWS Inferentia上で Hugging Face BertForQuestionAnswering モデルを実行するには、torch_neuronフレームワークをインポートする以外に、通常の Transformer の実装に1行のコードを追加するだけです。以下のスニペットに従って、フォワードパスメソッドの例を適応することができます。

from transformers import BertTokenizer, BertForQuestionAnswering
import torch
import torch_neuron

tokenizer = BertTokenizer.from_pretrained('twmkn9/bert-base-uncased-squad2')
model = BertForQuestionAnswering.from_pretrained('twmkn9/bert-base-uncased-squad2',return_dict=False)

question, text = "Who was Jim Henson?", "Jim Henson was a nice puppet"
inputs = tokenizer(question, text, return_tensors='pt')

neuron_model = torch.neuron.trace(model, 
                                  example_inputs = (inputs['input_ids'],inputs['attention_mask']),
                                  verbose=1)

outputs = neuron_model(*(inputs['input_ids'],inputs['attention_mask']))

前述のコードで1 行増えたのは、torch.neuron.trace()メソッドの呼び出しです。この呼び出しによりモデルがコンパイルされ、新しいneuron_model()メソッドが返され、スクリプトの最後の行に示されているように、元の入力に対して推論を実行することができます。この例をテストしてみたい方は、PyTorch Hugging Face pretrained BERT Tutorial を参照してください。

事前学習済みのモデルや、ファインチューニングされたモデルを、Hugging Face のモデルリポジトリから直接コンパイルして推論を実行できるようになったことは、本番環境でのデプロイを最適化するための最初のステップです。この最初のステップにより、GPU を使用した場合と比較して、70% 低いコストで 2 倍のパフォーマンスが得られます (この記事の後半で説明します)。Neuronコアグループとパイプライン機能を組み合わせると、1つの Inf1 インスタンス内でモデルをパッケージ化する様々な方法を検討することができます。

Neuronコアグループとパイプラインによるモデルデプロイの最適化

HuggingFace の質問応答モデルのデプロイは、いくつかのモデルのパラメータを事前に設定する必要があります。Neuron は、コンパイル時にテンソル形状の知識を必要とする Ahead-Of-Time (AOT) コンパイラです。そのため、モデルデプロイ時のバッチサイズとシーケンス長の両方を定義します。前述の例では、Neuron フレームワークは、トレース呼び出しで渡されたサンプル入力からそれらを推論しました: (inputs[‘input_ids’], inputs[‘attention_mask’])。

これら 2 つのモデルパラメーターのほかに、コンパイラ引数--neuroncore-pipeline-coresと環境変数NEURONCORE_GROUP_SIZESを設定して、モデルサーバーが AWS Inferentia チップ上の Neuronコアをどのように消費するかを微調整することができます。

例えば、1つの AWS Inferentia チップで推論リクエストを処理するサーバーワーカーの同時実行数を4コアで最大化するためには、環境変数NEURONCORE_GROUP_SIZES=”1,1,1,1”を設定し、コンパイラ引数--neuroncore-pipeline-coresを1に設定するか省略します。次の図は、この設定を描いたもので、完全なデータパラレルでのデプロイです。

レイテンシーを最小限に抑えるために、--neuroncore-pipeline-coresを 4 に設定し、環境変数 NEURONCORE_GROUP_SIZES=”4”を設定して、プロセスが一つのモデルに対して同時に4つの Neuronコアすべてを消費するようにすることができます。AWS Inferentia チップは、4 つの推論リクエストをストリームとして同時に処理できます。このモデルパイプラインパラレルでのデプロイは、次の図のようになります。

データパラレルでは、複数のワーカーが同時にリクエストを処理することでスループットが向上するのに対し、パイプラインパラレルでは、レイテンシーが優先されますが、ストリーム処理動作により、スループットも向上します。これらの2つのパラメータを追加することで、ユースケースで最も重要なサービングメトリクスに応じて、サービングアプリケーションのアーキテクチャを微調整することができます。

最小レイテンシーのための最適化: マルチコア パイプラインパラレル

オンラインチャットボットのワークフローの一部としてシーケンス分類を行うような、最小のレイテンシーを必要とするアプリケーションを考えてみましょう。ユーザーがテキストを送信すると、バックエンドで実行されているモデルが1つのユーザー入力の意図を分類しますが、その推論速度には限界があります。このモデルは、ほとんどの場合、単一の入力（バッチサイズ1）のリクエストに対して応答する必要があります。以下の表は、HuggingFace BERT base モデルをデータパラレルとパイプラインパラレルの構成で、バッチサイズ1で実行した場合の、Inf1インスタンスと g4dn.xlarge（クラウドでの推論用に最適化されたGPUインスタンスファミリー）のパフォーマンスとコストを比較したものです。95 パーセンタイル (p95) のレイテンシーを見ると、4 コアの inf1.xlarge と 16 コアの inf1.6xlarge の両方のインスタンスで、パイプラインモードの方が低い値を示しています。Inf1 インスタンス間の最適な構成は 16 コアのケースで、レイテンシーが 58% 削減され、6 ミリ秒に達しています。

テストしたモデルは、HuggingFace bert-base-uncase の PyTorch バージョンで、シーケンス長は128です。AWS Inferentia上では、利用可能なすべてのコアを使用するようにモデルをコンパイルし、フルパイプラインのパラレル実行を行いました。データパラレルの場合は、1つのコア用にモデルをコンパイルし、コアごとにワーカーモデルを実行するように Neuronコアグループを設定しました。GPU を利用したテストでは、AWS Inferentia と同じ設定を使用し、モデルを TorchScript JIT でトレースし、PyTorch AMP Autocast で混合精度にキャストしました。

スループットも AWS Inferentia のパイプラインモードで1.84倍になり、毎秒1,793文に達し、g4dn.xlarge の12倍のスループットになりました。また、この構成での推論コストは、時間当たりのコストが高くなっても、最もコスト効率の良い GPU オプションよりも inf1.6xlarge が有利です。100 万文あたりのコストは、Amazon Elastic Compute Cloud（Amazon EC2）のオンデマンドインスタンスの価格設定では 70% 低くなります。inf1.6xlarge のスループットをフルに活用できないレイテンシーに敏感なアプリケーションや、BERT Small などの小型モデルでは、費用対効果の高い導入を実現するために、inf1.xlarge のパイプラインモードを使用することをお勧めします。

最大スループットのための最適化: シングルコア データパラレル

低レイテンシーよりもスループットの向上が求められる NLP のユースケースとして、検索や文書検索のパイプラインの一部である抽出型の質問応答タスクがあります。この場合、並行して処理されるドキュメントセクションの数を増やすことで検索を高速化し、検索された回答の質や幅を向上させることができます。このような設定では、推論はバッチで実行される可能性が高くなります（バッチサイズは1よりも大きい）。

最大のスループットを達成するために、実験を通して、AWS Inferentia での最適なバッチサイズは、先程テストした同じモデルで 6であることがわかりました。g4dn.xlarge では、バッチ 64を実行しても GPU メモリが不足することはありませんでした。以下の結果は、inf1.6xlarge において、バッチサイズ 6が GPU と比較して、61%の低コストで 9.2倍のスループットを提供できることを示しています。

このアプリケーションでは、コストを考慮することも、最終的なサービングインフラストラクチャの設計に影響を与えます。バッチ処理された推論を実行する最もコスト効率の高い方法は、inf1.xlarge インスタンスを使用することです。inf1.xlarge インスタンスは、GPU インスタンスと比較して、70%低いコストで 2.3倍高いスループットを実現しています。inf1.xlarge と inf1.6xlarge のどちらを選択するかは、コストを最小にするか、スループットを最大にするかという主な目的によって決まります。

Neuronコアのパイプライン及びグループ機能を自分でテストするには、PyTorch の最新の Utilizing Neuron Capabilities チュートリアルを参照してください。

まとめ

この記事では、Neuronコア のグループおよびパイプライン機能を使用して、NLP デプロイを最適化する方法を検討しました。AWS Neuron SDK と PyTorch のネイティブな統合により、最小限のコード変更で HuggingFace Transformers モデルをコンパイルし、AWS Inferentia 上で実行できるように最適化することができました。デプロイするアーキテクチャをパイプラインパラレルになるように調整することで、BERT モデルはリアルタイムアプリケーションのレイテンシーを最小限に抑え、g4dn.xlarge と比べて 12倍のスループットを実現しながら、実行コストは 70％削減されました。また、バッチ推論では、60％のコスト削減で 9.2倍のスループットを達成しています。

この記事で説明する Neuron SDK の機能は、他の ML モデルタイプおよびフレームワークにも適用されます。詳細については、AWS Neuron Documentation を参照してください。

AWS Inferentia チップとAmazon EC2 Inf1インスタンスの詳細をご覧いただき、Neuron SDK を使用して AWS Inferentia 上で独自のカスタムMLパイプラインを実行し始めましょう。

翻訳はアンナプルナラボの常世が担当しました。