メモリ使用量の削減

Flux や Wan などの最新の拡散モデルは数十億のパラメータを持ち、推論時にハードウェア上で大量のメモリを消費します。一般的な GPU には十分なメモリがないことが多いため、これは課題となります。メモリ制限を克服するには、複数の GPU を使用したり（利用可能な場合）、パイプラインの一部を CPU にオフロードしたり、その他の方法があります。

このガイドでは、メモリ使用量を削減する方法について説明します。

これらの手法は、モデルによって調整が必要な場合があります。例えば、Transformerベースの拡散モデルは、UNetベースのモデルと比べて、これらのメモリ最適化から同程度の恩恵を受けられない可能性があります。

複数のGPU

複数のGPUにアクセスできる場合、大規模なモデルをハードウェア全体に効率的にロードして分散するためのいくつかのオプションがあります。これらの機能はAccelerateライブラリによってサポートされているため、まずインストールされていることを確認してください。

pip install -U accelerate

シャード化されたチェックポイント

複数のシャードに分割された大きなチェックポイントをロードすることは有用です。シャードは1つずつロードされるため、メモリ使用量を低く抑えられ、モデルサイズと最大のシャードサイズに必要なメモリだけで済みます。fp32のチェックポイントが5GBを超える場合は、シャード化することを推奨します。デフォルトのシャードサイズは5GBです。

max_shard_size パラメータを指定して、save_pretrained() でチェックポイントをシャード化します。

from diffusers import AutoModel

unet = AutoModel.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0", subfolder="unet"
)
unet.save_pretrained("sdxl-unet-sharded", max_shard_size="5GB")

これにより、通常のチェックポイントの代わりにシャード化されたチェックポイントを使用することで、メモリを節約できます。

import torch
from diffusers import AutoModel, StableDiffusionXLPipeline

unet = AutoModel.from_pretrained(
    "username/sdxl-unet-sharded", dtype=torch.float16
)
pipeline = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0",
    unet=unet,
    dtype=torch.float16
).to("cuda")  # or "mps", "xpu", "cpu"

デバイスの配置

デバイス配置は実験的な機能であり、APIが変更される可能性があります。現時点では、balanced戦略のみがサポートされています。今後、追加のマッピング戦略をサポートする予定です。

device_map パラメータは、パイプライン内のモデルコンポーネント、または個別モデルのレイヤーがデバイス間でどのように分散されるかを制御します。

パイプラインレベル

モデルレベル

独自のdevice_mapを設計する際は、モデルの特定のモジュール名またはレイヤーと、デバイス識別子（GPUの場合は整数、CPUの場合はcpu、ディスクの場合はdisk）の辞書である必要があります。

モデルに対してhf_device_mapを呼び出して、モデルのレイヤーがどのように分散されているかを確認し、その後、独自の分散設計を行ってください。

print(transformer.hf_device_map)
{'pos_embed': 0, 'time_text_embed': 0, 'context_embedder': 0, 'x_embedder': 0, 'transformer_blocks': 0, 'single_transformer_blocks.0': 0, 'single_transformer_blocks.1': 0, 'single_transformer_blocks.2': 0, 'single_transformer_blocks.3': 0, 'single_transformer_blocks.4': 0, 'single_transformer_blocks.5': 0, 'single_transformer_blocks.6': 0, 'single_transformer_blocks.7': 0, 'single_transformer_blocks.8': 0, 'single_transformer_blocks.9': 0, 'single_transformer_blocks.10': 'cpu', 'single_transformer_blocks.11': 'cpu', 'single_transformer_blocks.12': 'cpu', 'single_transformer_blocks.13': 'cpu', 'single_transformer_blocks.14': 'cpu', 'single_transformer_blocks.15': 'cpu', 'single_transformer_blocks.16': 'cpu', 'single_transformer_blocks.17': 'cpu', 'single_transformer_blocks.18': 'cpu', 'single_transformer_blocks.19': 'cpu', 'single_transformer_blocks.20': 'cpu', 'single_transformer_blocks.21': 'cpu', 'single_transformer_blocks.22': 'cpu', 'single_transformer_blocks.23': 'cpu', 'single_transformer_blocks.24': 'cpu', 'single_transformer_blocks.25': 'cpu', 'single_transformer_blocks.26': 'cpu', 'single_transformer_blocks.27': 'cpu', 'single_transformer_blocks.28': 'cpu', 'single_transformer_blocks.29': 'cpu', 'single_transformer_blocks.30': 'cpu', 'single_transformer_blocks.31': 'cpu', 'single_transformer_blocks.32': 'cpu', 'single_transformer_blocks.33': 'cpu', 'single_transformer_blocks.34': 'cpu', 'single_transformer_blocks.35': 'cpu', 'single_transformer_blocks.36': 'cpu', 'single_transformer_blocks.37': 'cpu', 'norm_out': 'cpu', 'proj_out': 'cpu'}

