Kurangi penggunaan memori

Model difusi modern seperti Flux dan Wan memiliki miliaran parameter yang memakan banyak memori pada perangkat keras Anda untuk inferensi. Hal ini menantang karena GPU umum sering kali tidak memiliki memori yang cukup. Untuk mengatasi keterbatasan memori, Anda dapat menggunakan lebih dari satu GPU (jika tersedia), memindahkan beberapa komponen pipeline ke CPU, dan lain-lain.

Panduan ini akan menunjukkan cara mengurangi penggunaan memori Anda.

Perhatikan bahwa teknik-teknik ini mungkin perlu disesuaikan tergantung pada model. Misalnya, model difusi berbasis transformer mungkin tidak mendapatkan manfaat yang sama dari optimasi memori ini seperti model berbasis UNet.

Beberapa GPU

Jika Anda memiliki akses ke lebih dari satu GPU, ada beberapa opsi untuk memuat dan mendistribusikan model besar secara efisien di perangkat keras Anda. Fitur-fitur ini didukung oleh pustaka Accelerate, jadi pastikan pustaka tersebut terinstal terlebih dahulu.

pip install -U accelerate

Checkpoint terbagi

Memuat checkpoint besar dalam beberapa shard berguna karena shard dimuat satu per satu. Hal ini menjaga penggunaan memori tetap rendah, hanya membutuhkan memori sebesar ukuran model dan ukuran shard terbesar. Kami merekomendasikan sharding ketika checkpoint fp32 lebih dari 5GB. Ukuran shard default adalah 5GB.

Shard checkpoint di save_pretrained() dengan parameter max_shard_size.

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

Sekarang Anda dapat menggunakan checkpoint terbagi, alih-alih checkpoint reguler, untuk menghemat memori.

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"

Penempatan perangkat

Penempatan perangkat adalah fitur eksperimental dan API-nya dapat berubah. Saat ini, hanya strategi balanced yang didukung. Kami berencana untuk mendukung strategi pemetaan tambahan di masa depan.

Parameter device_map mengontrol bagaimana komponen model dalam pipeline atau lapisan dalam model individual didistribusikan di antara perangkat.

level pipeline

level model

Saat merancang device_map Anda sendiri, ia harus berupa kamus yang memetakan nama modul atau layer spesifik dari model ke pengidentifikasi perangkat (bilangan bulat untuk GPU, cpu untuk CPU, dan disk untuk disk).

Panggil hf_device_map pada model untuk melihat bagaimana lapisan model didistribusikan, lalu rancang distribusi Anda sendiri.

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'}

Sebagai contoh, device_map di bawah menempatkan single_transformer_blocks.10 hingga single_transformer_blocks.20 pada GPU kedua (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
)

Lewatkan dictionary yang memetakan penggunaan memori maksimum ke setiap perangkat untuk menerapkan batas. Jika perangkat tidak ada di max_memory, perangkat tersebut diabaikan dan komponen pipeline tidak akan didistribusikan ke perangkat tersebut.

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 menggunakan memori maksimum dari semua perangkat secara default, tetapi jika model tidak muat di GPU, Anda perlu menggunakan satu GPU dan melakukan offload ke CPU dengan metode di bawah ini.

enable_model_cpu_offload() hanya berfungsi pada satu GPU, tetapi model yang sangat besar mungkin tidak muat di GPU tersebut
enable_sequential_cpu_offload() mungkin berfungsi tetapi sangat lambat dan juga terbatas pada satu GPU

Gunakan metode reset_device_map() untuk mereset device_map. Hal ini diperlukan jika Anda ingin menggunakan metode seperti .to(), enable_sequential_cpu_offload(), dan enable_model_cpu_offload() pada pipeline yang telah dipetakan perangkatnya.

pipeline.reset_device_map()

Pemotongan VAE

Slicing VAE menghemat memori dengan memecah batch input besar menjadi satu batch data dan memprosesnya secara terpisah. Metode ini paling efektif saat menghasilkan lebih dari satu gambar sekaligus.

Sebagai contoh, jika Anda menghasilkan 4 gambar sekaligus, proses decoding akan meningkatkan memori aktivasi puncak sebesar 4x. VAE slicing mengurangi hal ini dengan hanya melakukan decoding 1 gambar pada satu waktu, bukan semua 4 gambar sekaligus.

Panggil enable_slicing() untuk mengaktifkan VAE yang dipotong (sliced). Anda dapat mengharapkan peningkatan kinerja kecil saat mendekode batch multi-gambar dan tidak ada dampak kinerja untuk batch single-gambar.

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

Kelas AutoencoderKLWan dan AsymmetricAutoencoderKL tidak mendukung pemotongan (slicing).

Tiling VAE

Tiling VAE menghemat memori dengan membagi gambar menjadi ubin-ubin kecil yang saling tumpang tindih, alih-alih memproses seluruh gambar sekaligus. Hal ini juga mengurangi penggunaan memori puncak karena GPU hanya memproses satu ubin pada satu waktu.

