AWS Inferentia を使った大規模推論環境の構築

これは完全なビジネス上の都合ですが、プロジェクトスタート時点で我々に与えられていた猶予は 2 ヶ月ほどしかありませんでした。
人数も限られている状況で、この短期間にすべてのシステムをリプレースするということは非常に困難なタスクでした。

HiTTO が提供するチャットボットサービスの特性によるものですが、推論処理におけるレスポンスのレイテンシ要求が非常に高いものとなっていました。
チャットボットサービスなので、実際に使用している機械学習モデルは自然言語処理のモデルとなっています。

深層学習を使った自然言語処理の機械学習モデルを CPU 上で高速に推論することは非常に困難です。これは深層学習を使ったモデル等が同時に多量の計算を求めるつくりとなっており、CPU のようなコア数の少ないものでは計算に時間がかかってしまうのが理由です。

そのため深層学習ではコア数が多く同時に多量の計算をすることが得意な GPU を使うことで高速な推論を実現することがよくありますが、実際に今回のケースでは GPU を採用しても達成することが難しいレベルの要求速度となっていました。

なにはともあれ、試験的に利用してみた結果、AWS Inferentia チップによる推論処理は非常に高速でした。先述の通りなるべくレイテンシーを小さくする必要があった我々にとって、AWS Inferentia チップの処理速度の速さは必要不可欠なものでした。

当初は AWS Inferentia を搭載した Inf1 インスタンスの他にも東京リージョンで利用可能な NVIDIA GPU を搭載したG4dn インスタンスも検討しており、 G4dn インスタンスとの比較検証も行いました。残念ながら、検討時点ではまだ東京リージョンには G5 インスタンス は来ておりませんでした。
その結果、G4dn.xlarge と比較してなんと推論処理にかかる時間は 97% も減ったのです !

AWS Inferentia による推論速度の向上はもちろん採用しているモデルによって大きく変わると思いますが、今回我々が採用していた NLP のモデルに対してはかなり AWS Inferentia の強みが活かせました。

そのうえコストも G4dn の 50% 以下となっていて、今回のタスクでは AWS Inferentia チップを使うしかない ! と感じました。

AWS Inferentia チップには 4 つの Neuron コアがあり、それを有効活用することで 1 台 で複数モデルをサービングすることが非常に容易となっています。メモリに乗り切る分であれば複数モデルをそのまま乗せることが可能です。

今回の事例ではむしろリクエスト数が集中することで捌ききれなくなることを理由に、乗せるモデル数を制限したくらいですが、それでも複数モデルを乗せることができたおかげで全体としてインスタンス台数の削減 = 大幅なコスト減につながっています。

また、AWS Inferentia を利用することが非常に簡単だったことも採用の理由となっています。

先述の通り、我々にはもう時間があまりないという状況だったので、どんなに速くて安かったとしてもすぐに実装できる自信がなければ採用することは難しいと考えていました。しかしながら、AWS Inferentia チップを利用する方法は既存のモデルをコンパイルするだけと非常に簡単で、事前の用意等もほとんど必要なかったので、最終的に推論処理で AWS Inferentia チップを利用することにしました。

ここからは、これまでの体験にもとづいて、AWS Inferentia を利用する方法を説明します。

今回我々がもともと用意していたモデルは PyTorch のモデルで、これを専用の AWS Neuron SDK を利用してコンパイルする必要があります。このコンパイル作業自体は Inf1 インスタンス以外でも利用可能なので、たとえば P4d インスタンスで学習してできたモデルをコンパイルして Inf1 インスタンスでサービングするといったことは可能です。
Inf1 インスタンス上でしかコンパイルできないといった制限があると困ってしまいますが、ここは一安心ですね。

コンパイル作業では先述の AWS Neuron SDK の neuron-cc というものを利用します。neuron-cc はモデルコンパイル用のコマンドラインツールです。PyTorchの場合 torch-neuron というものもあるので、これを使います。torchscript や ONNX をやったことがある方はおなじみの感覚で、サクッとコンパイルできます。

pip config set global.extra-index-url https://pip.repos.neuron.amazonaws.com
pip install torch-neuron neuron-cc[tensorflow] "protobuf<4" torchvision

