AWS Inferentia を使った大規模推論環境の構築
Author : 安立 健人 (株式会社マネーフォワード)
こんにちは。株式会社マネーフォワードで主に機械学習関連の開発をしている安立と申します。
本稿では、グループ会社の HiTTO 株式会社との取り組みにおいて、機械学習推論処理に特化した AWS Inferentia チップを利用して大規模推論環境を構築したことについて、またその際に直面した問題や実現するために行った工夫についてお話したいと思います。
AWS Inferentia チップを弊社は今回初めて利用してみましたが、想像以上に速く利用法も簡単で驚きました。
弊社が直面していた課題
まず、AWS Inferentia を採用した理由の前に、本プロジェクトにおいて我々が直面していた課題についてご説明したいと思います。
もともと本プロジェクトは、HiTTO 株式会社が提供しているチャットボットサービスの学習・推論システムリプレースプロジェクトとしてスタートしました。(本稿では学習環境については触れませんが、こちらも Amazon SageMaker Training を利用したものにリプレースしました !)
このプロジェクトが始まった当初、我々は以下のような課題を抱いていました。
短期間でチャットボットの学習・推論環境を内製する必要があった
これは完全なビジネス上の都合ですが、プロジェクトスタート時点で我々に与えられていた猶予は 2 ヶ月ほどしかありませんでした。
人数も限られている状況で、この短期間にすべてのシステムをリプレースするということは非常に困難なタスクでした。
低レイテンシーの要求
HiTTO が提供するチャットボットサービスの特性によるものですが、推論処理におけるレスポンスのレイテンシ要求が非常に高いものとなっていました。
チャットボットサービスなので、実際に使用している機械学習モデルは自然言語処理のモデルとなっています。
深層学習を使った自然言語処理の機械学習モデルを CPU 上で高速に推論することは非常に困難です。これは深層学習を使ったモデル等が同時に多量の計算を求めるつくりとなっており、CPU のようなコア数の少ないものでは計算に時間がかかってしまうのが理由です。
そのため深層学習ではコア数が多く同時に多量の計算をすることが得意な GPU を使うことで高速な推論を実現することがよくありますが、実際に今回のケースでは GPU を採用しても達成することが難しいレベルの要求速度となっていました。
AWS Inferentia を採用した理由
以上のような課題があるなかで、AWS の方々と相談した結果、候補としてあがったのが AWS Inferentia チップでした。この AWS Inferentia チップは AWS 独自の機械学習推論処理に特化したチップとなっており、 Amazon Elastic Compute Cloud (EC2) Inf1 インスタンスと呼ばれるインスタンスファミリーで利用可能です。
今回、最終的に我々はこの AWS Inferentia を採用することで課題を解決することができました。
AWS Inferentia チップを採用した大きな理由は以下の 3 つでした。
推論実行速度
なにはともあれ、試験的に利用してみた結果、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 を利用する方法を説明します。
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" torchvisionimport 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 を使ってコンパイル済モデルの評価をできない点は注意が必要です。
AWS Inferentia を使った推論構成
さて、ではこのコンパイルできたモデルを弊社でどのように利用しているかご紹介します。
本プロジェクトでは、利用する顧客ごとにモデルを作成・推論環境を提供する必要があるので、モデルは顧客数分存在し、膨大な数になります。実際には ECS on EC2 で、大量のサービスを並べるという手法を選択しています。
AWS Fargate が利用できないので若干のメンテナンスコストはあるものの、開発チームが SageMaker よりも ECS に関する知識・スキルがあった点などを理由にこの手法を選択しました。
大規模推論ならではの工夫したこと
このように多数のモデルを多数の ECS on EC2 で管理して推論サービスを提供するという形をとっていますが、いくつか問題点や詰まった点もありました。
キャパシティー予約
このサービスで利用しているのはこれまでに述べてきたとおり Inf1 インスタンスファミリーのインスタンスですが、このサービスでは、ある程度の台数 (数十台以上) を着実に用意しておく必要がありました。そのため本プロジェクトではまず最初に Inf1 インスタンスのオンデマンドキャパシティ予約を行いました。
クリックすると拡大します
キャパシティ予約自体は EC2 のメニューから実行できます。
これを行うことで、サービス提供に必要な台数を予め確保するようにしました。
また今回我々は、希望する AWS Inferentia チップ枚数分のインスタンスを確実に予約するために、複数のサイズのキャパシティープールを混在させる戦略としました。具体的には、 Inf1.xlarge 、 Inf1.2xlarge を混在させて予約するといった方法で必要な台数を確保しました。
また、予約可能な台数というのは日々変わり続け (最初に希望の台数に届かなくても、1 週間後には予約できたりする)、オンデマンドキャパシティー予約は任意に変更できるので、コスト削減のために小まめに予約数量を更新してみて可能な限り Inf1.xlarge に変更していきました。
複数インスタンスタイプの AWS Auto Scaling
前述の通り、我々は必要なインスタンスを確保するために 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 をキックするだけで、デプロイが始まるようになっています。
まとめ
株式会社マネーフォワード・HiTTO株式会社において AWS Inferentia チップを活用して大規模推論環境を構築した話をさせていただきました。
いくつかの注意点はありましたが、このように AWS Inferentia を採用することで 短期間に高速な推論サービスを立ち上げることが可能 です。
正直なところ我々も最初は「Inferentia は推論専用チップという触れ込みだが、果たしてそんなに速くなるのか ?」と思っていましたが、今回使ってみてその性能にとても驚きました。
AWS では他にも独自チップの開発が進んでおり、直近では学習専用チップ AWS Traininum というチップを搭載した Trn1 インスタンスがプレビューアナウンスされており、こちらの性能もとても気になるところです !
ここまで読んでくださった方ありがとうございました。
本記事が機械学習サービスにおけるインフラ選定の一助となれば幸いです。
筆者プロフィール
安立 健人 (@ken11)
株式会社マネーフォワード
CTO室 AI推進部
旅するエンジニア。マネーフォワードにて ML 関連システムの開発を担当しています。
趣味で自然言語処理系の事前学習済みモデルの公開等もしています。
最近のマイブームは一晩中 YouTube を見続けることです。
AWS を無料でお試しいただけます