Panggil enable_tiling() untuk mengaktifkan tiling VAE. Gambar yang dihasilkan mungkin memiliki variasi nada dari satu tile ke tile lainnya karena didekode secara terpisah, tetapi seharusnya tidak ada sambungan yang jelas antar tile. Tiling dinonaktifkan untuk resolusi yang lebih rendah dari batas yang ditentukan sebelumnya (namun dapat dikonfigurasi). Misalnya, batas ini adalah 512x512 untuk VAE di StableDiffusionPipeline.

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 dan AsymmetricAutoencoderKL tidak mendukung tiling.

Untuk sepenuhnya melewati memori VAE di GPU Anda, lakukan enkripsi dan dekripsi pada endpoint jarak jauh dengan Inferensi jarak jauh.

Offloading

Strategi offloading memindahkan lapisan atau model yang saat ini tidak aktif ke CPU untuk menghindari peningkatan memori GPU. Strategi-strategi ini dapat dikombinasikan dengan kuantisasi dan torch.compile untuk menyeimbangkan kecepatan inferensi dan penggunaan memori.

Lihat panduan Compile and offloading quantized models untuk detail lebih lanjut.

Offloading CPU

CPU offloading secara selektif memindahkan bobot dari GPU ke CPU. Ketika suatu komponen diperlukan, komponen tersebut dipindahkan ke GPU, dan ketika tidak diperlukan, komponen tersebut dipindahkan ke CPU. Metode ini bekerja pada submodul, bukan pada model secara keseluruhan. Metode ini menghemat memori dengan menghindari penyimpanan seluruh model di GPU.

CPU offloading secara drastis mengurangi penggunaan memori, tetapi juga sangat lambat karena submodul dipindahkan bolak-balik beberapa kali antar perangkat. Hal ini sering kali tidak praktis karena kecepatannya yang lambat.

Jangan memindahkan pipeline ke CUDA sebelum memanggil enable_sequential_cpu_offload(), karena jika tidak, jumlah memori yang dihemat hanya minimal (lihat issue ini untuk detail lebih lanjut). Ini adalah operasi stateful yang menginstal hook pada model.

Panggil enable_sequential_cpu_offload() untuk mengaktifkannya pada pipeline.

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

Offloading model

Model offloading memindahkan seluruh model ke GPU, bukan secara selektif memindahkan beberapa lapisan atau komponen model. Salah satu model pipeline utama, biasanya text encoder, UNet, dan VAE, ditempatkan di GPU sementara komponen lainnya disimpan di CPU. Komponen seperti UNet yang dijalankan beberapa kali tetap berada di GPU hingga prosesnya benar-benar selesai dan tidak lagi dibutuhkan. Hal ini menghilangkan overhead komunikasi dari CPU offloading dan menjadikan model offloading sebagai alternatif yang lebih cepat. Komprominya adalah penghematan memori tidak akan sebesar itu.

Perlu diingat bahwa jika model digunakan kembali di luar pipeline setelah hook terpasang (lihat Menghapus Hook untuk detail lebih lanjut), Anda perlu menjalankan seluruh pipeline dan model dalam urutan yang diharapkan agar offload dapat dilakukan dengan benar. Ini adalah operasi stateful yang memasang hook pada model.

Panggil enable_model_cpu_offload() untuk mengaktifkannya pada pipeline.

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() juga membantu saat Anda menggunakan metode encode_prompt() secara mandiri untuk menghasilkan hidden state dari text encoder.

Offloading grup

Group offloading memindahkan grup lapisan internal (torch.nn.ModuleList atau torch.nn.Sequential) ke CPU. Metode ini menggunakan memori lebih sedikit daripada model offloading dan lebih cepat daripada CPU offloading karena mengurangi overhead komunikasi.

Group offloading mungkin tidak berfungsi dengan semua model jika implementasi forward mengandung device casting input yang bergantung pada bobot, karena hal ini dapat bertentangan dengan mekanisme device casting dari group offloading.

Aktifkan group offloading dengan mengonfigurasi parameter offload_type menjadi block_level atau leaf_level.

block_level memindahkan kelompok lapisan berdasarkan parameter num_blocks_per_group. Sebagai contoh, jika num_blocks_per_group=2 pada model dengan 40 lapisan, 2 lapisan dimuat dan dipindahkan pada satu waktu (total 20 kali pemuatan/pemindahan). Hal ini secara drastis mengurangi kebutuhan memori.
leaf_level melakukan offload lapisan individu pada level terendah dan setara dengan CPU offloading. Namun, prosesnya dapat dibuat lebih cepat jika Anda menggunakan stream tanpa mengorbankan kecepatan inferensi.

