HuggingFace Diffusers 0.12 : 最適化 : メモリと速度 (翻訳/解説)

翻訳 : (株)クラスキャットセールスインフォメーション
作成日時 : 02/16/2023 (v0.12.1)

* 本ページは、HuggingFace Diffusers の以下のドキュメントを翻訳した上で適宜、補足説明したものです：

Optimization / Special Hardware : Memory and speed

* サンプルコードの動作確認はしておりますが、必要な場合には適宜、追加改変しています。
* ご自由にリンクを張って頂いてかまいませんが、sales-info@classcat.com までご一報いただけると嬉しいです。

クラスキャット人工知能研究開発支援サービス

◆ クラスキャットは人工知能・テレワークに関する各種サービスを提供しています。お気軽にご相談ください :

人工知能研究開発支援
1. 人工知能研修サービス(経営者層向けオンサイト研修)
2. テクニカルコンサルティングサービス
3. 実証実験(プロトタイプ構築)
4. アプリケーションへの実装
人工知能研修サービス
PoC(概念実証)を失敗させないための支援

◆ 人工知能とビジネスをテーマに WEB セミナーを定期的に開催しています。スケジュール。

お住まいの地域に関係なく Web ブラウザからご参加頂けます。事前登録 が必要ですのでご注意ください。

◆ お問合せ : 本件に関するお問い合わせ先は下記までお願いいたします。

株式会社クラスキャット セールス・マーケティング本部セールス・インフォメーション
sales-info@classcat.com ; Web: www.classcat.com ; ClassCatJP

HuggingFace Diffusers 0.12 : 最適化 : メモリと速度

メモリとスピードのために 🤗 Diffusers 推論を最適化する幾つかのテクニックとアイデアを紹介します。原則的には、メモリ効率アテンションのために xFormers の使用を勧めます、推奨されるインストール手順をご覧ください。

以下の設定がパフォーマンスとメモリにどのように影響するかを説明します。

obtained on NVIDIA TITAN RTX by generating a single image of size 512×512 from the prompt “a photo of an astronaut riding a horse on mars” with 50 DDIM steps.

cuDNN 自動チューナーを有効にする

NVIDIA cuDNN は畳み込みを計算するための多くのアルゴリズムをサポートしています。Autotuner は短いベンチマークを実行して、与えられたハードウェア上で与えられた入力サイズに対して最善の性能を持つカーネルを選択します。

畳み込みネットワークを使用していますので (他のタイプは現在サポートされていません)、以下を設定することにより推論を起動する前に cuDNN autotuner を有効にすることができます :

import torch

torch.backends.cudnn.benchmark = True

fp32 の代わりに tf32 を使用する (Ampere とそれ以後の CUDA デバイス上で)

Ampere とそれ以後の CUDA デバイスでは、行列乗算と畳込みは高速化のために TensorFloat32 (TF32) モードが使用できますが、計算の正確さは僅かに下がります。PyTorch は畳み込みのために TF32 を有効にしますが、行列乗算に対してはそうではありません、そしてネットワークが全 float32 精度を必要としない限りは、行列演算に対してもこの設定を有効にすることを勧めます。それは通常は数値精度の無視できるほどの損失を伴うものの計算を大幅にスピードアップすることができます。それについてここで更に読むことができます。All you need to do is to add this before your inference:

import torch

torch.backends.cuda.matmul.allow_tf32 = True

半精度重み

より多くの GPU メモリを節約して更にスピードを得るため、モデル重みを直接半精度でロードして実行することができます。これは fp16 という名前のブランチにセーブされた、float16 版の重みをロードして、それらをロードするときに PyTorch に float16 型を使用するように伝えることによります :

pipe = StableDiffusionPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5",
    
    torch_dtype=torch.float16,
)
pipe = pipe.to("cuda")

prompt = "a photo of an astronaut riding a horse on mars"
image = pipe(prompt).images[0]

