AWS Inferentia2 で安定した拡散のパフォーマンスを最大化し、推論コストを削減 | アマゾン ウェブ サービス

生成 AI モデルは、リアルなテキスト、画像、コード、オーディオを作成する優れた機能により、ここ数カ月で急速に成長しています。 これらのモデルの中でも、安定拡散モデルは、テキスト プロンプトに基づいて高品質の画像を作成するという独自の強みで際立っています。 安定した拡散により、リアルなポートレート、風景、さらには抽象芸術など、さまざまな高品質の画像を生成できます。 また、他の生成 AI モデルと同様に、安定拡散モデルには、低遅延の推論を提供する強力なコンピューティングが必要です。

この投稿では、安定拡散モデルを実行し、低コストで高いパフォーマンスを達成する方法を示します。 アマゾン エラスティック コンピューティング クラウド (Amazon EC2) の使用 Amazon EC2 Inf2インスタンス を搭載 AWS インフェレンシア 2。 安定拡散モデルのアーキテクチャを見て、次を使用して安定拡散モデルをコンパイルする手順を順を追って説明します。 AWS ニューロン そしてそれを Inf2 インスタンスにデプロイします。 また、Neuron SDK がパフォーマンスを向上させるために自動的に行う最適化についても説明します。 Stable Diffusion 2.1 と 1.5 の両方のバージョンを AWS Inferentia2 でコスト効率よく実行できます。 最後に、安定した拡散モデルを Inf2 インスタンスにデプロイする方法を示します。 アマゾンセージメーカー.

Stable Diffusion 2.1 モデルのサイズは、浮動小数点 32 (FP32) では 5 GB、bfoat2.5 (BF16) では 16 GB です。 単一の inf2.xlarge インスタンスには、2 GB の HBM メモリを備えた 32 つの AWS Inferentia2.1 アクセラレータがあります。 Stable Diffusion 2 モデルは、単一の infXNUMX.xlarge インスタンスに適合します。 Stable Diffusion は、入力としてテキスト プロンプトを提供するだけで、さまざまなスタイルやコンテンツの画像を作成するために使用できるテキストから画像へのモデルです。 安定拡散モデルのアーキテクチャの詳細については、を参照してください。 安定拡散モデルで高品質のイメージを作成し、Amazon SageMaker でコスト効率よくデプロイします。.

Neuron SDK が安定した拡散のパフォーマンスを最適化する方法

Stable Diffusion 2.1 モデルを AWS Inferentia2 インスタンスにデプロイする前に、 ニューロン SDK。 Neuron SDK にはディープ ラーニング コンパイラー、ランタイム、ツールが含まれており、ディープ ラーニング モデルをコンパイルして自動的に最適化するため、Inf2 インスタンス上で効率的に実行し、AWS Inferentia2 アクセラレータのパフォーマンスを最大限に引き出すことができます。 Stable Diffusion 2.1 モデルで利用可能な例が、 GitHubレポ。 このノートブックでは、安定拡散モデルをコンパイルし、コンパイルされた Neuron モデルを保存し、推論のためにランタイムにロードする方法のエンドツーエンドの例を示します。

を使用しております StableDiffusionPipeline 抱き合う顔から diffusers モデルをロードしてコンパイルするためのライブラリ。 次に、次を使用して Neuron のモデルのすべてのコンポーネントをコンパイルします。 torch_neuronx.trace() そして最適化されたモデルをTorchScriptとして保存します。 コンパイル プロセスはメモリを大量に消費する可能性があり、大量の RAM が必要になります。 これを回避するには、各モデルをトレースする前に、 deepcopy トレースされているパイプラインの部分。 これに続いて、次を使用してメモリからパイプライン オブジェクトを削除します。 del pipe。 この手法は、RAM が少ないインスタンスでコンパイルする場合に特に役立ちます。

さらに、安定拡散モデルの最適化も実行します。 UNet は、推論の中で最も計算量が多い側面を保持します。 UNet コンポーネントは、バッチ サイズ XNUMX の入力テンソルを操作し、バッチ サイズ XNUMX の対応する出力テンソルを生成して、単一のイメージを生成します。 これらのバッチ内の要素は互いに完全に独立しています。 この動作を利用して、各 Neuron コアで XNUMX つのバッチを実行することで、最適なレイテンシーを得ることができます。 (XNUMX つのバッチで入力テンソルを使用して) XNUMX つのバッチ用に UNet をコンパイルし、 torch_neuronx.DataParallel この単一のバッチ モデルを各コアにロードするための API。 この API の出力はシームレスな XNUMX バッチ モジュールです。UNet に XNUMX つのバッチの入力を渡すと、XNUMX バッチの出力が返されますが、内部的には XNUMX つの単一バッチ モデルが XNUMX つの Neuron コアで実行されます。 。 この戦略により、リソースの使用率が最適化され、待ち時間が短縮されます。

