Intel CPU上での安定な拡散推論の高速化

'Intel CPUでの高速な推論の安定化'

最近、私たちは最新世代のIntel Xeon CPU（コードネームSapphire Rapids）を紹介しました。これには、ディープラーニングの高速化に対応した新しいハードウェア機能があります。また、これらを使用して自然言語処理のトランスフォーマーの分散微調整と推論を加速する方法も紹介しました。

この投稿では、Sapphire Rapids CPU上で安定拡散モデルを加速するための異なる技術を紹介します。次の投稿では、分散微調整について同様の内容を紹介します。

執筆時点では、Sapphire Rapidsサーバーにアクセスする最も簡単な方法は、Amazon EC2 R7izインスタンスファミリーを使用することです。まだプレビュー段階ですので、アクセスするためにはサインアップする必要があります。前の投稿と同様に、私はUbuntu 20.04 AMI（ami-07cd3e6c4915b2d18）を使用してr7iz.metal-16xlインスタンス（64 vCPU、512GB RAM）を使用しています。

さあ、始めましょう！コードサンプルはGitlabで利用できます。

• Substraを使用してプライバシーを保護するAIの作成
• トランスフォーマーによるグラフ分類
• AWS Inferentia2を使用してHugging Face Transformersを高速化する

Diffusersライブラリ

Diffusersライブラリは、安定拡散モデルを使用して画像を生成するのが非常に簡単です。これらのモデルに詳しくない場合は、こちらの素晴らしいイラスト入りの紹介をご覧ください。

まず、必要なライブラリ（Transformers、Diffusers、Accelerate、PyTorch）を使用して仮想環境を作成しましょう。

virtualenv sd_inference
source sd_inference/bin/activate
pip install pip --upgrade
pip install transformers diffusers accelerate torch==1.13.1

次に、推論を繰り返し実行し、単一画像生成の平均レイテンシを返すシンプルなベンチマーク関数を作成します。

import time

def elapsed_time(pipeline, prompt, nb_pass=10, num_inference_steps=20):
    # ウォームアップ
    images = pipeline(prompt, num_inference_steps=10).images
    start = time.time()
    for _ in range(nb_pass):
        _ = pipeline(prompt, num_inference_steps=num_inference_steps, output_type="np")
    end = time.time()
    return (end - start) / nb_pass

では、デフォルトのfloat32データ型を使用したStableDiffusionPipelineを構築し、推論のレイテンシを測定しましょう。

from diffusers import StableDiffusionPipeline

model_id = "runwayml/stable-diffusion-v1-5"
pipe = StableDiffusionPipeline.from_pretrained(model_id)
prompt = "sailing ship in storm by Rembrandt"
latency = elapsed_time(pipe, prompt)
print(latency)

平均レイテンシは32.3秒です。このIntel Spaceによるデモンストレーションによると、同じコードは以前の世代のIntel Xeon（コードネームIce Lake）では約45秒かかります。

コードの変更なしにSapphire Rapids CPUの方がはるかに高速であることがわかります！

さあ、加速しましょう！

Optimum IntelとOpenVINO

Optimum Intelは、Intelアーキテクチャ上でエンドツーエンドのパイプラインを高速化します。そのAPIは、バニラのDiffusers APIと非常に似ており、既存のコードを簡単に適合させることができます。

Optimum Intelは、高性能推論のためのIntelオープンソースツールキットであるOpenVINOをサポートしています。

Optimum IntelとOpenVINOは、次のようにインストールできます。

pip install optimum[openvino]

上記のコードから始めると、StableDiffusionPipelineをOVStableDiffusionPipelineに置き換えるだけです。PyTorchモデルをロードし、オンザフライでOpenVINO形式に変換するために、モデルをロードする際にexport=Trueを設定することができます。

from optimum.intel.openvino import OVStableDiffusionPipeline
...
ov_pipe = OVStableDiffusionPipeline.from_pretrained(model_id, export=True)
latency = elapsed_time(ov_pipe, prompt)
print(latency)

# エクスポートされたモデルを保存するのを忘れないでください
ov_pipe.save_pretrained("./openvino")

OpenVINOは、モデルをbfloat16形式に自動的に最適化します。そのおかげで、平均レイテンシは16.7秒となり、2倍の高速化が実現されます。

上記のパイプラインは、動的な入力形状をサポートしており、画像の数や解像度に制限はありません。Stable Diffusionでは、通常、512×512や256×256などの1つ（またはいくつか）の異なる出力解像度に制限されることが多いため、パイプラインを固定解像度に変形することで大幅な加速が可能です。複数の出力解像度が必要な場合は、各解像度に対して1つのパイプラインインスタンスを維持するだけです。

ov_pipe.reshape(batch_size=1, height=512, width=512, num_images_per_prompt=1)
latency = elapsed_time(ov_pipe, prompt)

静的な形状であれば、平均レイテンシは4.7秒に短縮され、さらに3.5倍の高速化が実現されます。

ご覧の通り、OpenVINOはStable Diffusion推論の高速化においてシンプルかつ効率的な方法です。Sapphire Rapids CPUと組み合わせることで、Ice Lake Xeons上のバニラ推論と比較してほぼ10倍の高速化が実現されます。

OpenVINOを使用することができない場合、または使用したくない場合、この記事の残りの部分では他の最適化技術の一連の手法を紹介します。シートベルトを締めてください！

システムレベルの最適化

Diffuserモデルは大きなギガバイト単位のモデルであり、画像生成はメモリを多く必要とする操作です。高性能なメモリ割り当てライブラリをインストールすることで、メモリ操作の高速化およびXeonコア間での並列化が可能になるはずです。なお、これによりシステムのデフォルトのメモリ割り当てライブラリが変更されます。もちろん、新しいライブラリをアンインストールすることでデフォルトのライブラリに戻すこともできます。

