AWS Inferentia2 で Stable Diffusion のパフォーマンスを最大化し、推論コストを削減する

基盤モデルは急速に進歩しています。その中でも、Stable Diffusion モデルはテキストの入力から高品質な画像を生成する機能に優れています。Stable Diffusion を使えば、リアルな肖像画、風景画、抽象画など、さまざまな高品質な画像を作ることができます。 Stable Diffusion モデルアーキテクチャの詳細については、「Create high-quality images with Stable Diffusion models and deploy them cost-efficiently with Amazon SageMaker.」をご参照ください。

一方で、Stable Diffusion モデルを動かすには、高性能なコンピューターが必要です。数枚の画像を生成するには気になりませんが、繰り返し試行錯誤を行う際は待ち時間が発生し、動画やアニメーションを制作する時には大量の画像の生成にコストも時間もかかることは想像に難くありません。そこで、本記事では、Amazon Elastic Compte Cloud (Amazon EC2) の Amazon EC2 Inf2 インスタンスを使って Stable Diffusion モデルの推論を高性能かつ低コストに実行する方法を説明します。Inf2 インスタンスには AWS Inferentia2 という機械学習の推論向けアクセラレータチップが搭載されています。AWS Inferentia2 では Stable Diffusion のバージョン 2.1 と 1.5 の両方を低コストで実行できます。

この記事の前半では、AWS Neuron を使って Stable Diffusion モデルをコンパイルし、最適化する手順を説明します。記事の後半では、コンパイルしたモデルをAmazon SageMaker を使って Inf2 インスタンスにデプロイする方法を示します。Stable Diffusion 2.1 モデルのサイズは float32(FP32) で 5 GB、bfoat16(BF16) で 2.5 GB です。1つの inf2.xlarge インスタンスには 32GB の HBM メモリと 1 つの AWS Inferentia2 アクセラレータが搭載されているため、Stable Diffusion 2.1 モデルは 1 つの inf2.xlarge インスタンスに収まります。

Stable Diffusion モデルを Inf2 EC2 インスタンス向けにコンパイルする

はじめに、 Neuron SDK を使ってモデルのコンポーネントをコンパイルします。Neuron SDK には、ディープラーニングモデルをコンパイルして AWS Inferentia2 アクセラレータの性能を最大限に引き出すためのディープラーニングコンパイラ、ランタイム、ツールが含まれています。コンパイルする環境として、inf2.8xlarge インスタンスを AWS マネジメントコンソールにサインインして作成します。 inf2.8xlarge を使用するのは、モデルをコンパイルするのに多くのホストメモリが必要なためで、推論する際は最もコスト効率が高い inf2.xlarge インスタンスでホストできます。AWS コマンドラインインターフェイス (AWS CLI) コマンドを使用して、Neuron ライブラリを含む最新の AMI を見つけることができます。

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

この例では、ディープラーニング AMI Neuron PyTorch 1.13 (Ubuntu 20.04) を使用して EC2 インスタンスを作成しました。その後、インスタンスに接続して次のステップを実行することで、JupyterLab 環境を作成できます。

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

モデルをコンパイルしてホストするためのすべての手順が GitHub にあります。

テキストエンコーダブロックの 1 つのコンパイル手順を見てみましょう。 Stable Diffusion パイプラインの他のブロックも同様にコンパイルできます。最初のステップは、Hugging Face から事前にトレーニングされたモデルをロードすることです。

StableDiffusionPipeline.from_pretrained メソッドは、事前に学習されたモデルを pipe というパイプラインオブジェクトに読み込みます。ここでは、モデルの重みをBF16で量子化しています。次に、そのパイプラインからテキストエンコーダのコピーを作成し、効果的にクローンを作ります。次に del pipe コマンドを使って元のパイプラインオブジェクトを削除し、そのオブジェクトが使っていたメモリを解放します。

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) をコンパイルするための同様の手順が含まれています。すべてのモデルをコンパイルしたら、次の手順に従ってモデルをロードして実行できます。

コンパイルされたモデルのパスを定義します。
bfloat16 データ型を使用するように構成を指定して、事前にトレーニングされた StableDiffusionPipeline モデルを読み込みます。
torch_neuronx.DataParallel を使用して UNet モデルを 2 つの Neuron コアにロードします。これにより、データの並列推論が可能になり、モデルのパフォーマンスを大幅に向上させることができます。 UNet は、推論の中で最も計算負荷の高いコンポーネントなのでこの最適化は重要です。
モデルの残りの部分 (text_encoder、decoder、post_quant_conv) を単一の Neuron コアにロードします。

その後、プロンプトとして入力テキストを指定してパイプラインを実行できます。以下は、モデルの入出力例です。

A castle in the middle of a forest （森の真ん中にある城）

An image of A castle in the middle of a forest

コンパイルは、モデルをデプロイする前に行うか、実行時に行うことが可能です。事前にコンパイルしておくと、モデルの読み込み時間が短縮されるため、推奨されるオプションです。

