Percepat inferensi

Model difusi lambat saat inferensi karena generasi adalah proses iteratif di mana noise secara bertahap disempurnakan menjadi gambar atau video dalam sejumlah “langkah” tertentu. Untuk mempercepat proses ini, Anda dapat mencoba bereksperimen dengan berbagai schedulers, mengurangi presisi bobot model untuk komputasi yang lebih cepat, menggunakan mekanisme perhatian yang lebih hemat memori, dan lain-lain.

Gabungkan dan gunakan teknik-teknik ini secara bersamaan agar inferensi lebih cepat dibandingkan menggunakan satu teknik saja.

Panduan ini akan membahas cara mempercepat inferensi.

Tipe data model

Presisi dan tipe data bobot model memengaruhi kecepatan inferensi karena presisi yang lebih tinggi membutuhkan lebih banyak memori untuk dimuat dan lebih banyak waktu untuk melakukan komputasi. PyTorch memuat bobot model dalam float32 atau presisi penuh secara default, sehingga mengubah tipe data adalah cara sederhana untuk mendapatkan inferensi yang lebih cepat dengan cepat.

import torch
from diffusers import StableDiffusionXLPipeline

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

prompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"
pipeline(prompt, num_inference_steps=30).images[0]

import torch
from diffusers import StableDiffusionXLPipeline

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

prompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"
pipeline(prompt, num_inference_steps=30).images[0]

import torch
from diffusers import StableDiffusionXLPipeline

torch.backends.cuda.matmul.allow_tf32 = True

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

prompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"
pipeline(prompt, num_inference_steps=30).images[0]

Perhatian hasil kali titik berskala

Perhatian hemat memori mengoptimalkan kecepatan inferensi dan penggunaan memori!

Scaled dot product attention (SDPA) mengimplementasikan beberapa backend perhatian, FlashAttention, xFormers, dan implementasi C++ native. Backend yang paling optimal untuk perangkat keras Anda akan dipilih secara otomatis.

SDPA diaktifkan secara default jika Anda menggunakan PyTorch >= 2.0 dan tidak diperlukan perubahan tambahan pada kode Anda. Namun, Anda dapat bereksperimen dengan backend attention lainnya jika ingin memilih sendiri. Contoh di bawah ini menggunakan context manager torch.nn.attention.sdpa_kernel untuk mengaktifkan attention yang efisien.

from torch.nn.attention import SDPBackend, sdpa_kernel
import torch
from diffusers import StableDiffusionXLPipeline

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

prompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"

with sdpa_kernel(SDPBackend.EFFICIENT_ATTENTION):
  image = pipeline(prompt, num_inference_steps=30).images[0]

torch.compile

torch.compile mempercepat inferensi dengan mengompilasi kode dan operasi PyTorch menjadi kernel yang dioptimalkan. Diffusers biasanya mengompilasi model yang lebih intensif komputasi seperti UNet, transformer, atau VAE.

Aktifkan pengaturan compiler berikut untuk kecepatan maksimal (lihat daftar lengkap untuk opsi lainnya).

import torch
from diffusers import StableDiffusionXLPipeline

torch._inductor.config.conv_1x1_as_mm = True
torch._inductor.config.coordinate_descent_tuning = True
torch._inductor.config.epilogue_fusion = False
torch._inductor.config.coordinate_descent_check_all_directions = True

Muat dan kompilasi UNet dan VAE. Ada beberapa mode berbeda yang dapat Anda pilih, tetapi "max-autotune" mengoptimalkan untuk kecepatan tercepat dengan mengompilasi ke CUDA graph. CUDA graphs secara efektif mengurangi overhead dengan meluncurkan beberapa operasi GPU melalui satu operasi CPU.

Dengan PyTorch 2.3.1, Anda dapat mengontrol perilaku caching dari torch.compile. Hal ini sangat bermanfaat untuk mode kompilasi seperti "max-autotune" yang melakukan grid-search pada beberapa flag kompilasi untuk menemukan konfigurasi optimal. Pelajari lebih lanjut di tutorial Compile Time Caching in torch.compile.

Mengubah tata letak memori menjadi channels_last juga mengoptimalkan memori dan kecepatan inferensi.

pipeline = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0", dtype=torch.float16
).to("cuda")  # or "mps", "xpu", "cpu"
pipeline.unet.to(memory_format=torch.channels_last)
pipeline.vae.to(memory_format=torch.channels_last)
pipeline.unet = torch.compile(
    pipeline.unet, mode="max-autotune", fullgraph=True
)
pipeline.vae.decode = torch.compile(
    pipeline.vae.decode,
    mode="max-autotune",
    fullgraph=True
)

prompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"
pipeline(prompt, num_inference_steps=30).images[0]

