[レポート] 機械学習のインスタンスどれ使えばいいの?に答える Chalk Talk セッションに参加してきました #reinvent #AIM407

本エントリはAWS re:Invent 2022のセッション「Choosing the right ML instance for training and inference on AWS (AIM407)」のレポートです。
2022.12.13

この記事は公開されてから1年以上経過しています。情報が古い可能性がありますので、ご注意ください。

はじめに

こんちには。

データアナリティクス事業本部 機械学習チームの中村です。

re:Invent 2022に現地で参加し、機械学習系のセッションをメインに回っていました。

本記事では「Choosing the right ML instance for training and inference on AWS」というセッションに参加しましたので、そのレポートをします。

セッションについて

  • タイトル
    • Choosing the right ML instance for training and inference on AWS
  • 登壇者
    • Raghu Ramesha, ML Solutions Architect, AWS
    • Samir Araujo, AI/ML Solutions Architect, AWS
  • セッション情報
    • 日時: 2022-11-30 (Wed) 11:30-12:30
    • 形式: Chalk Talk
    • 番号: AIM407
    • 会場: MGM Grand (Level 1, Room 101, MGM Grand)
    • レベル: 400 - Expert

なお、本セッションはChalk Talkであるため、セッションの動画は公開されていないようです。

セッション概要

事前の案内としては以下の通りです。

As a data scientist, choosing the right compute instance for your workload on AWS can be challenging. On AWS, you can choose from CPUs, GPUs, AWS Trainium, and Intel Habana Gaudi to accelerate training. For inference, you can choose from CPUs, GPUs, and AWS Inferentia. This chalk talk guides you through how to choose the right compute instance type on AWS for your deep learning projects. Explore the available options, such as the most performant instance for training, the best instance for prototyping, and the most cost-effective instance for inference deployments. Learn how to use Amazon SageMaker Inference Recommender to help you make the right decision for your inference workloads.

(日本語訳)

データサイエンティストとして、AWS上のワークロードに適したコンピュートインスタンスを選択することは困難です。AWSでは、CPU、GPU、AWS Trainium、Intel Habana Gaudiから選択し、学習を加速させることができます。推論については、CPU、GPU、AWS Inferentiaから選択することができます。このチョークトークでは、ディープラーニングプロジェクトに適したAWSのコンピュートインスタンスタイプを選択する方法をご案内します。トレーニングに最もパフォーマンスの高いインスタンス、プロトタイピングに最適なインスタンス、推論展開に最も費用対効果の高いインスタンスなど、利用可能な選択肢を探ります。Amazon SageMaker Inference Recommenderを使用して、推論ワークロードの正しい決定を支援する方法を学びます。

セッション聴講内容

機械学習における学習・推論においてどのようなインスタンスを選択するべきかを議論するセッションとなっています。

アジェンダ

  • Common scenarios for training and inference
  • Challenges
  • Picking the right training instances
  • Picking the right inference instances
  • Discussion and QA

アーキテクチャの一例

はじめに題材として、以下のようなアーキテクチャが提示されました。

JPEGの複数フレームを1つの画像にする、Mosaicなどの処理を入れているように見えます。 データ拡張としてMosaicを入れているのか、処理効率化のためにMosaicを入れているのか、ちょっと判別はできませんでした。 (後述のベンチマークを見る限り、推論時の話をしている気がします)

また、検出したいすべてのパーツのCADモデルやファイルを持っているため、Blenderでデータ拡張をしたりする処理も入れているようです。おそらく以下のような話と思われます。

推論には、Inferentiaを使用しています。Inferentiaは1チップに4つのコアがあり、モデルをスライスしてそれぞれの コアに割り当てるよう動作するようです。より低いレイテンシで予測値を返します。

以下がベンチマークの結果として示されました。

さらにここから改善するには以下のような選択肢があります。

まずはスケールアウト、スケールアップをすること。 これは、Inferentiaには、Neuron SDKと呼ばれるSDKがありますので、自動的に行うことができます。 コンパイル時に、Neuron core pipelinesというパラメータと、利用したいコア数を指定する必要があり、 そこで指定したコア数でモデルをスライスしてシャードしてくれるようになるようです。

他には、TurboJPEGを使用してデコードを高速化することなどが挙げられていました。

コードも提示されました。Mosaicにしたデータは、バッチ作成時に1つずつの画像に復元されているようですね。

ML training common scenarios

トレーニングについての3つのシナリオを挙げています。

  • Small : 古典的な機械学習を使用する場合
  • Intermediate/big : 画像解析レベルまで(ResNet101など)
  • Huge : NLPなどの大規模言語モデル