Inf2 EC2 インスタンス上で Stable Diffusion モデルをコンパイルしてデプロイする

Stable Diffusion モデルをコンパイルして Inf2 EC2 インスタンスにデプロイするには、 AWSマネジメントコンソール そして、inf2.8xlarge インスタンスを作成します。 コンパイルにはより多くのホスト メモリが必要となるため、inf2.8xlarge インスタンスはモデルのコンパイルにのみ必要であることに注意してください。 Stable Diffusion モデルは、inf2.xlarge インスタンスでホストできます。 次のコマンドを使用して、Neuron ライブラリを含む最新の AMI を見つけることができます。 AWSコマンドラインインターフェイス (AWS CLI) コマンド:

aws ec2 describe-images --region us-east-1 --owners amazon --filters 'Name=name,Values=Deep Learning AMI Neuron PyTorch 1.13.? (Amazon Linux 2) ????????' 'Name=state,Values=available' --query 'reverse(sort_by(Images, &CreationDate))[:1].ImageId' --output text

この例では、Deep Learning AMI Neuron PyTorch 2 (Ubuntu 1.13) を使用して EC20.04 インスタンスを作成しました。 その後、インスタンスに接続して次の手順を実行することで、JupyterLab ラボ環境を作成できます。

run source /opt/aws_neuron_venv_pytorch/bin/activate
pip install jupyterlab
jupyter-lab

モデルをコンパイルしてホストするためのすべての手順が記載されたノートブックは、次の場所にあります。 GitHubの.

テキスト エンコーダー ブロックの XNUMX つのコンパイル手順を見てみましょう。 Stable Diffusion パイプラインの一部である他のブロックも同様にコンパイルできます。

最初のステップは、Hugging Face から事前トレーニングされたモデルをロードすることです。 の StableDiffusionPipeline.from_pretrained メソッドは、事前トレーニングされたモデルをパイプライン オブジェクトにロードします。 pipe。 次に、 deepcopy パイプラインからテキスト エンコーダーを取得し、効果的にクローンを作成します。 の del pipe 次に、コマンドを使用して元のパイプライン オブジェクトを削除し、それによって消費されていたメモリを解放します。 ここでは、モデルを BF16 重みに量子化しています。

model_id = "stabilityai/stable-diffusion-2-1-base"
pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.bfloat16)
text_encoder = copy.deepcopy(pipe.text_encoder)
del pipe

このステップでは、テキスト エンコーダを次のようにラップします。 NeuronTextEncoder ラッパー。 コンパイルされたテキスト エンコーダ モジュールの出力は次のようになります。 dict。 それを次のように変換します。 list このラッパーを使用して入力します。

text_encoder = NeuronTextEncoder(text_encoder)

PyTorch tensor を初期化します emb いくつかの値を使用して。 の emb テンソルは、 torch_neuronx.trace 関数。 この関数はテキスト エンコーダーをトレースし、Neuron に最適化された形式にコンパイルします。 コンパイルされたモデルのディレクトリ パスは結合によって構築されます。 COMPILER_WORKDIR_ROOT サブディレクトリを使用して text_encoder:

emb = torch.tensor([...])
text_encoder_neuron = torch_neuronx.trace(
        text_encoder.neuron_text_encoder,
        emb,
        compiler_workdir=os.path.join(COMPILER_WORKDIR_ROOT, 'text_encoder'),
        )

コンパイルされたテキスト エンコーダーは次を使用して保存されます。 torch.jit.save。 ファイル名model.ptで保存されています。 text_encoder コンパイラのワークスペースのディレクトリ:

text_encoder_filename = os.path.join(COMPILER_WORKDIR_ROOT, 'text_encoder/model.pt')
torch.jit.save(text_encoder_neuron, text_encoder_filename)

　 ノート モデルの他のコンポーネント (UNet、VAE デコーダー、VAE) をコンパイルするための同様の手順が含まれています。 post_quant_conv。 すべてのモデルをコンパイルした後、次の手順に従ってモデルをロードして実行できます。