例えば、以下のdevice_mapでは、single_transformer_blocks.10からsingle_transformer_blocks.20までを2番目のGPU(1)に配置します。

import torch
from diffusers import AutoModel

device_map = {
    'pos_embed': 0, 'time_text_embed': 0, 'context_embedder': 0, 'x_embedder': 0, 'transformer_blocks': 0, 'single_transformer_blocks.0': 0, 'single_transformer_blocks.1': 0, 'single_transformer_blocks.2': 0, 'single_transformer_blocks.3': 0, 'single_transformer_blocks.4': 0, 'single_transformer_blocks.5': 0, 'single_transformer_blocks.6': 0, 'single_transformer_blocks.7': 0, 'single_transformer_blocks.8': 0, 'single_transformer_blocks.9': 0, 'single_transformer_blocks.10': 1, 'single_transformer_blocks.11': 1, 'single_transformer_blocks.12': 1, 'single_transformer_blocks.13': 1, 'single_transformer_blocks.14': 1, 'single_transformer_blocks.15': 1, 'single_transformer_blocks.16': 1, 'single_transformer_blocks.17': 1, 'single_transformer_blocks.18': 1, 'single_transformer_blocks.19': 1, 'single_transformer_blocks.20': 1, 'single_transformer_blocks.21': 'cpu', 'single_transformer_blocks.22': 'cpu', 'single_transformer_blocks.23': 'cpu', 'single_transformer_blocks.24': 'cpu', 'single_transformer_blocks.25': 'cpu', 'single_transformer_blocks.26': 'cpu', 'single_transformer_blocks.27': 'cpu', 'single_transformer_blocks.28': 'cpu', 'single_transformer_blocks.29': 'cpu', 'single_transformer_blocks.30': 'cpu', 'single_transformer_blocks.31': 'cpu', 'single_transformer_blocks.32': 'cpu', 'single_transformer_blocks.33': 'cpu', 'single_transformer_blocks.34': 'cpu', 'single_transformer_blocks.35': 'cpu', 'single_transformer_blocks.36': 'cpu', 'single_transformer_blocks.37': 'cpu', 'norm_out': 'cpu', 'proj_out': 'cpu'
}

transformer = AutoModel.from_pretrained(
    "black-forest-labs/FLUX.1-dev", 
    subfolder="transformer",
    device_map=device_map,
    dtype=torch.bfloat16
)

各デバイスの最大メモリ使用量をマッピングする辞書を渡すことで、上限を強制できます。max_memory に含まれていないデバイスは無視され、パイプラインのコンポーネントはそこに分配されません。

import torch
from diffusers import AutoModel, StableDiffusionXLPipeline

max_memory = {0:"1GB", 1:"1GB"}
pipeline = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0",
    dtype=torch.float16,
    device_map="balanced",
    max_memory=max_memory
)

Diffusers はデフォルトで全デバイスの最大メモリを使用しますが、GPU に収まらない場合は、以下の方法を使用して単一の GPU を使用し、CPU へオフロードする必要があります。

enable_model_cpu_offload() は単一の GPU でのみ動作しますが、非常に大きなモデルは 1 つの GPU に収まらない場合があります。
enable_sequential_cpu_offload() は動作する場合がありますが、非常に低速であり、かつ単一の GPU に限定されます。

device_mapをリセットするには、reset_device_map()メソッドを使用します。デバイスマップが設定されたパイプラインに対して.to()、enable_sequential_cpu_offload()、enable_model_cpu_offload()などのメソッドを使用する場合は、この処理が必要です。

pipeline.reset_device_map()

VAEスライシング

VAE slicingは、大きなバッチの入力を単一のデータバッチに分割して個別に処理することでメモリを節約します。この手法は、一度に複数の画像を生成する場合に最も効果的です。

例えば、一度に4枚の画像を生成する場合、デコードによってピーク時のアクティベーションメモリが4倍に増加します。VAEスライシングでは、4枚すべてを一度にデコードするのではなく、1枚ずつデコードすることで、この増加を抑制します。

スライスされた VAE を有効にするには、enable_slicing() を呼び出します。複数画像のバッチをデコードする際にパフォーマンスがわずかに向上し、単一画像のバッチではパフォーマンスへの影響はありません。