Kompilasi lambat pada percobaan pertama, tetapi setelah dikompilasi, prosesnya jauh lebih cepat. Usahakan untuk hanya menggunakan pipeline yang telah dikompilasi untuk jenis operasi inferensi yang sama. Memanggil pipeline yang telah dikompilasi pada ukuran gambar yang berbeda akan memicu ulang proses kompilasi, yang lambat dan tidak efisien.

Kompilasi bentuk dinamis

Pastikan untuk selalu menggunakan versi nightly dari PyTorch untuk dukungan yang lebih baik.

torch.compile melacak bentuk input dan kondisi, dan jika berbeda, model akan dikompilasi ulang. Misalnya, jika model dikompilasi pada gambar dengan resolusi 1024x1024 dan digunakan pada gambar dengan resolusi berbeda, kompilasi ulang akan dipicu.

Untuk menghindari kompilasi ulang, tambahkan dynamic=True untuk mencoba menghasilkan kernel yang lebih dinamis sehingga kompilasi ulang dapat dihindari ketika kondisi berubah.

+ torch.fx.experimental._config.use_duck_shape = False
+ pipeline.unet = torch.compile(
    pipeline.unet, fullgraph=True, dynamic=True
)

Menentukan use_duck_shape=False memberi instruksi kepada compiler apakah harus menggunakan variabel simbolik yang sama untuk merepresentasikan ukuran input yang sama. Untuk detail lebih lanjut, lihat komentar ini.

Tidak semua model mungkin mendapat manfaat dari kompilasi dinamis secara bawaan dan mungkin memerlukan perubahan. Lihat PR ini yang meningkatkan implementasi AuraFlowPipeline agar dapat memanfaatkan kompilasi dinamis.

Silakan buka issue jika kompilasi dinamis tidak berfungsi sesuai harapan untuk model Diffusers.

Kompilasi regional

Kompilasi regional memangkas latensi cold-start dengan hanya mengompilasi blok kecil yang sering diulang dari sebuah model - biasanya satu lapisan transformer - dan memungkinkan penggunaan ulang artefak terkompilasi untuk setiap kemunculan berikutnya. Untuk banyak arsitektur difusi, hal ini memberikan peningkatan kecepatan runtime yang sama dengan kompilasi grafik penuh dan mengurangi waktu kompilasi hingga 8–10x.

Gunakan metode compile_repeated_blocks(), sebuah helper yang membungkus torch.compile, pada komponen apa pun seperti model transformer seperti yang ditunjukkan di bawah ini.

# pip install -U diffusers
import torch
from diffusers import StableDiffusionXLPipeline

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

# compile only the repeated transformer layers inside the UNet
pipeline.unet.compile_repeated_blocks(fullgraph=True)

Untuk mengaktifkan kompilasi regional untuk model baru, tambahkan atribut _repeated_blocks ke kelas model yang berisi nama kelas (sebagai string) dari blok-blok yang ingin Anda kompilasi.

class MyUNet(ModelMixin):
    _repeated_blocks = ("Transformer2DModel",)  # ← compiled by default

Untuk contoh kompilasi regional lainnya, lihat PR referensi.

Ada juga metode compile_regions di Accelerate yang secara otomatis memilih blok kandidat dalam model untuk dikompilasi. Sisa graf dikompilasi secara terpisah. Ini berguna untuk eksperimen cepat karena tidak ada banyak opsi untuk Anda mengatur blok mana yang akan dikompilasi atau menyesuaikan flag kompilasi.

# pip install -U accelerate
import torch
from diffusers import StableDiffusionXLPipeline
from accelerate.utils import compile_regions

pipeline = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0", dtype=torch.float16
).to("cuda")  # or "mps", "xpu", "cpu"
pipeline.unet = compile_regions(pipeline.unet, mode="reduce-overhead", fullgraph=True)

compile_repeated_blocks() sengaja dibuat eksplisit. Daftarkan blok yang akan diulang dalam _repeated_blocks dan helper hanya akan mengompilasi blok-blok tersebut. Hal ini menawarkan perilaku yang dapat diprediksi dan memudahkan penalaran tentang penggunaan ulang cache dalam satu baris kode.

Graph breaks

Penting untuk menentukan fullgraph=True di torch.compile untuk memastikan tidak ada graph breaks pada model yang mendasarinya. Hal ini memungkinkan Anda memanfaatkan torch.compile tanpa degradasi kinerja. Untuk UNet dan VAE, ini mengubah cara Anda mengakses variabel return.

- latents = unet(
-   latents, timestep=timestep, encoder_hidden_states=prompt_embeds
-).sample

+ latents = unet(
+   latents, timestep=timestep, encoder_hidden_states=prompt_embeds, return_dict=False
+)[0]

Sinkronisasi GPU