1. コンパイルされたモデルのパスを定義します。
2. 事前トレーニングされたものをロードする StableDiffusionPipeline モデル。その構成は bfloat16 データ型を使用するように指定されています。
3. 次のコマンドを使用して、UNet モデルを XNUMX つの Neuron コアにロードします。 torch_neuronx.DataParallel API。 これにより、データの並列推論が実行できるようになり、モデルのパフォーマンスが大幅に向上します。
4. モデルの残りの部分をロードします(text_encoder, decoder, post_quant_conv) を単一の Neuron コアに配置します。

その後、入力テキストをプロンプトとして指定してパイプラインを実行できます。 以下は、プロンプト用にモデルによって生成されたいくつかの画像です。

• ルノー・セーチャンの肖像、ペンとインク、複雑な線画、クレイグ・マリンズ、ルアン・ジア、三浦健太郎、グレッグ・ルトコウスキー、ロウンドロー作

• 19 世紀の古い炭鉱夫の肖像画、美しい絵、グレッグ・ルトコウスキーによる非常に詳細な顔の絵

• 森の真ん中にあるお城

AWS Inferentia2.1 および SageMaker 上の Host Stable Diffusion 2

SageMaker で Stable Diffusion モデルをホストするには、Neuron SDK でコンパイルする必要もあります。 Large Model Inference (LMI) コンテナーを使用すると、事前にコンパイルを完了することも、実行中にコンパイルを完了することもできます。 事前にコンパイルすると、モデルの読み込み時間を短縮できるため、推奨されるオプションです。

SageMaker LMI コンテナでは、モデルをデプロイする XNUMX つの方法が提供されます。

• コード不要のオプションで、 serving.properties 必要な構成を含むファイル
• 独自の推論スクリプトを持ち込む

