Amazon Web Services ブログ

Amazon SageMaker 上で AWS Inferentia2 と AWS Trainium を使って、低コストで高性能な生成系 AI 推論を実現

人工知能 (AI) と機械学習 (ML) の世界では、人間のようなテキスト、画像、コード、音声を作成できる生成系 AI モデルの台頭により、パラダイムシフトが起きています。従来の ML モデルと比較して、生成系 AI モデルは大規模で、より複雑です。また、複雑さが増すにつれて、推論のコストが高くなり、強力なコンピューティングリソースに対するニーズも高まっています。生成系 AI モデルの推論コストが高いため、リソースが限られている企業や研究者にとって参入の障壁となる可能性があり、より効率的で費用対効果の高いソリューションが必要になります。さらに、生成系 AI のユースケースの大半は、人間との対話や現実世界のシナリオに関係しているため、低レイテンシー性能を実現できるハードウェアが必要です。AWS は、強力で効率的で費用対効果の高いコンピューティングハードウェアに対するニーズの高まりに応えるために、専用のチップで革新を進めてきました。

本日 (訳注: 原著公開日は2023年5月4日)、Amazon SageMaker が AWS Inferentia2 (ml.inf2) と AWS Trainium (ml.trn1) ベース (訳注: Trainium はモデル学習特化のチップですが、後述するように推論用途で利用することも可能です) の SageMaker インスタンスをサポートして、リアルタイムおよび非同期推論のための生成系 AI モデルをホストすることを発表できることを嬉しく思います。米国東部 (オハイオ) において SageMaker の ml.inf2 インスタンスを用いたモデルのデプロイを利用可能です。また、ml.trn1 インスタンスは米国東部 (バージニア北部) にて利用することができます。

これらのインスタンスを SageMaker で使用すると、大規模言語モデル (LLM)、Stable Diffusion、vision transformers などの生成系 AI モデルを高性能かつ低コストで実現できます。さらに、Amazon SageMaker Inference Recommender を使用すると、負荷テストを実行したり、これらのインスタンスにモデルをデプロイすることによるコストパフォーマンス上のメリットを評価したりすることができます。

ml.inf2 インスタンスと ml.trn1 インスタンスを使用して SageMaker で ML アプリケーションを実行し、テキストの要約、コード生成、ビデオと画像の生成、音声認識、パーソナライズ、不正検出などを行うことができます。SageMaker エンドポイントを設定するときに ml.trn1 または ml.inf2 インスタンスを指定するだけで簡単に開始できます。PyTorch、TensorFlow、Hugging Face、および大規模モデル推論 (LMI) 用の ml.trn1 および ml.inf2 互換の AWS Deep Learning Containers (DLC) を使用すると、簡単に始めることができます。すべてのバージョン一覧については、「Available Deep Learning Containers Images」を参照してください。

この記事では、LMI コンテナを活用して、追加のコーディングを必要とせずに SageMaker を使用して AWS Inferentia2 に大規模な言語モデルをデプロイするプロセスを示します。GPT4ALL-J というファインチューニングした GPT-J 7B モデルを使用しており、チャットボットスタイルの対話が可能です。

ml.trn1 インスタンスと ml.inf2 インスタンスの概要

ml.trn1 インスタンスは、主に LLM を含む生成系 AI モデルの高性能深層学習トレーニングを目的として構築された AWS Trainium アクセラレータを搭載しています。ただし、これらのインスタンスは、Inf2 よりもさらに大きなモデルを用いた推論ワークロードもサポートします。最大のインスタンスサイズである trn1.32xlarge インスタンスでは、1 つのインスタンスに 512 GB のアクセラレータメモリを備えた 16 個の Trainium アクセラレータが搭載されており、最大 3.4 ペタフロップスの FP16/BF16 の演算性能を実現します。16 台の Trainium アクセラレータ間を超高速 NeuronLinkV2 で接続することで、効率的な集団通信 (collective communications) を実現しています。

ml.inf2 インスタンスは、推論用のアクセラレータである AWS Inferentia2 アクセラレータを搭載しています。第 1 世代の AWS Inferentia と比較して、コンピューティング性能が 3 倍、スループットが最大 4 倍、10 分の 1 以下の低レイテンシーを実現します。最大のインスタンスサイズである inf2.48xlarge では、1 つのインスタンスに 384 GB のアクセラレータメモリを備えた 12 個の AWS Inferentia2 アクセラレータが搭載されており、BF16/FP16 の合計演算性能は 2.3 ペタフロップスです。これにより、1 つのインスタンスに最大 1750 億パラメータ規模のモデルをデプロイできます。Inf2 は、これまで高価なトレーニング向けインスタンスでのみ利用できた相互接続を提供する唯一の推論最適化インスタンスです。単一のアクセラレータに収まらない超大規模モデルの場合、データは NeuronLink を使用してアクセラレータ間を直接流れ、CPU を完全にバイパスします。NeuronLink を活用することで、Inf2 はより高速な分散推論をサポートし、スループットとレイテンシーを向上させます。