Amazon SageMaker で Stable Diffusion 2.1 を AWS Inferentia2 にホストする

SageMaker の大規模モデル推論コンテナでモデルをデプロイするには次の 2 つの方法があります。

必要な構成を含む serving.properties ファイルを提供するだけのコードなしのオプション
独自の推論スクリプトを用意するオプション

両方のオプションの構成と推論スクリプト (model.py) を見ていきます。この記事では、Amazon Simple Storage Service (Amazon S3) バケットに保存されているコンパイル済みモデルを使用したデプロイについて説明します。このコンパイル済みモデルをデプロイに使用できます。

事前に用意されたスクリプトとモデルを使用してデプロイする

このセクションでは、Stable Diffusion モデルをホストするように大規模モデル推論コンテナ（LMI コンテナ）を設定する方法を示します。SD2.1 ノートブックは GitHub で入手できます。最初のステップとして、以下のディレクトリ構造に従ってモデル構成パッケージを作成します。必要なディレクトリ構造は次のとおりです。

<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 コンテナは s5cmd を使用して S3 の場所からモデルをロードするため、コンパイルされたウェイトがある場所を指定する必要があります。
option.entryPoint — 組み込みハンドラーを使用するには、transformers-neuronx クラスを指定します。カスタム推論スクリプトがある場合は、代わりにそれを指定する必要があります。
option.dtype — 特定のサイズでウェイトをロードするように指定します。この記事では、BF16 を使用しています。これにより、FP32 と比較してメモリ要件がさらに減少し、そのためレイテンシが短縮されます。
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): # which returns an object of type Outputs

付属のノートブックの重要な部分をいくつか見てみましょう。Bring Your Own Script 機能の説明がされています。

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 コアの数も指定していることに注意してください。ここでは、2 つのコアを持つ 1 つのアクセラレータにロードします。

    # 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)

プロンプトで推論を実行すると、pipe オブジェクトが呼び出されて画像が生成されます。

SageMaker モデルを作成する

SageMaker モデルを作成するために、まずは serving.properties と独自のスクリプトの model.py ファイルを圧縮し model.tar.gz を作成し、モデルアーティファクトとして Amazon S3 にアップロードします。

続いて 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 インスタンスを使用します。エンドポイント起動可能数が 0 場合は Service Quotas から引き上げることが可能です。 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",
)

Stable Diffusion on Inf2 のベンチマーク

Stable Diffusion のモデルを BF16 データ型を使用し、Inf2 上でいくつかのベンチマークテストを行いました。その結果、他のアクセラレーターと同じかそれ以上の低遅延を実現できました。これに加えて AWS Inferentia2 チップの低コストが相まって、非常に価値のある提案になっています。

以下の数値は Inf2.xl インスタンスにデプロイされた Stable Diffusion モデルからのものです。コストの詳細は Amazon EC2 Inf2 インスタンスを参照してください。

Model	Resolution	Data type	Iterations	P95 Latency (ms)	Inf2.xl On-Demand cost per hour	Inf2.xl (Cost per image)
Stable Diffusion 1.5	512×512	bf16	50	2,427.4	$0.76	$0.0005125
Stable Diffusion 1.5	768×768	bf16	50	8,235.9	$0.76	$0.0017387
Stable Diffusion 1.5	512×512	bf16	30	1,456.5	$0.76	$0.0003075
Stable Diffusion 1.5	768×768	bf16	30	4,941.6	$0.76	$0.0010432

Stable Diffusion 2.1	512×512	bf16	50	1,976.9	$0.76	$0.0004174
Stable Diffusion 2.1	768×768	bf16	50	6,836.3	$0.76	$0.0014432
Stable Diffusion 2.1	512×512	bf16	30	1,186.2	$0.76	$0.0002504
Stable Diffusion 2.1	768×768	bf16	30	4,101.8	$0.76	$0.0008659

まとめ

この記事では、Inf2 インスタンスを使用した Stable Diffusion 2.1 モデルのコンパイル、最適化、およびデプロイについて詳しく説明しました。また、SageMaker を用いた Stable Diffusion モデルの導入についてもデモを行いました。 Inf2 インスタンスは、Stable Diffusion 1.5 のコストパフォーマンスも優れています。 Inf2 インスタンスが生成 AI や大規模言語モデルに最適な理由の詳細については、「低コストで高性能な生成系 AI 推論用の Amazon EC2 Inf2 インスタンスが一般公開されました」を参照してください。パフォーマンスの詳細については、Inf2 パフォーマンスを参照してください。 AWS Trainium や AWS Inferentia などの Neuron Accelerator を利用したその他の例については、AWS Neuron の GitHub リポジトリをご覧ください。また、その他の日本語のサンプルコードについては AWS ML JP の GitHub レポジトリをご覧ください。

本記事は「Maximize Stable Diffusion performance and lower inference costs with AWS Inferentia2」をもとにソリューションアーキテクト前川が執筆しました。

Amazon Web Services ブログ