AWS Inferentia2を使用して、安定したディフュージョンのパフォーマンスを最大化し、推論コストを低減します

使用するAWS Inferentia2で、ディフュージョンのパフォーマンスを最大化し、推論コストを低減します

生成AIモデルは、最近の数ヶ月間で急速な成長を遂げており、リアルなテキスト、画像、コード、音声を生成する驚異的な能力を持っています。これらのモデルの中でも、Stable Diffusionモデルは、テキストのプロンプトに基づいて高品質な画像を生成するという独自の強みで注目されています。Stable Diffusionは、リアルな肖像画、風景、抽象的なアートなど、さまざまな高品質な画像を生成することができます。そして、他の生成AIモデルと同様に、Stable Diffusionモデルは、低レイテンシの推論を提供するために強力なコンピューティングを必要とします。

この記事では、Amazon Elastic Compute Cloud（Amazon EC2）でStable Diffusionモデルを実行し、最低コストで高性能を実現する方法を紹介します。これには、AWS Inferentia2によってパワードされたAmazon EC2 Inf2インスタンスを使用します。Stable Diffusionモデルのアーキテクチャと、AWS Neuronを使用してStable Diffusionモデルをコンパイルし、Inf2インスタンスにデプロイする手順を説明します。また、Neuron SDKが自動的に行う最適化についても説明します。Stable Diffusion 2.1および1.5のバージョンを、AWS Inferentia2上で効率的に実行することができます。最後に、Amazon SageMakerを使用してStable DiffusionモデルをInf2インスタンスにデプロイする方法を示します。

Stable Diffusion 2.1モデルの浮動小数点32（FP32）のサイズは5GBであり、bfoat16（BF16）では2.5GBです。単一のinf2.xlargeインスタンスには、32GBのHBMメモリを持つ1つのAWS Inferentia2アクセラレータが搭載されています。Stable Diffusion 2.1モデルは、単一のinf2.xlargeインスタンスに収まるサイズです。Stable Diffusionは、テキストから画像への変換モデルであり、テキストプロンプトを入力として提供するだけで、さまざまなスタイルとコンテンツの画像を作成することができます。Stable Diffusionモデルアーキテクチャの詳細については、「Stable Diffusionモデルを使用して高品質な画像を作成し、Amazon SageMakerでコスト効率よくデプロイする」を参照してください。

Neuron SDKがStable Diffusionのパフォーマンスを最適化する方法

AWS Inferentia2インスタンスにStable Diffusion 2.1モデルをデプロイする前に、Neuron SDKを使用してモデルコンポーネントをコンパイルする必要があります。Neuron SDKには、ディープラーニングコンパイラ、ランタイム、ツールが含まれており、ディープラーニングモデルをコンパイルし、Inf2インスタンスで効率的に実行し、AWS Inferentia2アクセラレータのパフォーマンスを最大限に引き出すように自動的に最適化します。Stable Diffusion 2.1モデルの例は、GitHubリポジトリで利用可能です。このノートブックでは、Stable Diffusionモデルをコンパイルし、コンパイルされたNeuronモデルを保存し、推論のためにランタイムにロードするエンドツーエンドの例を紹介します。

• 「ICML 2023でのGoogle」
• 「TidyBotでの掃除」
• 「OpenAIがオープンソースのGPTモデルのリリースを予告」

私たちは、Hugging FaceのdiffusersライブラリからStableDiffusionPipelineを使用してモデルをロードし、コンパイルします。そして、torch_neuronx.trace()を使用してモデルのすべてのコンポーネントをNeuron用にコンパイルし、最適化されたモデルをTorchScriptとして保存します。コンパイルプロセスは、かなりのメモリを要求するため、メモリが低いインスタンスでコンパイルする場合には、各モデルのトレースの前にトレースされるパイプラインのdeepcopyを作成します。その後、del pipeを使用してパイプラインオブジェクトをメモリから削除します。このテクニックは、低いRAMを持つインスタンスでコンパイルする場合に特に有用です。

さらに、Stable Diffusionモデルに最適化も行っています。UNetは推論の計算量が最も大きい部分です。UNetコンポーネントは、バッチサイズ2の入力テンソルで動作し、バッチサイズ2の対応する出力テンソルを生成して単一の画像を生成します。これらのバッチ内の要素は互いに完全に独立しています。この動作を利用して、各Neuronコアで1つのバッチを実行することで、最適なレイテンシを得ることができます。バッチサイズ1のUNetをコンパイルし（バッチサイズ1の入力テンソルを使用）、torch_neuronx.DataParallel APIを使用してこの単一バッチモデルを各コアにロードします。このAPIの出力はシームレスな2バッチモジュールであり、2つのバッチの入力をUNetに渡すと、2バッチの出力が返されますが、内部では2つのシングルバッチモデルが2つのNeuronコア上で実行されています。この戦略はリソースの利用効率を最適化し、レイテンシを削減します。

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