AWS Inferentia2 と Trainium のアクセラレータはどちらも、2 つの NeuronCores-v2、32 GB HBM メモリスタック、および専用の集合計算エンジンを搭載しており、マルチアクセラレータ推論を行う際に計算と通信を重複させることでランタイムを自動的に最適化します。アーキテクチャの詳細については、「Trainium and Inferentia devices」を参照してください。

次の図は、AWS Inferentia2 を使用したアーキテクチャの例を示しています。

AWS Neuron SDK

AWS Neuron は、AWS Inferentia および Trainium ベースのインスタンスで深層学習ワークロードを実行するために使用されるSDK です。AWS Neuron には、TensorFlow と PyTorch にネイティブに統合された深層学習コンパイラ、ランタイム、ツールが含まれています。Neuron を使用すると、高性能 ML ワークロードを ml.trn1 と ml.inf2 で開発、プロファイリング、デプロイできます。

Neuron コンパイラは、さまざまな形式 (TensorFlow、PyTorch、XLA HLO) の ML モデルを Neuron デバイスで実行するように最適化します。Neuron コンパイラは ML フレームワーク内で呼び出され、ML モデルは Neuron フレームワークプラグインによってコンパイラに送信されます。生成されるコンパイラアーティファクトは NEFF ファイル (Neuron 実行ファイル形式) と呼ばれ、次に Neuron ランタイムによって Neuron デバイスにロードされます。

Neuron ランタイムは、AWS Inferentia および Trainium Neuron デバイスにアクセスするための API を提供するカーネルドライバーと C/C++ ライブラリで構成されています。TensorFlow と PyTorch 用の Neuron ML フレームワークプラグインは、Neuron ランタイムを使用して Neuron コアにモデルをロードして実行します。Neuron ランタイムは、コンパイルされた深層学習モデル (NEFF) をNeuron デバイスにロードし、高スループットと低レイテンシーに最適化されています。

SageMaker ml.inf2 インスタンスを使用して NLP モデルをホストする

大きなモデルの重み行列を複数の Neuron コアにシャードするオープンソースのライブラリである transformers-neuronx を利用して LLM をサービングする方法について詳しく説明する前に、単一の Neuron コアに収まるモデルの一般的なデプロイフローを見ていきましょう。

サポートされているモデルのリストをチェックして、そのモデルが AWS Inferentia2 でサポートされていることを確認してください。次に、モデルを Neuron コンパイラでプリコンパイルする必要があります。SageMaker ノートブックまたは Amazon Elastic Compute Cloud (Amazon EC2) インスタンスを使用してモデルをコンパイルできます。次のコードに示すように、SageMaker Python SDK を使用して、PyTorch などの一般的な深層学習フレームワークを使用してモデルをデプロイできます。モデルを SageMaker ホスティングサービスにデプロイし、推論に使用できるエンドポイントを取得できます。これらのエンドポイントはフルマネージド型で、自動スケーリングをサポートしています。

from sagemaker.pytorch.model import PyTorchModel

pytorch_model = PyTorchModel(
    model_data=s3_model_uri,
    role=role,
    source_dir="code",
    entry_point="inference.py",
    image_uri=ecr_image
)

predictor = pytorch_model.deploy(
    initial_instance_count=1, 
    instance_type="ml.inf2.xlarge"
)

SageMaker 上の Inf2 の一般的な開発フローとサンプルスクリプトの詳細については、「Developer Flows」を参照してください。

SageMaker ml.inf2 インスタンスを使用して LLM をホストする

何十億ものパラメータを持つ大規模な言語モデルは、大きすぎて単一のアクセラレータに収まらないことがよくあります。そのため、複数のアクセラレータにわたって LLM をホストするには、モデル並列技術を使用する必要があります。LLM をホスティングするためのもう 1 つの重要な要件は、高性能なモデルサービングソリューションの実装です。このソリューションでは、モデルを効率的にロードし、パーティショニングを管理し、HTTP エンドポイント経由でシームレスにリクエストを処理する必要があります。