It is strongly discouraged to make use of [`torch.autocast`]( https://pytorch.org/docs/stable/amp.html#torch.autocast ) in any of the pipelines as it can lead to black images and is always slower than using pure float16 precision.

更なるメモリ節約のための sliced アテンション

更ならメモリ節約のために、アテンションの sliced 版を利用できます、これは計算のすべてを一度に実行する代わりにステップ毎に行います。

アテンション・スライシングは単に 1 のバッチサイズが使用された場合でさえ、モデルが一つより多いアテンションヘッドを使用している限りは有用です。1 つより多いアテンションヘッドがある場合、*QK^T* アテンション行列は各ヘッドに対してシーケンシャルに計算できて、これはメモリのかなりの量を節約できます。

各ヘッドに対してアテンション計算をシーケンシャルに実行するためには、推論の前にパイプラインで enable_attention_slicing() を呼び出す必要があるだけです、ここでのようにです :

import torch
from diffusers import StableDiffusionPipeline

pipe = StableDiffusionPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5",
    
    torch_dtype=torch.float16,
)
pipe = pipe.to("cuda")

prompt = "a photo of an astronaut riding a horse on mars"
pipe.enable_attention_slicing()
image = pipe(prompt).images[0]

約 10% 遅い推論時間という小さなパフォーマンス・ペナルティはありますが、この方法は Stable Diffusion を 3.2 GB ほどに少ない VRAM で利用することを可能にします！

大規模なバッチのための Sliced VAE デコード

限られた VRAM で画像の大規模バッチをデコードしたり、32 以上の画像を持つバッチを有効にするには、sliced VAE デコードを使用できます、これはバッチ潜在変数 (batch latents) を一度に 1 画像デコードできます。

メモリ使用量を更に最小化するため enable_attention_slicing() や enable_xformers_memory_efficient_attention() とともにこれを組み合わせるを望むかもしれません。

一度に 1 画像の VAE デコードを実行するには、推論の前にパイプラインで enable_vae_slicing() を呼び出します。例えば :

import torch
from diffusers import StableDiffusionPipeline

pipe = StableDiffusionPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5",
    
    torch_dtype=torch.float16,
)
pipe = pipe.to("cuda")

prompt = "a photo of an astronaut riding a horse on mars"
pipe.enable_vae_slicing()
images = pipe([prompt] * 32).images

複数画像バッチ上では VAE デコードで小さい性能ブーストを見る可能性があります。単一画像バッチではパフォーマンスへの影響はないはずです。

メモリ節約のために accelerate で CPU にオフロードする

更なるメモリ節約のために、重みを CPU にオフロードして forward パスを実行するときだけそれらを GPU にロードすることが可能です。

CPU オフロードを実行するには、必要なことは enable_sequential_cpu_offload() を呼び出すことだけです :

import torch
from diffusers import StableDiffusionPipeline

pipe = StableDiffusionPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5",
    
    torch_dtype=torch.float16,
)

prompt = "a photo of an astronaut riding a horse on mars"
pipe.enable_sequential_cpu_offload()
image = pipe(prompt).images[0]

するとメモリ消費量を 3GB 未満にすることができます。

最小限のメモリ消費 (< 2GB) のためにアテンション slicing とそれを連鎖することも可能です。

import torch
from diffusers import StableDiffusionPipeline

pipe = StableDiffusionPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5",
    
    torch_dtype=torch.float16,
)

prompt = "a photo of an astronaut riding a horse on mars"
pipe.enable_sequential_cpu_offload()
pipe.enable_attention_slicing(1)

image = pipe(prompt).images[0]

Note : enable_sequential_cpu_offload() を使用するとき、パイプラインを前もって CUDA に移動しないことが重要です、そうでないならばメモリ消費のゲインは最低限にしかなりません。詳細はこの issue をご覧ください。

チャネル last メモリ形式の使用

チャネル last メモリ形式は、次元順序を保持したままメモリで NCHW テンソルを順序付ける別の方法です。チャネル last テンソルは、チャネルが最も密な次元になるような方法で順序付けられます (aka 画像をピクセル単位でストアします)。現在すべての演算子がチャネル last 形式をサポートするわけではないので、最悪のパフォーマンスという結果になる可能性がありますので、それを試してモデルに対して動作するかを確認するのが良いです。

例えば、パイプラインの UNet モデルがチャネル last 形式を使用するように設定するには、以下を使用できます :

print(pipe.unet.conv_out.state_dict()["weight"].stride())  # (2880, 9, 3, 1)
pipe.unet.to(memory_format=torch.channels_last)  # in-place operation
print(
    pipe.unet.conv_out.state_dict()["weight"].stride()
)  # (2880, 1, 960, 320) having a stride of 1 for the 2nd dimension proves that it works