Inf2 EC2インスタンスにStable Diffusionモデルをコンパイルしてデプロイするには、AWS Management Consoleにサインインし、inf2.8xlargeインスタンスを作成します。なお、モデルのコンパイルにはホストメモリがより多く必要とされるため、inf2.8xlargeインスタンスが必要です。Stable Diffusionモデルはinf2.xlargeインスタンスでホストできます。次のAWS Command Line Interface（AWS CLI）コマンドを使用して、Neuronライブラリが最新のAMIを見つけることができます：

aws ec2 describe-images --region us-east-1 --owners amazon \
--filters 'Name=name,Values=ディープラーニング 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 1.13 (Ubuntu 20.04) を使用して EC2 インスタンスを作成しました。次の手順を実行して、JupyterLab ラボ環境を作成できます。

source /opt/aws_neuron_venv_pytorch/bin/activate を実行します
pip install jupyterlab を実行します
jupyter-lab を実行します

モデルのコンパイルとホスティングの手順についてのノートブックは GitHub にあります。

テキストエンコーダーブロックのコンパイル手順を見てみましょう。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 テンソル 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 を使用して保存されます。それは、コンパイラのワークスペースの text_encoder ディレクトリ内の model.pt というファイル名で保存されます：

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. bfloat16 データ型を使用するように構成された事前学習済みの StableDiffusionPipeline モデルをロードします。
3. UNet モデルを 2 つの Neuron コアにロードします。これにより、データ並列推論が実行され、モデルのパフォーマンスが大幅に向上します。
4. モデルの残りの部分（text_encoder、decoder、post_quant_conv）を単一の Neuron コアにロードします。

それから、入力テキストをプロンプトとして提供することでパイプラインを実行できます。以下は、モデルがプロンプトに対して生成したいくつかの画像です：

• ルノー・セヤンのポートレート、ペンとインク、複雑な線画、クレイグ・マリンズ、ルアン・ジア、ケンタロウ・ミウラ、グレッグ・ルトコウスキ、ラウンドロー

• 19世紀の古い石炭鉱夫の肖像画、グレッグ・ルトコウスキによる高詳細な顔の描画が美しい絵画です

• 森の中にある城

AWS Inferentia2およびSageMakerでStable Diffusion 2.1をホストする

SageMakerでStable Diffusionモデルをホストするには、Neuron SDKでのコンパイルが必要です。コンパイルは事前に完了するか、Large Model Inference（LMI）コンテナを使用してランタイム中に完了することができます。事前にコンパイルすると、モデルの読み込み時間が短くなり、推奨されるオプションです。

SageMaker LMIコンテナには、モデルを展開する2つの方法があります：

• 必要な構成を持つserving.propertiesファイルのみを提供するノーコードオプション
• 独自の推論スクリプトを使用する

両方の解決策を見て、構成と推論スクリプト（model.py）を説明します。この記事では、Amazon Simple Storage Service（Amazon S3）バケットに保存されている事前コンパイル済みモデルを使用してデプロイメントをデモンストレーションします。これを使用してデプロイメントを行うことができます。

提供されたスクリプトでモデルを構成する

このセクションでは、LMIコンテナを設定してStable Diffusionモデルをホストする方法を示します。GitHubで利用可能なSD2.1ノートブックを使用します。最初のステップは、次のディレクトリ構造に従ってモデル構成パッケージを作成することです。モデルをホストするために必要な最小限のモデル構成を使用することを目指しています。必要なディレクトリ構造は次のようになります：

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

次に、次のパラメータを使用してserving.propertiesファイルを作成します：

%%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

これらのパラメータは以下を指定します：

• option.model_id – LMIコンテナはS3の場所からモデルをロードするため、コンパイル済みの重みの場所を指定する必要があります。
• option.entryPoint – 組み込みハンドラを使用するため、transformers-neuronxクラスを指定します。カスタムの推論スクリプトを使用する場合は、代わりにそれを提供する必要があります。
• option.dtype – これは特定のサイズで重みをロードすることを指定します。この記事ではBF16を使用し、メモリ要件をさらに低減し、レイテンシを低くします。
• option.tensor_parallel_degree – このパラメータはこのモデルに使用するアクセラレータの数を指定します。AWS Inferentia2チップアクセラレータには2つのNeuronコアがあり、値を2に指定することで1つのアクセラレータ（2つのコア）を使用します。これにより、エンドポイントのスループットを向上させるために複数のワーカーを作成できるようになります。
• option.engine – これはPythonに設定されており、このホスティングにDeepSpeedやFaster Transformerなどの他のコンパイラを使用しないことを示します。