SageMaker には、モデル並列処理と大規模モデル推論に特化した AWS Deep Learning Containers (DLC)、ライブラリ、ツールが含まれています。SageMaker で LMI を使い始めるためのリソースについては、「Model parallelism and large model inference」を参照してください。SageMaker は、GPT、T5、OPT、BLOOM、Stable Diffusion などの大規模モデルを AWS インフラストラクチャでホストするために、広く使われている複数のオープンソースライブラリと共に DLC を管理しています。これらの特殊な DLC は SageMaker LMI コンテナと呼ばれます。

SageMaker LMI コンテナは DJLServing を使用しています。これは transformers-neuronx ライブラリと統合されたモデルサーバーで、Neuron コア間のテンソル並列処理をサポートしています。DJLServing の仕組みの詳細については、「Deploy large models on Amazon SageMaker using DJLServing and DeepSpeed model parallel inference」を参照してください。DJL モデルサーバーと transformers-neuronx ライブラリはコンテナのコアコンポーネントとして機能し、Neuron SDK も含まれています。この設定により、AWS Inferentia2 アクセラレータへのモデルのロードが容易になり、複数の Neuron コアにわたってモデルを並列化し、HTTP エンドポイント経由でサービスを提供できるようになります。

LMI コンテナは、Amazon Simple Storage Service (Amazon S3) バケットまたは Hugging Face Hub からのモデルのロードをサポートしています。デフォルトのハンドラースクリプトは、モデルを読み込み、コンパイルして Neuron に最適化された形式に変換し、ロードします。LMI コンテナを使用して LLM をホストするには、次の 2 つのオプションがあります。

  • ノーコード (推奨) — LMI コンテナを使用して LLM をデプロイする最も簡単な方法です。この方法では、提供されているデフォルトのハンドラーを使用して、モデル名と serving.properties ファイルに必要なパラメータを渡すだけで、モデルをロードしてホストできます。デフォルトのハンドラーを使用するには、entryPoint パラメータに djl_python.transformers-neuronx という名前を付けます。
  • 独自のスクリプトを使用 — この方法では、モデルのロードと提供に必要なコードを含む独自の model.py ファイルを作成することができます。このファイルは、DJLServing API と transformers-neuronx APIの間の仲介役として機能します。モデルの読み込みプロセスをカスタマイズするには、設定可能なパラメータを serving.properties に指定できます。設定可能なパラメータの包括的なリストについては、「All DJL configuration options」を参照してください。こちらに model.py ファイルの例を示します。

ランタイムアーキテクチャ

tensor_parallel_degree プロパティの値は、複数の Neuron コアにわたるテンソル並列モジュールの配分を決定します。たとえば、inf2.24xlarge には 6 つの AWS Inferentia2 アクセラレータが搭載されています。各 AWS Inferentia2 アクセラレータには 2 つの Neuron コアがあります。各 Neuron コアには、テンソル並列モジュールを格納する 16 GB の専用高帯域幅メモリ (HBM) があります。テンソルの並列度が 4 の場合、LMI は同じモデルの 3 つのモデルコピーを割り当て、それぞれが 4 つの Neuron コアを使用します。次の図に示すように、LMI コンテナが起動すると、モデルが最初に CPU のアドレス指定可能なメモリに読み込まれ、トレースされます。トレースが完了すると、モデルはテンソルの並列度に基づいてNeuron コア全体に分割されます。

LMI はモデルサービングスタックとして DJLServing を使用しています。コンテナのヘルスチェックが SageMaker でパスすると、コンテナは推論リクエストを処理する準備が整います。DJLServing は、TOTAL NUMBER OF NEURON CORES/TENSOR_PARALLEL_DEGREE に相当する複数の Python プロセスを起動します。各 Python プロセスには、TENSOR_PARALLEL_DEGREE と同等の C++ スレッドが含まれています。各 C++ スレッドは、1 つの Neuron コアにモデルのシャードを 1 つ保持します。

しばしば、推論サーバーが複数の独立したリクエストで呼び出されたときに、それぞれを独立に順次処理する実装が取られます。この実装はセットアップが簡単であるものの、アクセラレータの使用効率を考えるとベストな選択肢では有りません。これに対応するため、DJLServing には動的バッチ処理の最適化機能が組み込まれています。サーバー側でこれらの独立した推論リクエストを組み合わせて、より大きなバッチを動的に作成してスループットを向上させます。すべてのリクエストは、推論を待つために実際のジョブキューに入る前に、まず動的バッチャーに到達します。serving.properties の batch_size 設定を使用して、動的バッチ処理に適したバッチサイズを設定できます。また、レイテンシー要件に基づいて、他のリクエストがバッチに参加するのを待つまでのバッチャーの最大遅延時間を指定するように max_batch_delay を設定することもできます。スループットは、モデルコピーの数とコンテナで起動された Python プロセスグループにも依存します。次の図に示すように、テンソルの並列度を 4 に設定すると、LMI コンテナは 3 つの Python プロセスグループを起動し、それぞれがモデルのフルコピーを保持します。これにより、バッチサイズを大きくしてスループットを向上させることができます。