トレーシング

Tracing はサンプル入力テンソルをモデルを通して実行し、その入力がモデルの層を進むときに呼び出される演算を捕捉し、その結果 just-in-time コンパイルを使用して最適化された executable や ScriptFunction が返されるようにします。

UNet モデルをトレースするには、以下を利用できます :

import time
import torch
from diffusers import StableDiffusionPipeline
import functools

# torch disable grad
torch.set_grad_enabled(False)

# set variables
n_experiments = 2
unet_runs_per_experiment = 50

# load inputs
def generate_inputs():
    sample = torch.randn(2, 4, 64, 64).half().cuda()
    timestep = torch.rand(1).half().cuda() * 999
    encoder_hidden_states = torch.randn(2, 77, 768).half().cuda()
    return sample, timestep, encoder_hidden_states


pipe = StableDiffusionPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5",
    torch_dtype=torch.float16,
).to("cuda")
unet = pipe.unet
unet.eval()
unet.to(memory_format=torch.channels_last)  # use channels_last memory format
unet.forward = functools.partial(unet.forward, return_dict=False)  # set return_dict=False as default

# warmup
for _ in range(3):
    with torch.inference_mode():
        inputs = generate_inputs()
        orig_output = unet(*inputs)

# trace
print("tracing..")
unet_traced = torch.jit.trace(unet, inputs)
unet_traced.eval()
print("done tracing")


# warmup and optimize graph
for _ in range(5):
    with torch.inference_mode():
        inputs = generate_inputs()
        orig_output = unet_traced(*inputs)


# benchmarking
with torch.inference_mode():
    for _ in range(n_experiments):
        torch.cuda.synchronize()
        start_time = time.time()
        for _ in range(unet_runs_per_experiment):
            orig_output = unet_traced(*inputs)
        torch.cuda.synchronize()
        print(f"unet traced inference took {time.time() - start_time:.2f} seconds")
    for _ in range(n_experiments):
        torch.cuda.synchronize()
        start_time = time.time()
        for _ in range(unet_runs_per_experiment):
            orig_output = unet(*inputs)
        torch.cuda.synchronize()
        print(f"unet inference took {time.time() - start_time:.2f} seconds")

# save the model
unet_traced.save("unet_traced.pt")

Then we can replace the unet attribute of the pipeline with the traced model like the following

from diffusers import StableDiffusionPipeline
import torch
from dataclasses import dataclass


@dataclass
class UNet2DConditionOutput:
    sample: torch.FloatTensor


pipe = StableDiffusionPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5",
    torch_dtype=torch.float16,
).to("cuda")

# use jitted unet
unet_traced = torch.jit.load("unet_traced.pt")
# del pipe.unet
class TracedUNet(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.in_channels = pipe.unet.in_channels
        self.device = pipe.unet.device

    def forward(self, latent_model_input, t, encoder_hidden_states):
        sample = unet_traced(latent_model_input, t, encoder_hidden_states)[0]
        return UNet2DConditionOutput(sample=sample)


pipe.unet = TracedUNet()

with torch.inference_mode():
    image = pipe([prompt] * 1, num_inference_steps=50).images[0]

メモリ効率的アテンション

アテンション・ブロックの処理能力 (bandwidth) の最適化における最近のワークは大幅なスピードアップと GPU メモリ使用量のゲインを生み出しました。最も最近のものは @tridao の Flash アテンションです : コード, 論文

バッチサイズ 1 (1 プロンプト) の 512×512 で推論を実行したときの幾つかの Nvidia GPU 上で獲得したスピードアップがここで示されます :

To leverage it just make sure you have:

PyTorch > 1.12
Cuda available
Installed the xformers library.

from diffusers import StableDiffusionPipeline
import torch

pipe = StableDiffusionPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5",
    torch_dtype=torch.float16,
).to("cuda")

pipe.enable_xformers_memory_efficient_attention()

with torch.inference_mode():
    sample = pipe("a small cat")

# optional: You can disable it via
# pipe.disable_xformers_memory_efficient_attention()

以上

月	火	水	木	金	土	日
		1	2	3	4	5
6	7	8	9	10	11	12
13	14	15	16	17	18	19
20	21	22	23	24	25	26
27	28