jemallocとtcmallocは同様に興味深いです。ここでは、わずかな性能向上があるため、jemallocをインストールしています。また、特定のワークロードに最適化するために調整することもできます。詳細については、チューニングガイドを参照してください。

sudo apt-get install -y libjemalloc-dev
export LD_PRELOAD=$LD_PRELOAD:/usr/lib/x86_64-linux-gnu/libjemalloc.so
export MALLOC_CONF="oversize_threshold:1,background_thread:true,metadata_thp:auto,dirty_decay_ms: 60000,muzzy_decay_ms:60000"

次に、並列処理の最適化のためにlibiompライブラリをインストールします。これはIntel OpenMP* Runtimeの一部です。

sudo apt-get install intel-mkl
export LD_PRELOAD=$LD_PRELOAD:/usr/lib/x86_64-linux-gnu/libiomp5.so
export OMP_NUM_THREADS=32

最後に、numactlコマンドラインツールをインストールします。これにより、Pythonプロセスを特定のコアにピン留めし、コンテキストスイッチに関連するオーバーヘッドを回避することができます。

numactl -C 0-31 python sd_blog_1.py

これらの最適化により、元のDiffusersコードは11.8秒で予測するようになりました。これはコードの変更なしでほぼ3倍の速さです。これらのツールは、32コアのXeon上で非常に優れた性能を発揮しています。

まだまだ終わりではありません。次に、Intel Extension for PyTorchを追加しましょう。

IPEXとBF16

Intel Extension for PyTorch（IPEX）はPyTorchを拡張し、Intel CPUのハードウェアアクセラレーション機能（AVX-512 Vector Neural Network Instructions（AVX512 VNNI）およびAdvanced Matrix Extensions（AMX））を利用します。

それをインストールしましょう。

pip install intel_extension_for_pytorch==1.13.100

次に、IPEXを使用して各パイプライン要素を最適化するためにコードを更新します（pipeオブジェクトを印刷してリストすることができます）。これには、それらをchannels-last形式に変換する必要があります。

import torch
import intel_extension_for_pytorch as ipex
...
pipe = StableDiffusionPipeline.from_pretrained(model_id)

# to channels last
pipe.unet = pipe.unet.to(memory_format=torch.channels_last)
pipe.vae = pipe.vae.to(memory_format=torch.channels_last)
pipe.text_encoder = pipe.text_encoder.to(memory_format=torch.channels_last)
pipe.safety_checker = pipe.safety_checker.to(memory_format=torch.channels_last)

# Create random input to enable JIT compilation
sample = torch.randn(2,4,64,64)
timestep = torch.rand(1)*999
encoder_hidden_status = torch.randn(2,77,768)
input_example = (sample, timestep, encoder_hidden_status)

# optimize with IPEX
pipe.unet = ipex.optimize(pipe.unet.eval(), dtype=torch.bfloat16, inplace=True, sample_input=input_example)
pipe.vae = ipex.optimize(pipe.vae.eval(), dtype=torch.bfloat16, inplace=True)
pipe.text_encoder = ipex.optimize(pipe.text_encoder.eval(), dtype=torch.bfloat16, inplace=True)
pipe.safety_checker = ipex.optimize(pipe.safety_checker.eval(), dtype=torch.bfloat16, inplace=True)

また、Sapphire Rapids CPU 上に存在する AMX タイル行列乗算ユニット (TMMU) アクセラレータを活用するために、bloat16 データ形式も有効になります。

with torch.cpu.amp.autocast(enabled=True, dtype=torch.bfloat16):
    latency = elapsed_time(pipe, prompt)
    print(latency)

この更新版では、推論の遅延時間が 11.9 秒から5.4 秒にさらに短縮されます。これは IPEX と AMX のおかげで 2 倍以上の高速化です。

さらにパフォーマンスを少し引き出すことはできますか？ はい、スケジューラを使用することで可能です！

スケジューラ

Diffusers ライブラリを使用すると、Stable Diffusion パイプラインにスケジューラを追加することができます。スケジューラは、ノイズ除去の速度と品質の最適なトレードオフを見つけようとします。

ドキュメントによると、”このドキュメントを作成した時点では、DPMSolverMultistepScheduler がおそらく最適な速度/品質のトレードオフを提供し、最低でも 20 ステップで実行できます。”

試してみましょう。

from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler
...
dpm = DPMSolverMultistepScheduler.from_pretrained(model_id, subfolder="scheduler")
pipe = StableDiffusionPipeline.from_pretrained(model_id, scheduler=dpm)

この最終版では、推論の遅延時間が5.05 秒になります。初期の Sapphire Rapids ベースライン (32.3 秒) と比較すると、これは約 6.5 倍高速です！

*環境：Amazon EC2 r7iz.metal-16xl、Ubuntu 20.04、Linux 5.15.0-1031-aws、libjemalloc-dev 5.2.1-1、intel-mkl 2020.0.166-1、PyTorch 1.13.1、Intel Extension for PyTorch 1.13.1、transformers 4.27.2、diffusers 0.14、accelerate 0.17.1、openvino 2023.0.0.dev20230217、optimum 1.7.1、optimum-intel 1.7*

結論

数秒で高品質の画像を生成できる能力は、顧客向けのアプリ、マーケティングやメディア向けのコンテンツ生成、データセットの拡張のための合成データなど、多くのユースケースに適しています。

以下は、始めるためのいくつかのリソースです：

• Diffusers ドキュメント
• Optimum Intel ドキュメント
• Intel IPEX on GitHub
• Intel と Hugging Face の開発者リソース

ご質問やフィードバックがありましたら、Hugging Face フォーラムでお待ちしています。

お読みいただき、ありがとうございました！

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