LLM をデプロイするための SageMaker ノートブック

このセクションでは、FP32 で 24GB になる 60億パラメータのモデルである GPT4all-J のデプロイの手順を段階的に説明します。GPT4all-J は、単語問題、対話、コード、詩、歌、物語など、膨大な種類の対話内容でトレーニングされた人気のチャットボットです。GPT4all-J は、GPT-J モデルをファインチューニングしたモデルで、人間の対話に近い応答を生成します。

この例の完全なノートブックは GitHub で提供されています。SageMaker Python SDK を使用してモデルを Inf2 インスタンスにデプロイすることができます。提供されているデフォルトのハンドラーを使用してモデルをロードします。これで、servings.properties ファイルを提供するだけで済みます。このファイルには、DJL モデルサーバーがモデルをダウンロードしてホストするために必要な構成が含まれています。model_id パラメータを使用して Hugging Face モデルの名前を指定すると、Hugging Face リポジトリから直接モデルをダウンロードできます。または、s3url パラメータを指定して Amazon S3 からモデルをダウンロードすることもできます。entryPoint パラメータは、モデルをロードするライブラリを指すように構成されています。djl_python.fastertransformer の詳細については、GitHub コードを参照してください。

tensor_parallel_degree プロパティの値は、複数のデバイスにわたるテンソル並列モジュールの分散方法を決定します。たとえば、Neuron コアが 12 個、テンソル並列度数が 4 の場合、LMI はモデルコピーを 3 つ割り当て、それぞれ 4 つの Neuron コアを使用します。dtype プロパティを使用して精度タイプを定義することもできます。n_position パラメータは、モデルの最大入力シーケンス長と出力シーケンス長の合計を定義します。次のコードを参照してください。

%%writefile serving.properties# Start writing content here
engine=Python
option.entryPoint=djl_python.transformers-neuronx
#option.model_id=nomic-ai/gpt4all-j
option.s3url = {{s3url}}
option.tensor_parallel_degree=2
option.model_loading_timeout=2400
option.n_positions=512

serving.properties を含む tarball を作成し、S3 バケットにアップロードします。この例ではデフォルトのハンドラーを使用していますが、ロードとサービングのプロセスをカスタマイズするための model.py ファイルを開発できます。インストールが必要なパッケージがある場合は、それらを requirements.txt ファイルに含めてください。次のコードを参照してください。

%%sh
mkdir mymodel
mv serving.properties mymodel/
tar czvf mymodel.tar.gz mymodel/
rm -rf mymodel

s3_code_prefix = "large-model-lmi/code"
bucket = sess.default_bucket()  # bucket to house artifacts
code_artifact = sess.upload_data("mymodel.tar.gz", bucket, s3_code_prefix)print(f"S3 Code or Model tar ball uploaded to --- > {code_artifact}")

DJL コンテナイメージを取得して SageMaker モデルを作成します。

##Retrieve djl container image
image_uri = image_uris.retrieve(
        framework="djl-deepspeed",
        region=sess.boto_session.region_name,
        version="0.21.0"
    )
image_uri = image_uri.split(":")[0] + ":" + "0.22.1-neuronx-sdk2.9.0"

model = Model(image_uri=image_uri, model_data=code_artifact, env=env, role=role)

次に、先に定義したモデル構成で SageMaker エンドポイントを作成します。SageMaker が /tmp を Amazon Elastic Block Store (Amazon EBS) にマップするため、コンテナはモデルを /tmp スペースにダウンロードします。volume_size パラメータを追加して、/tmp ディレクトリにモデルをダウンロードしてコンパイルするのに十分なスペースを確保する必要があります。container_startup_health_check_timeout を 3,600 秒に設定して、モデルの準備が整った後にヘルスチェックが開始されるようにします。ここでは ml.inf2.8xlarge インスタンスを使用しています。次のコードを参照してください。

instance_type = "ml.inf2.8xlarge"
endpoint_name = sagemaker.utils.name_from_base("lmi-model")