独自のスクリプトを使用する

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

自分自身の推論スクリプトを作成する（model.py）

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

def handle(inputs: Input) -> Outputsの型でオブジェクトを返す

付属のノートブックの重要な領域のいくつかを見てみましょう。これによって独自のスクリプト機能がデモンストレーションされます。

cross_attentionモジュールを最適化バージョンで置き換えてください：

# オリジナルのクロスアテンションモジュールをパフォーマンス向上のためにカスタムクロスアテンションモジュールで置き換える
    CrossAttention.get_attention_scores = get_attention_scores
以下のコンパイル済みのウェイトをロードする
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コアの数も指定しています。ここでは、2つのコアを持つシングルアクセラレータにロードしています：

# 2つのNeuronコアにコンパイル済みのUNetをロードする
    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)
   
 
    # 他のコンパイル済みモデルをシングルなNeuronコアにロードする。
 
    # - エンコーダをロードする
    pipe.text_encoder = NeuronTextEncoder(pipe.text_encoder)
    clip_compiled = torch.jit.load(text_encoder_filename)
    pipe.text_encoder.neuron_text_encoder = clip_compiled
    # - デコーダをロードする
    pipe.vae.decoder = torch.jit.load(decoder_filename)
    pipe.vae.post_quant_conv = torch.jit.load(post_quant_conv_filename)

プロンプトを使用して推論を実行すると、pipeオブジェクトが画像を生成します。

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

SageMakerエンドポイントを作成するためにBoto3 APIを使用します。次の手順を完了します：

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

SageMakerモデルを作成する

モデル.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 Store (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でのStable Diffusionモデルのベンチマーク

Inf2上でBF 16データ型を使用してStable Diffusionモデルのベンチマークテストを実行し、他のアクセラレータに匹敵するまたは超えるレイテンシの数値を得ることができました。これに加えて、AWS Inferentia2チップの低コストもあり、これは非常に価値ある提案です。

以下の数値は、inf2.xlインスタンスに展開されたStable Diffusionモデルのものです。コストに関する詳細については、Amazon EC2 Inf2インスタンスを参照してください。

モデル Stable Diffusion 1.5 Stable Diffusion 1.5 Stable Diffusion 1.5 Stable Diffusion 1.5	解像度 512×512 768×768 512×512 768×768	データ型 bf16 bf16 bf16 bf16	イテレーション数 50 50 30 30	P95レイテンシ（ms） 2,427.4 8,235.9 1,456.5 4,941.6	Inf2.xlオンデマンドコスト（1時間あたり） $0.76 $0.76 $0.76 $0.76	Inf2.xl（1枚あたりのコスト） $0.0005125 $0.0017387 $0.0003075 $0.0010432
Stable Diffusion 2.1 Stable Diffusion 2.1 Stable Diffusion 2.1 Stable Diffusion 2.1	512×512 768×768 512×512 768×768	bf16 bf16 bf16 bf16	50 50 30 30	1,976.9 6,836.3 1,186.2 4,101.8	$0.76 $0.76 $0.76 $0.76	$0.0004174 $0.0014432 $0.0002504 $0.0008659

結論

この投稿では、Inf2インスタンスを使用してStable Diffusion 2.1モデルのコンパイル、最適化、および展開について詳しく説明しました。また、SageMakerを使用したStable Diffusionモデルの展開も示しました。Inf2インスタンスはStable Diffusion 1.5に対しても優れた価格パフォーマンスを提供します。ジェネラティブAIや大規模言語モデルにおいてInf2インスタンスが優れている理由については、Amazon EC2 Inf2 Instances for Low-Cost, High-Performance Generative AI Inference are Now Generally Availableを参照してください。パフォーマンスの詳細については、Inf2 Performanceを参照してください。GitHubリポジトリで追加の例もご覧ください。

Matthew Mcclain、Beni Hegedus、Kamran Khan、Shruti Koparkar、およびQing Lanには、レビューして貴重な意見を提供していただきました。特に感謝いたします。

We will continue to update VoAGI; if you have any questions or suggestions, please contact us!