import torch
from diffusers import AutoModel, StableDiffusionXLPipeline

pipeline = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0",
    dtype=torch.float16,
).to("cuda")  # or "mps", "xpu", "cpu"
pipeline.vae.enable_slicing()
pipeline(["An astronaut riding a horse on Mars"]*32).images[0]
print(f"Max memory reserved: {torch.cuda.max_memory_allocated() / 1024**3:.2f} GB")

AutoencoderKLWan クラスと AsymmetricAutoencoderKL クラスは、スライシングをサポートしていません。

VAEのタイル分割

VAE tilingは、画像全体を一度に処理するのではなく、より小さな重複するタイルに分割して処理することでメモリを節約します。これにより、GPUが一度に1つのタイルのみを処理するため、ピーク時のメモリ使用量も削減されます。

enable_tiling() を呼び出して VAE のタイル分割を有効にします。タイルごとに個別にデコードされるため、生成された画像にはタイル間でわずかなトーンのばらつきが生じることがありますが、タイルの境界に明らかな継ぎ目が出ることはありません。解像度が事前に指定された（ただし設定可能）閾値を下回る場合は、タイル分割は無効になります。例えば、StableDiffusionPipeline の VAE では、この閾値は 512x512 です。

import torch
from diffusers import AutoPipelineForImage2Image
from diffusers.utils import load_image

pipeline = AutoPipelineForImage2Image.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0", dtype=torch.float16
).to("cuda")  # or "mps", "xpu", "cpu"
pipeline.vae.enable_tiling()

init_image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/img2img-sdxl-init.png")
prompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"
pipeline(prompt, image=init_image, strength=0.5).images[0]
print(f"Max memory reserved: {torch.cuda.max_memory_allocated() / 1024**3:.2f} GB")

AutoencoderKLWan と AsymmetricAutoencoderKL は、タイル分割をサポートしていません。

GPU上のVAEメモリを完全にスキップするには、Remote inferenceでリモートエンドポイント上でエンコードとデコードを実行します。

オフローディング

オフローディング戦略では、現在アクティブでないレイヤーやモデルをCPUに移動し、GPUメモリを増加させないようにします。これらの戦略は、量子化やtorch.compileと組み合わせることで、推論速度とメモリ使用量のバランスを取ることができます。

詳細については、Compile and offloading quantized modelsガイドを参照してください。

CPUオフロード

CPU offloadingは、重みをGPUからCPUへ選択的に移動します。コンポーネントが必要になった場合はGPUに転送され、不要になった場合はCPUに移動します。この手法は、モデル全体ではなくサブモジュール単位で動作します。モデル全体をGPUに保持することを避けることで、メモリを節約します。

CPU offloadingはメモリ使用量を大幅に削減しますが、サブモジュールがデバイス間で何度もやり取りされるため、非常に遅いです。その遅さゆえに、実用的でない場合もよくあります。

enable_sequential_cpu_offload()を呼び出す前に、パイプラインをCUDAに移動しないでください。そうしないと、メモリ削減の効果がほとんど得られません（詳細はこのissueを参照してください）。これはステートフルな操作であり、モデルにフックをインストールします。

パイプラインでこれを有効にするには、enable_sequential_cpu_offload()を呼び出します。

import torch
from diffusers import DiffusionPipeline

pipeline = DiffusionPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-schnell", dtype=torch.bfloat16
)
pipeline.enable_sequential_cpu_offload()

pipeline(
    prompt="An astronaut riding a horse on Mars",
    guidance_scale=0.,
    height=768,
    width=1360,
    num_inference_steps=4,
    max_sequence_length=256,
).images[0]
print(f"Max memory reserved: {torch.cuda.max_memory_allocated() / 1024**3:.2f} GB")

モデルのオフロード

Model offloadingでは、一部のレイヤーやモデルコンポーネントを選択的に移動するのではなく、モデル全体をGPUに移動します。通常、テキストエンコーダー、UNet、VAEといった主要なパイプラインモデルの1つがGPUに配置され、その他のコンポーネントはCPU上に保持されます。UNetのように複数回実行されるコンポーネントは、完全に処理が完了して不要になるまでGPU上に保持されます。これにより、CPU offloadingに伴う通信オーバーヘッドが排除され、Model offloadingはより高速な代替手段となります。その代わりとして、メモリ削減の幅はCPU offloadingほど大きくはなりません。

フックがインストールされた後、パイプラインの外でモデルが再利用される場合（詳細はフックの削除を参照）、モデルを適切にオフロードするには、パイプライン全体とモデルを想定された順序で実行する必要があります。これはモデルにフックをインストールする状態付きの操作です。