import torch.neuron

# trained_model: 学習済みのPyTorchモデルです
# inputs: モデルの入力サンプルを用意します
model_neuron = torch.neuron.trace(trained_model, inputs, strict=False)
model_neuron.save("neuron_model.pt")

コンパイル自体はこれだけです !
torchscript のコンパイルが

traced_model = torch.jit.trace(trained_model, inputs)
torch.jit.save(traced_model, "model.pt")

なので、ほとんど変わらない手順で簡単にコンパイルできることがわかると思います。

注意点

コンパイルしたモデルを実際に推論で使うには、AWS Inferentia チップ搭載の Inf1 インスタンスファミリーのインスタンスを使う必要があります。(他のインスタンスタイプでは動作しません) この Inf1 を使うには Amazon SageMaker の推論エンドポイント機能か、EC2 の 2 つの方法があります。どちらを利用するかはプロジェクトの状況等によると思いますが、弊社では今回 EC2 を使って Amazon ECS on EC2 で推論サーバーを実行しています。また、当然ながらコンパイルによって精度には変化が出る可能性があるのでそこは気をつける必要があります。

コンパイル時に --fast-math というオプションでどの程度パフォーマンスと精度をトレードオフするか選択することができます。(デフォルトではパフォーマンス最優先)

それもふまえて、コンパイル後のモデルを評価したいケースもあると思いますが、先述の通り推論の実行は Inf1 上でしかできないため、コンパイル済モデルの評価プロセスも Inf1 上で行うことになります。

Inf1 インスタンスは現状は SageMaker の推論エンドポイントか EC2 しかないので、SageMaker Processing を使ってコンパイル済モデルの評価をできない点は注意が必要です。

このサービスで利用しているのはこれまでに述べてきたとおり Inf1 インスタンスファミリーのインスタンスですが、このサービスでは、ある程度の台数 (数十台以上) を着実に用意しておく必要がありました。そのため本プロジェクトではまず最初に Inf1 インスタンスのオンデマンドキャパシティ予約を行いました。

キャパシティ予約自体は EC2 のメニューから実行できます。
これを行うことで、サービス提供に必要な台数を予め確保するようにしました。

また今回我々は、希望する AWS Inferentia チップ枚数分のインスタンスを確実に予約するために、複数のサイズのキャパシティープールを混在させる戦略としました。具体的には、 Inf1.xlarge 、 Inf1.2xlarge を混在させて予約するといった方法で必要な台数を確保しました。

また、予約可能な台数というのは日々変わり続け (最初に希望の台数に届かなくても、1 週間後には予約できたりする)、オンデマンドキャパシティー予約は任意に変更できるので、コスト削減のために小まめに予約数量を更新してみて可能な限り Inf1.xlarge に変更していきました。

前述の通り、我々は必要なインスタンスを確保するために Inf1.xlarge と Inf1.2xlarge の 2 つのインスタンスタイプを混合で、Amazon EC2 Auto Scaling グループ を用意する必要がありました。

これは Auto Scaling グループ の インスタンスタイプの要件 という項目で設定することが可能です。

また、多数の ECS on EC2 のサービスを並べていると、デプロイ作業も難しくなります。

私は CI/CD においてよくサードパーティー製ツールである ecspresso を利用します。CircleCI で ecspresso を使ってデプロイをキックする形です。

しかし今回のように大量の ECS on EC2 のサービスを CircleCI からキックすると、デプロイに長時間を要します。ECS on EC2 の場合まずインスタンスの起動から必要ということもありデプロイにはとにかく時間がかかるので、その間 CircleCI をロックするのは避けたいと思っていました。

対策として、今回この ECS サービス群をデプロイ管理するための AWS Step Functions を作成しました。

このような Step Functions を用意することで、大量のサービスに対して同時に ECS の UpdateService 操作を要求することが可能となります。

また、UpdateService をキックして終わりではなく、その後 DescribeServices を使ってデプロイが完了するまで待機するようにしています。

CI からはこの Step Functions をキックするだけで、デプロイが始まるようになっています。