選択肢としては4つ示されました。

  • CPU instances
    • c5, c6, c6g, c7gなど
    • Gravitonは安いものの、x86_64ではなく、ARM CPUをターゲットにする必要がある。
  • GPU instances
    • p4, それほど大きなパワーが必要でない場合はp3
    • より小さい選択肢はg4やg5など
  • AWS Trainium
    • トレーニングプロセスを高速化し、コストを削減し、パフォーマンスを向上させた
    • GPUが汎用的な行列演算が目的であるのと違い、機械学習のためだけに作成された
    • GPUよりも性能が良く、消費電力も少ない。
    • そのため、同じサーバーの中で、GPUよりもトレーニングの方がアクセラレーターを多く搭載でき、消費電力は同じです。
    • TensorFlowやPyTorchと互換性があり、XLAによりそれが可能となっている。
    • トレーニング後の最終的なモデルは、通常のPyTorchやTensorFlowのモデルと同じで、他のinstanceにデプロイして使用可能
  • Habana Gaudi-based instances
    • dl1

上記のうち、CPU instances、GPU instances、AWS TrainiumはSageMakerから利用可能となっています。

まずは、CPU instancesにフォーカスして考えるべき点が紹介されていました。

次にGPU instancesについて、ハイエンドからT4などのColabでもよく見るGPUが搭載されるG4インスタンスまで。

p4はマルチGPUなので、シングルGPUで良い場合は、p3、g5, g4が選択肢となるようです。

Trainiumについては時間を割いて多く説明がありました。

Trainiumは16このチップ、32コアのインスタンスを起動できます。 GPUメモリがp4よりもさらに大規模である点は特徴としてありそうです。 また特に小さいバッチサイズでは、Trainiumにパフォーマンスの優位性があるようですね。 更に大きなものが必要な場合は32コア、512コアを搭載したTrainiumインスタンスも使用可能となっています。

チップはすべて、Neural Linkと呼ばれる低密度・高スループットのリンクで接続されます。 各コアは勾配情報を共有するため、コア間も低遅延のリンクが必要になります。

分散学習についての説明がありました。

  • Data Paralelism
    • モデルがアクセラレータのメモリに収まる場合はこの手法が使える
    • 全てのアクセラレータにモデルを複製し、それぞれ異なるバッチを処理してその後平均化される
  • Pipeline Parelelism
    • アクセラレータに収まらない場合に、モデルをレイヤ毎にスライスする
  • Tensor Parelelism
    • アクセラレータに収まらない場合に、モデルをTensor毎にスライスする

これらの並列化に対応するためには、全ての演算(オペレータ)を対応させる必要があるが、 通常は順次増やしていくような形となっているそうです。新しいオペレータを追加するニーズがあれば、要望をすると 早く実装される可能性もあるようです。

SageMakerについては知っている人も多いだろうということでスキップされました。

次にデータセットをどのように構築するかの選択肢として、S3、EFS、FSx for Lustreを比較されていました。

最適化の戦略について3軸が挙げられていました。

一つは、適切なインフラ(インスタンスを選択すること)、もう一つはCustom Stackを活用すること、最後にモデルを圧縮やランタイムの選択です。

圧縮には、pruning(枝刈り)、quantization(量子化)などがあり、ランタイムには、TVM、TreeLite、TensorRT、AWS Neuron、Amazon SageMaker Neoなど様々な選択肢があるようです。

デモとして、TensorRTを使用した場合の推論速度の向上が紹介されました。かなり高速になっていることが分かります。

Instances and accelerators for ML Inference

次に推論時の選択肢について紹介されました。

AWS Inferentiaも4つのNeuron Core間でチップ間通信をしているとのこと。混合精度にも対応しています。

AWS Neuron SDKも提供されており、モデルのベンチマークや、 推論スコアを最大化するためのツールセットも提供されています。

3つの選択肢の総合的なまとめです。

Inferentiaは、モデルに互換性があればコスト効率が良い選択肢になるとのこと。 プレビューですが、もちろんinf2インスタンスも使用することが可能です。以下の弊社のブログでも紹介されています。

最後にSageMakerを使った推論オプションについて説明がありました。

AWS SageMakerの推論はゼロへのスケールダウンを支援します。 これはマネージドキューの提供で、トラフィックがないときだけゼロにスケールダウンしてくれるようです。

また、インスタンスを選択したくない場合は、SageMaker Serverless Inferenceという機能もあります。 必要なモデルやメモリを用意するだけで、同等のインスタンスをスピンアップして管理します。

その他、SageMaker Asynchronous InferenceとSageMaker Batch Transformが紹介されていました。

また、Inference Recommenderは、これはスループット要件に基づいて適切なインスタンスを選択するのに役立つもので、 これにより、最小限のコストで必要なパフォーマンスが得られる環境を見つけることが可能です。

まとめ

いかがでしたでしょうか。 Chalk Talkということで、機械学習専用チップの紹介や使いどころなどについても聞けて良かったと思います。 実業務でも機械学習でどのインスタンスを選択するかは、コスト・パフォーマンスの観点からも重要な要素なので、 これからも着目していきたいと思います!