パイプラインでこれを有効にするには、enable_model_cpu_offload()を呼び出します。

import torch
from diffusers import DiffusionPipeline

pipeline = DiffusionPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-schnell", dtype=torch.bfloat16
)
pipeline.enable_model_cpu_offload()

pipeline(
    prompt="An astronaut riding a horse on Mars",
    guidance_scale=0.,
    height=768,
    width=1360,
    num_inference_steps=4,
    max_sequence_length=256,
).images[0]
print(f"Max memory reserved: {torch.cuda.max_memory_allocated() / 1024**3:.2f} GB")

enable_model_cpu_offload() は、encode_prompt() メソッドを単独で使用してテキストエンコーダーの隠れ状態を生成する場合にも役立ちます。

グループオフローディング

Group offloadingは、内部レイヤーのグループ（torch.nn.ModuleListまたはtorch.nn.Sequential）をCPUに移動します。model offloadingよりもメモリ使用量が少なく、通信オーバーヘッドを削減できるため、CPU offloadingよりも高速です。

グループオフローディングは、フォワード実装が重みに依存する入力のデバイスキャストを含む場合、グループオフローディングのデバイスキャスト機構と競合する可能性があるため、すべてのモデルで動作しない場合があります。

offload_type パラメータを block_level または leaf_level に設定することで、グループオフロードを有効にできます。

block_level は、num_blocks_per_group パラメータに基づいてレイヤーのグループをオフロードします。例えば、40 レイヤーのモデルで num_blocks_per_group=2 の場合、一度に 2 レイヤーがオンロードおよびオフロードされます（オンロード/オフロードの総数は 20 回）。これにより、メモリ要件が大幅に削減されます。
leaf_level は最下位レベルで個々のレイヤーをオフロードし、CPU オフロードと同等です。ただし、推論速度を犠牲にすることなくストリームを使用することで、より高速にできます。

グループオフロードは、パイプライン全体または個々のモデルに対して適用できます。パイプライン全体にグループオフロードを適用するのが最も簡単ですが、個々のモデルに対して選択的に適用することで、ユーザーはモデルごとに異なるオフロード手法を使用する柔軟性を得られます。

pipeline

model

CUDA ストリーム

CUDA デバイスで非同期データ転送ストリームがサポートされている場合、use_stream パラメータを有効にすることで、CPU オフロードと比較して全体の実行時間を短縮できます。レイヤープリフェッチを使用して、データ転送と計算をオーバーラップさせます。現在のレイヤーの実行中に、次に実行するレイヤーを GPU に読み込みます。CPU メモリが大幅に増加する可能性があるため、モデルサイズの 2 倍のメモリを確保してください。

メモリ使用量がわずかに増加する代わりに、より大きな高速化を得るには、record_stream=True を設定してください。詳細については、torch.Tensor.record_stream のドキュメントを参照してください。

タイル分割を有効にした VAE で use_stream=True を使用する場合、デバイス不一致エラーを回避するために、推論の前にダミーのフォワードパス（ダミー入力でも可）を実行してください。すべての実装で動作するとは限らないため、問題が発生した場合は、お気軽に issue を作成してください。

block_level グループオフロードを use_stream を有効にして使用している場合、num_blocks_per_group パラメータは 1 に設定する必要があります。そうしないと、警告が発生します。

pipeline.transformer.enable_group_offload(onload_device=onload_device, offload_device=offload_device, offload_type="leaf_level", use_stream=True, record_stream=True)

low_cpu_mem_usage パラメータを True に設定すると、グループオフローディング時にストリームを使用している場合の CPU メモリ使用量を削減できます。leaf_level オフローディングや、CPU メモリがボトルネックになっている場合に特に有効です。ピン留めされたテンソルを事前に作成するのではなく、必要に応じてピン留めされたテンソルを動的に作成することでメモリを節約します。ただし、これにより全体の実行時間が長くなる可能性があります。

ディスクへのオフロード

Group offloading はモデルサイズによって、システムメモリを大量に消費する可能性があります。メモリが限られているシステムでは、セカンダリメモリとしてディスクへの group offloading を試してください。

モデルをディスクにオフロードするには、enable_group_offload() または apply_group_offloading() のいずれかで offload_to_disk_path 引数を設定します。

pipeline.transformer.enable_group_offload(onload_device=onload_device, offload_device=offload_device, offload_type="leaf_level", offload_to_disk_path="path/to/disk")

apply_group_offloading(pipeline.text_encoder, onload_device=onload_device, offload_type="block_level", num_blocks_per_group=2, offload_to_disk_path="path/to/disk")