Group offloading didukung untuk pipeline utuh atau model individual. Menerapkan group offloading ke seluruh pipeline adalah opsi yang paling mudah, sementara menerapkannya secara selektif ke model individual memberikan pengguna lebih banyak fleksibilitas untuk menggunakan teknik offloading yang berbeda untuk model yang berbeda.

pipeline

model

CUDA stream

Parameter use_stream dapat diaktifkan untuk perangkat CUDA yang mendukung stream transfer data asinkron untuk mengurangi waktu eksekusi keseluruhan dibandingkan dengan CPU offloading. Parameter ini menumpuk transfer data dan komputasi dengan menggunakan prefetching layer. Layer berikutnya yang akan dieksekusi dimuat ke GPU sementara layer saat ini masih dalam proses eksekusi. Hal ini dapat meningkatkan penggunaan memori CPU secara signifikan, jadi pastikan Anda memiliki memori sebanyak 2x ukuran model.

Set record_stream=True untuk percepatan yang lebih besar dengan sedikit peningkatan penggunaan memori. Lihat dokumentasi torch.Tensor.record_stream untuk mempelajari lebih lanjut.

Saat use_stream=True pada VAE dengan tiling yang diaktifkan, pastikan untuk melakukan dummy forward pass (juga memungkinkan dengan dummy inputs) sebelum inferensi untuk menghindari error device mismatch. Hal ini mungkin tidak berfungsi pada semua implementasi, jadi silakan buka issue jika Anda menemui masalah apa pun.

Jika Anda menggunakan offloading grup block_level dengan use_stream diaktifkan, parameter num_blocks_per_group harus diatur ke 1, jika tidak peringatan akan muncul.

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

Parameter low_cpu_mem_usage dapat diatur ke True untuk mengurangi penggunaan memori CPU saat menggunakan stream selama group offloading. Pengaturan ini paling cocok untuk offloading leaf_level dan ketika memori CPU menjadi bottleneck. Memori dihemat dengan membuat tensor pinned secara on the fly alih-alih mem-pinning-nya sebelumnya. Namun, hal ini dapat meningkatkan waktu eksekusi keseluruhan.

Offloading ke disk

Group offloading dapat mengonsumsi memori sistem yang signifikan tergantung pada ukuran model. Pada sistem dengan memori terbatas, coba gunakan group offloading ke disk sebagai memori sekunder.

Atur argumen offload_to_disk_path di enable_group_offload() atau apply_group_offloading() untuk meng-offload model ke disk.

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

Lihat dua tabel ini untuk membandingkan trade-off antara kecepatan dan memori.

Casting per lapisan

Gabungkan layerwise casting dengan group offloading untuk penghematan memori yang lebih besar.

Layerwise casting menyimpan bobot dalam format data yang lebih kecil (misalnya, torch.float8_e4m3fn dan torch.float8_e5m2) untuk menggunakan lebih sedikit memori, lalu menaikkan presisi bobot tersebut ke presisi yang lebih tinggi seperti torch.float16 atau torch.bfloat16 untuk komputasi. Lapisan tertentu (bobot terkait normalisasi dan modulasi) dilewati karena menyimpannya dalam fp8 dapat menurunkan kualitas generasi.

Casting per lapisan mungkin tidak berfungsi dengan semua model jika implementasi forward mengandung typecasting internal pada bobot. Implementasi saat ini dari casting per lapisan mengasumsikan bahwa jalur forward independen dari presisi bobot dan tipe data input selalu ditentukan dalam compute_dtype (lihat di sini untuk implementasi yang tidak kompatibel).

Casting per lapisan juga dapat gagal pada implementasi pemodelan kustom dengan lapisan PEFT. Ada beberapa pemeriksaan yang tersedia, tetapi mereka tidak diuji secara luas atau dijamin berfungsi dalam semua kasus.

Panggil enable_layerwise_casting() untuk mengatur tipe data penyimpanan dan komputasi.

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)

Metode apply_layerwise_casting() juga dapat digunakan jika Anda membutuhkan kontrol dan fleksibilitas yang lebih besar. Metode ini dapat diterapkan secara parsial pada lapisan model dengan memanggilnya pada modul internal tertentu. Gunakan parameter skip_modules_pattern atau skip_modules_classes untuk menentukan modul yang harus dihindari, seperti lapisan normalisasi dan modulasi.

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 membalik cara tensor disimpan dari (batch size, channels, height, width) menjadi (batch size, height, width, channels). Ini menyelaraskan tensor dengan cara hardware mengakses secara berurutan tensor yang disimpan di memori dan menghindari lompatan di memori untuk mengakses nilai piksel.

Tidak semua operator saat ini mendukung format channels-last dan dapat menghasilkan kinerja yang lebih buruk, tetapi tetap layak dicoba.

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

Perhatian hemat memori

Diffusers mendukung beberapa backend perhatian hemat memori (FlashAttention, xFormers, SageAttention, dan lainnya) melalui set_attention_backend(). Lihat panduan Attention backends untuk mempelajari cara beralih di antara mereka.

Update on GitHub