model.deploy(initial_instance_count=1,
             instance_type=instance_type,
             endpoint_name=endpoint_name,
             container_startup_health_check_timeout=3600,
             volume_size=256
            )

SageMaker エンドポイントが作成されたら、Predictor オブジェクトを使用して SageMaker エンドポイントに対してリアルタイムで予測を行うことができます。

# our requests and responses will be in json format so we specify the serializer and the deserializer
predictor = sagemaker.Predictor(
    endpoint_name=endpoint_name,
    sagemaker_session=sess,
    serializer=serializers.JSONSerializer(),
    deserializer=deserializers.JSONDeserializer(),
)

predictor.predict(
    {"inputs": "write a blog on new York", "parameters": {}}
)

後片付け

テストが完了したら、エンドポイントを削除してコストを節約してください。

# - Delete the end point
sess.delete_endpoint(endpoint_name)
sess.delete_endpoint_config(endpoint_name)
model.delete_model()

まとめ

本記事では、新たに導入された SageMaker の機能について紹介しました。この新機能は、生成系 AI モデルをホストするための ml.inf2 および ml.trn1 インスタンスをサポートするようになりました。SageMaker と LMI コンテナを使用して、生成系 AI モデルである GPT4ALL-J を AWS Inferentia2 にデプロイする方法を、コードをまったく記述せずにデモンストレーションしました。また、DJLServing と transformers-neuronx を使用してモデルをロード、パーティショニング、サーブする方法も紹介しました。

Inf2 インスタンスは、生成系 AI モデルを AWS で実行する最も費用対効果の高い方法です。パフォーマンスの詳細については、「Inf2 Performance」を参照してください。

ノートブックの例については、GitHub リポジトリをご覧ください。試してみて、ご不明な点がございましたらお知らせください。
________________________________________
本稿は、2023年5月4日に公開された “Achieve high performance with lowest cost for generative AI inference using AWS Inferentia2 and AWS Trainium on Amazon SageMaker” を翻訳したものです。翻訳はアンナプルナラボの常世が担当しました。

著者について

Vivek Gangasani は、アマゾンウェブサービスのシニア機械学習ソリューションアーキテクトです。機械学習スタートアップ企業と協力して、AWS で AI/ML アプリケーションを構築してデプロイしています。現在は、MLOps、ML 推論、ローコード ML 向けのソリューションの提供に注力しています。彼は自然言語処理やコンピュータービジョンなど、さまざまな分野のプロジェクトに取り組んできました。

常世 大史 は、AWS アンナプルナラボのソリューションアーキテクトです。日本を拠点とし、AWS に買収される前からアンナプルナラボに入社し、一貫してアンナプルナラボの技術でお客様を支援してきました。最近は、専用のシリコン、AWS Inferentia、Trainiumをベースにした機械学習ソリューションに焦点を当てています。

Dhawal Patel は、AWS のプリンシパル機械学習アーキテクトです。彼は、分散コンピューティングや人工知能に関連する問題について、大企業から中規模のスタートアップまで、さまざまな組織と協力してきました。彼は自然言語処理やコンピュータービジョンの分野を含む深層学習に焦点を当てています。お客様が SageMaker で高性能なモデル推論を実現できるよう支援しています。

Qing Lan は、AWS のソフトウェア開発エンジニアです。高性能な ML 推論ソリューションや高性能なロギングシステムなど、Amazon におけるいくつかの挑戦的な製品に携わってきました。Qing のチームは、非常に低いレイテンシーを要求される Amazon Advertising において、数十億規模のパラメータを持つモデルの初めての立ち上げに成功しました。インフラストラクチャーの最適化と深層学習の高速化に関する深い知識を持っています。

Qingwei Li は、Amazon Web Services の機械学習スペシャリストです。オペレーションズ・リサーチの領域で博士号を取得しましたが、約束したノーベル賞の受賞を果たさないまま、指導教官の研究助成金を使い切ってしまいました。現在は、金融サービスや保険業界のお客様が AWS 上で機械学習ソリューションを構築するのを支援しています。読書と教えることが趣味。

Alan Tan は、SageMaker のシニアプロダクトマネージャーで、大規模モデル推論の取り組みを主導しています。機械学習をアナリティクスの分野に応用することに情熱を注いでいます。仕事以外では、彼はアウトドアを楽しんでいます。

Varun Syal は、Amazon Sagemaker のソフトウェア開発エンジニアで、ML 推論プラットフォームの重要な顧客向け機能に取り組んでいます。分散システムと AI の分野での仕事に情熱を注いでいます。余暇には、読書とガーデニングが好きです。