Fungsi step() dipanggil pada scheduler setiap kali setelah denoiser membuat prediksi, dan variabel sigmas diindeks. Ketika ditempatkan di GPU, hal ini memperkenalkan latensi karena sinkronisasi komunikasi antara CPU dan GPU. Hal ini menjadi lebih jelas ketika denoiser sudah dikompilasi.

Secara umum, sigmas sebaiknya tetap di CPU untuk menghindari sinkronisasi komunikasi dan latensi.

Lihat posting blog torch.compile and Diffusers: A Hands-On Guide to Peak Performance untuk memaksimalkan kinerja dengan torch.compile untuk model difusi.

Benchmarks

Lihat dataset diffusers/benchmarks untuk melihat data latensi inferensi dan penggunaan memori untuk pipeline yang dikompilasi.

Repositori diffusers-torchao juga berisi hasil benchmarking untuk versi terkompilasi dari Flux dan CogVideoX.

Kernels

Kernels adalah pustaka untuk membangun, mendistribusikan, dan memuat kernel komputasi yang dioptimalkan di Hub. Pustaka ini mendukung kernel attention dan kernel CUDA kustom untuk operasi seperti RMSNorm, GEGLU, RoPE, dan AdaLN.

Panduan Integrasi Pipeline Diffusers menunjukkan cara mengintegrasikan kernel dengan skill add cuda-kernels. Skill ini memungkinkan agen, seperti Claude atau Codex, untuk menulis kernel khusus yang ditargetkan ke model tertentu dan perangkat keras Anda.

Instal skill add cuda-kernels untuk mengajarkan agen cara menulis kernel. Artikel blog Custom kernels for all from Codex and Claude membahas hal ini secara lebih detail.

Sebagai contoh, kernel RMSNorm kustom (yang dihasilkan oleh skill add cuda-kernels) dengan torch.compile mempercepat generasi LTX-Video sebesar 1,43x pada H100.

Kuantisasi dinamis

Kuantisasi dinamis meningkatkan kecepatan inferensi dengan mengurangi presisi untuk memungkinkan operasi matematika yang lebih cepat. Jenis kuantisasi ini menentukan cara menskalakan aktivasi berdasarkan data pada saat runtime, alih-alih menggunakan faktor penskalaan yang tetap. Akibatnya, faktor penskalaan lebih akurat selaras dengan data.

Contoh di bawah ini menerapkan kuantisasi int8 dinamis pada UNet dan VAE dengan pustaka torchao.

Lihat dokumentasi torchao kami untuk mempelajari lebih lanjut tentang cara menggunakan integrasi torchao Diffusers.

Konfigurasikan tag compiler untuk kecepatan maksimal.

import torch
from torchao import apply_dynamic_quant
from diffusers import StableDiffusionXLPipeline

torch._inductor.config.conv_1x1_as_mm = True
torch._inductor.config.coordinate_descent_tuning = True
torch._inductor.config.epilogue_fusion = False
torch._inductor.config.coordinate_descent_check_all_directions = True
torch._inductor.config.force_fuse_int_mm_with_mul = True
torch._inductor.config.use_mixed_mm = True

Saring beberapa lapisan linear di UNet dan VAE yang tidak diuntungkan oleh kuantisasi dinamis dengan dynamic_quant_filter_fn.

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

apply_dynamic_quant(pipeline.unet, dynamic_quant_filter_fn)
apply_dynamic_quant(pipeline.vae, dynamic_quant_filter_fn)

prompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"
pipeline(prompt, num_inference_steps=30).images[0]

Matriks proyeksi yang digabungkan

Metode fuse_qkv_projections bersifat eksperimental dan dukungannya terbatas terutama pada pipeline Stable Diffusion. Lihat PR ini untuk mempelajari lebih lanjut tentang cara mengaktifkannya untuk pipeline lain

Sebuah input diproyeksikan ke tiga subruang, yang direpresentasikan oleh matriks proyeksi Q, K, dan V, dalam sebuah blok perhatian. Proyeksi-proyeksi ini biasanya dihitung secara terpisah, tetapi Anda dapat menggabungkannya secara horizontal menjadi satu matriks tunggal dan melakukan proyeksinya dalam satu langkah. Hal ini meningkatkan ukuran perkalian matriks dari proyeksi input dan juga meningkatkan dampak dari kuantisasi.

pipeline.fuse_qkv_projections()

Sumber daya

• Baca artikel blog Presenting Flux Fast: Making Flux go brrr on H100s untuk mempelajari lebih lanjut tentang cara Anda menggabungkan semua optimasi ini dengan TorchInductor dan AOTInductor untuk mendapatkan percepatan sekitar 2,5x menggunakan resep dari flux-fast.

  Resep-resep ini mendukung perangkat keras AMD dan Flux.1 Kontext Dev.

• Baca artikel blog torch.compile and Diffusers: A Hands-On Guide to Peak Performance untuk memaksimalkan kinerja saat menggunakan torch.compile.