速度とメモリ使用量のトレードオフを比較するには、これらの2つのテーブルを参照してください。

レイヤーごとのキャスト

レイヤーごとのキャストをグループオフローディングと組み合わせることで、さらにメモリを節約できます。

Layerwise castingは、重みをより小さいデータ形式（例えば、torch.float8_e4m3fnやtorch.float8_e5m2）で保存することでメモリ使用量を削減し、計算時にはtorch.float16やtorch.bfloat16などの高精度な形式にアップキャストします。正規化やモジュレーションに関連する重みを含む特定のレイヤーは、fp8で保存すると生成品質が低下する可能性があるため、スキップされます。

レイヤーワイズキャストは、フォワード実装が重みの内部型キャストを含む場合、すべてのモデルで動作しない可能性があります。現在のレイヤーワイズキャストの実装は、フォワードパスが重みの精度に依存しないこと、および入力データ型が常にcompute_dtypeで指定されることを前提としています（非互換な実装の例はこちらを参照）。

PEFTレイヤーを含むカスタムモデリング実装では、レイヤーごとのキャストが失敗する可能性があります。いくつかのチェックは用意されていますが、これらは十分にテストされておらず、すべてのケースで動作が保証されるものではありません。

enable_layerwise_casting()を呼び出して、ストレージと計算のデータ型を設定します。

import torch
from diffusers import CogVideoXPipeline, CogVideoXTransformer3DModel
from diffusers.utils import export_to_video

transformer = CogVideoXTransformer3DModel.from_pretrained(
    "THUDM/CogVideoX-5b",
    subfolder="transformer",
    dtype=torch.bfloat16
)
transformer.enable_layerwise_casting(storage_dtype=torch.float8_e4m3fn, compute_dtype=torch.bfloat16)

pipeline = CogVideoXPipeline.from_pretrained("THUDM/CogVideoX-5b",
    transformer=transformer,
    dtype=torch.bfloat16
).to("cuda")  # or "mps", "xpu", "cpu"
prompt = (
    "A panda, dressed in a small, red jacket and a tiny hat, sits on a wooden stool in a serene bamboo forest. "
    "The panda's fluffy paws strum a miniature acoustic guitar, producing soft, melodic tunes. Nearby, a few other "
    "pandas gather, watching curiously and some clapping in rhythm. Sunlight filters through the tall bamboo, "
    "casting a gentle glow on the scene. The panda's face is expressive, showing concentration and joy as it plays. "
    "The background includes a small, flowing stream and vibrant green foliage, enhancing the peaceful and magical "
    "atmosphere of this unique musical performance."
)
video = pipeline(prompt=prompt, guidance_scale=6, num_inference_steps=50).frames[0]
print(f"Max memory reserved: {torch.cuda.max_memory_allocated() / 1024**3:.2f} GB")
export_to_video(video, "output.mp4", fps=8)

より細かな制御と柔軟性が必要な場合は、apply_layerwise_casting() メソッドも使用できます。特定の内部モジュールに対して呼び出すことで、モデルのレイヤーに部分的に適用できます。正規化レイヤーやモジュレーションレイヤーなど、回避するモジュールを指定するには、skip_modules_pattern または skip_modules_classes パラメータを使用します。

import torch
from diffusers import CogVideoXTransformer3DModel
from diffusers.hooks import apply_layerwise_casting

transformer = CogVideoXTransformer3DModel.from_pretrained(
    "THUDM/CogVideoX-5b",
    subfolder="transformer",
    dtype=torch.bfloat16
)

# skip the normalization layer
apply_layerwise_casting(
    transformer,
    storage_dtype=torch.float8_e4m3fn,
    compute_dtype=torch.bfloat16,
    skip_modules_classes=["norm"],
    non_blocking=True,
)

torch.channels_last

torch.channels_last は、テンソルの保存形式を (batch size, channels, height, width) から (batch size, height, width, channels) に反転します。これにより、テンソルはメモリに保存されたテンソルにハードウェアがシーケンシャルにアクセスする方式と整合し、メモリ内でピクセル値にアクセスするために前後に飛び回ることを回避できます。

現在、すべてのオペレーターが channels-last 形式をサポートしているわけではなく、パフォーマンスが最悪になる場合もありますが、試してみる価値はあります。

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

メモリ効率の高いアテンション

Diffusers は set_attention_backend() を通じて、複数のメモリ効率に優れたアテンションバックエンド（FlashAttention、xFormers、SageAttention など）をサポートしています。切り替え方については、Attention backends ガイドを参照してください。

Update on GitHub