両方のソリューションを検討し、構成と推論スクリプトを確認します (model.py）。 この投稿では、 Amazon シンプル ストレージ サービス (Amazon S3) バケット。 このコンパイル済みモデルを展開に使用できます。

提供されたスクリプトを使用してモデルを構成する

このセクションでは、安定拡散モデルをホストするために LMI コンテナを構成する方法を示します。 SD2.1 ノートブックは次の場所で入手できます。 GitHubの。 最初のステップは、次のディレクトリ構造に従ってモデル構成パッケージを作成することです。 私たちの目的は、モデルをホストするために必要な最小限のモデル構成を使用することです。 必要なディレクトリ構造は次のとおりです。

<config-root-directory> / 
    ├── serving.properties
    │   
    └── model.py [OPTIONAL]

次に、 サービング.プロパティ 次のパラメータを含むファイル:

%%writefile code_sd/serving.properties
engine=Python
option.entryPoint=djl_python.transformers-neuronx
option.use_stable_diffusion=True
option.model_id=s3url
option.tensor_parallel_degree=2
option.dtype=bf16

パラメータでは以下を指定します。

• オプション.モデルID – LMI コンテナは s5cmd を使用して S3 の場所からモデルをロードするため、コンパイルされた重みが存在する場所を指定する必要があります。
• オプション.エントリポイント – 組み込みハンドラーを使用するには、transformers-neuronx クラスを指定します。 カスタム推論スクリプトがある場合は、代わりにそれを提供する必要があります。
• オプション.dtype – これは、特定のサイズでウェイトをロードするように指定します。 この記事では、BF16 を使用します。これにより、FP32 と比較してメモリ要件がさらに削減され、それによりレイテンシーが短縮されます。
• オプション.tensor_Parallel_degree – このパラメータは、このモデルに使用するアクセラレータの数を指定します。 AWS Inferentia2 チップ アクセラレータには 2 つの Neuron コアがあるため、値 XNUMX を指定すると、XNUMX つのアクセラレータ (XNUMX つのコア) を使用することになります。 これは、複数のワーカーを作成してエンドポイントのスループットを向上できることを意味します。
• オプション.エンジン – これは、このホスティングに DeepSpeed や Faster Transformer などの他のコンパイラを使用しないことを示すために Python に設定されています。

自分のスクリプトを持ってくる

独自のカスタム推論スクリプトを使用する場合は、 option.entryPoint から serving.properties。 この場合の LMI コンテナは、 model.py ファイルと同じ場所にあります serving.properties それを使用して推論を実行します。

独自の推論スクリプト (model.py) を作成する

独自の推論スクリプトの作成は、LMI コンテナを使用すると比較的簡単です。 コンテナにはあなたの必要があります model.py ファイルに次のメソッドを実装します。

def handle(inputs: Input) which returns an object of type Outputs

の重要な領域のいくつかを調べてみましょう。 付属のノートこれは、独自のスクリプトの使用機能を示しています。

置き換える cross_attention 最適化されたバージョンのモジュール:

# Replace original cross-attention module with custom cross-attention module for better performance
    CrossAttention.get_attention_scores = get_attention_scores
Load the compiled weights for the following
text_encoder_filename = os.path.join(COMPILER_WORKDIR_ROOT, 'text_encoder.pt')
decoder_filename = os.path.join(COMPILER_WORKDIR_ROOT, 'vae_decoder.pt')
unet_filename = os.path.join(COMPILER_WORKDIR_ROOT, 'unet.pt')
post_quant_conv_filename =. os.path.join(COMPILER_WORKDIR_ROOT, 'vae_post_quant_conv.pt')

これらは、コンパイルを作成するときに使用したコンパイル済みウェイト ファイルの名前です。 ファイル名は自由に変更できますが、ウェイト ファイル名がここで指定したものと一致していることを確認してください。

次に、Neuron SDK を使用してそれらをロードし、実際のモデルの重みに設定する必要があります。 UNet に最適化された重みをロードするときは、これらをロードする必要がある Neuron コアの数も指定していることに注意してください。 ここでは、XNUMX つのコアを持つ単一のアクセラレータにロードします。

# Load the compiled UNet onto two neuron cores.
    pipe.unet = NeuronUNet(UNetWrap(pipe.unet))
    logging.info(f"Loading model: unet:created")
    device_ids = [idx for idx in range(tensor_parallel_degree)]
   
    pipe.unet.unetwrap = torch_neuronx.DataParallel(torch.jit.load(unet_filename), device_ids, set_dynamic_batching=False)
   
 
    # Load other compiled models onto a single neuron core.
 
    # - load encoders
    pipe.text_encoder = NeuronTextEncoder(pipe.text_encoder)
    clip_compiled = torch.jit.load(text_encoder_filename)
    pipe.text_encoder.neuron_text_encoder = clip_compiled
    #- load decoders
    pipe.vae.decoder = torch.jit.load(decoder_filename)
    pipe.vae.post_quant_conv = torch.jit.load(post_quant_conv_filename)

プロンプトを使用して推論を実行すると、パイプ オブジェクトが呼び出されてイメージが生成されます。

SageMaker エンドポイントを作成する

Boto3 API を使用して SageMaker エンドポイントを作成します。 次の手順を実行します。

1. サービングとオプションのみを使用して tarball を作成します model.py ファイルを作成し、Amazon S3 にアップロードします。
2. イメージ コンテナーと前にアップロードしたモデル tarball を使用してモデルを作成します。
3. 次の主要なパラメータを使用してエンドポイント構成を作成します。
   1. 使用する ml.inf2.xlarge インスタンス。
   2. 作成セッションプロセスで ContainerStartupHealthCheckTimeoutInSeconds モデルのデプロイ後にヘルスチェックが確実に開始されるようにするには、240 に設定します。
   3. 作成セッションプロセスで VolumeInGB サイズが 32 GB のモデルの重みをロードするために使用できるように、より大きな値に設定します。

SageMaker モデルを作成する

model.tar.gz ファイルを作成して Amazon S3 にアップロードした後、SageMaker モデルを作成する必要があります。 LMI コンテナと前のステップのモデル アーティファクトを使用して、SageMaker モデルを作成します。 SageMaker を使用すると、さまざまな環境変数をカスタマイズして挿入できます。 このワークフローでは、すべてをデフォルトのままにすることができます。 次のコードを参照してください。

inference_image_uri = (
    f"763104351884.dkr.ecr.{region}.amazonaws.com/djl-inference:0 djl-serving-inf2"
)

モデル オブジェクトを作成します。これは基本的に、インスタンスにロードされて推論に使用されるロックダウン コンテナーを作成します。

model_name = name_from_base(f"inf2-sd")
create_model_response = boto3_sm_client.create_model(
    ModelName=model_name,
    ExecutionRoleArn=role,
    PrimaryContainer={"Image": inference_image_uri, "ModelDataUrl": s3_code_artifact},
)

SageMakerエンドポイントを作成する

このデモでは、ml.inf2.xlarge インスタンスを使用します。 を設定する必要があります。 VolumeSizeInGB パラメータを使用して、モデルと重みをロードするために必要なディスク領域を提供します。 このパラメータは、 Amazon Elastic Blockストア (Amazon EBS) ボリュームのアタッチメント。 モデルのダウンロードのタイムアウトとコンテナの起動ヘルスチェックをより高い値のままにすることができます。これにより、コンテナが Amazon S3 から重みを取得して AWS Inferentia2 アクセラレータにロードするのに十分な時間が得られます。 詳細については、を参照してください。 CreateEndpointConfig.

endpoint_config_response = boto3_sm_client.create_endpoint_config( EndpointConfigName=endpoint_config_name,
    ProductionVariants=[
        {
            "VariantName": "variant1",
            "ModelName": model_name,
            "InstanceType": "ml.inf2.xlarge", # - 
            "InitialInstanceCount": 1,
            "ContainerStartupHealthCheckTimeoutInSeconds": 360, 
            "VolumeSizeInGB": 400
        },
    ],
)

最後に、SageMaker エンドポイントを作成します。

create_endpoint_response = boto3_sm_client.create_endpoint(
    EndpointName=f"{endpoint_name}", EndpointConfigName=endpoint_config_name
)

モデルエンドポイントを呼び出す

これは生成モデルであるため、モデルが画像を生成するために使用するプロンプトを渡します。 ペイロードのタイプは JSON です。

response_model = boto3_sm_run_client.invoke_endpoint( EndpointName=endpoint_name,
    Body=json.dumps(
        {
            "prompt": "Mountain Landscape", 
            "parameters": {} # 
        }
    ), 
    ContentType="application/json",
)

Inf2 での安定拡散モデルのベンチマーク

Inf16 で BF 2 データ型を使用して安定拡散モデルのベンチマークを行うためにいくつかのテストを実行しました。その結果、安定拡散の他のアクセラレータに匹敵する、またはそれを超えるレイテンシの数値を導き出すことができました。 これは、AWS Inferentia2 チップの低コストと相まって、非常に価値のある提案になります。

次の数値は、inf2.xl インスタンスにデプロイされた安定拡散モデルからのものです。 費用の詳細については、を参照してください。 Amazon EC2 Inf2インスタンス.

まとめ

この投稿では、Inf2.1 インスタンスを使用した Stable Diffusion 2 モデルのコンパイル、最適化、デプロイメントについて詳しく説明します。 また、SageMaker を使用した安定拡散モデルのデプロイメントも実証しました。 Inf2 インスタンスは、Stable Diffusion 1.5 に対して優れたコスト パフォーマンスも提供します。 Inf2 インスタンスが生成 AI および大規模言語モデルに最適な理由の詳細については、以下を参照してください。 低コスト、高性能の生成 AI 推論用の Amazon EC2 Inf2 インスタンスが一般提供開始。 性能の詳細については、を参照してください。 Inf2 パフォーマンス。 追加の例を確認してください。 GitHubレポ.

貴重な情報をレビューし、提供してくださった Matthew Mcclain、Beni Hegedus、Kamran Khan、Shruti Koparkar、Qing Lan に心より感謝いたします。

著者について

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

KC トゥン AWS Annapurna Labs のシニア ソリューション アーキテクトです。 彼は、クラウドでの大規模なディープ ラーニング モデルのトレーニングと展開を専門としています。 彼は博士号を持っています。 ダラスのテキサス大学サウスウェスタン医療センターで分子生物物理学の博士号を取得しています。 彼は AWS Summits と AWS Reinvent で講演しました。 現在、彼はお客様が AWS クラウドで大規模な PyTorch および TensorFlow モデルをトレーニングおよびデプロイするのを支援しています。 彼は次の XNUMX 冊の本の著者です。 TensorFlow Enterprise を学ぶ および TensorFlow 2 ポケット リファレンス.

ルピンダー・グレワル AWS のシニア Ai/ML スペシャリスト ソリューション アーキテクトです。 彼は現在、SageMaker でのモデルと MLOps の提供に注力しています。 この役職に就く前は、モデルの構築とホスティングを行う機械学習エンジニアとして働いていました。 仕事以外では、テニスや山道でのサイクリングを楽​​しんでいます。

• SEO を活用したコンテンツと PR 配信。 今日増幅されます。
• PlatoData.Network 垂直生成 Ai。 自分自身に力を与えましょう。 こちらからアクセスしてください。
• プラトアイストリーム。 Web3 インテリジェンス。 知識増幅。 こちらからアクセスしてください。
• プラトンESG。 自動車/EV、 カーボン、 クリーンテック、 エネルギー、 環境、 太陽、 廃棄物管理。 こちらからアクセスしてください。
• ブロックオフセット。 環境オフセット所有権の近代化。 こちらからアクセスしてください。
• 情報源： https://aws.amazon.com/blogs/machine-learning/maximize-stable-diffusion-performance-and-lower-inference-costs-with-aws-inferentia2/