diffusers/main/id/optimization/fp16.md

Presisi dan kompilasi

Presisi rendah dan kompilasi adalah dua cara untuk mempercepat inferensi Diffusers. Muat bobot dalam bfloat16 atau float16, lalu kompilasi denoiser (UNet atau transformer) dengan torch.compile atau kompilasi regional.

Tipe data model

Presisi dan tipe data dari bobot model memengaruhi kecepatan inferensi karena presisi yang lebih tinggi membutuhkan lebih banyak memori untuk dimuat dan lebih banyak waktu untuk melakukan komputasi. Diffusers memuat bobot model dalam float32 ketika Anda mengabaikan dtype, sehingga mengubah tipe data adalah cara sederhana untuk mendapatkan inferensi yang lebih cepat.

bfloat16 mirip dengan float16 tetapi lebih tahan terhadap kesalahan numerik. Dukungan perangkat keras untuk bfloat16 bervariasi, tetapi sebagian besar GPU modern mampu mendukung bfloat16.

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]

float16 mirip dengan bfloat16 tetapi lebih rentan terhadap kesalahan numerik.

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]

Mode TensorFloat-32 (tf32) didukung pada GPU NVIDIA Ampere dan yang lebih baru, serta menghitung operasi konvolusi dan perkalian matriks dalam tf32. Penyimpanan dan operasi lainnya tetap dalam float32. Mode ini mempercepat operasi yang berjalan dalam float32, sehingga paling membantu ketika pipeline atau sebagian darinya tetap dalam float32.

PyTorch hanya mengaktifkan mode tf32 untuk konvolusi secara bawaan dan Anda perlu mengaktifkannya secara eksplisit untuk perkalian matriks.

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.float32
).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]

Lihat dokumentasi pelatihan presisi campuran untuk detail lebih lanjut.

Perhatian hasil kali titik berskala

Scaled dot product attention (SDPA) adalah perhatian default pada PyTorch 2.0 dan versi setelahnya, melalui AttnProcessor2_0 atau backend native dari dispatcher perhatian. PyTorch memilih kernel SDPA untuk perangkat keras Anda. Untuk FlashAttention, SageAttention, xFormers, kernel Hub, dan backend lainnya, lihat Attention backends.

torch.compile

torch.compile mempercepat inferensi dengan mengompilasi kode dan operasi PyTorch menjadi kernel yang dioptimalkan. Anda biasanya mengompilasi denoiser yang mendominasi waktu eksekusi, yaitu pipeline.unet atau pipeline.transformer, dan terkadang juga 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 denoiser serta VAE. Gunakan pipeline.unet pada pipeline UNet seperti Stable Diffusion XL, atau pipeline.transformer pada pipeline Flux dan gaya DiT. Ada beberapa mode berbeda yang dapat Anda pilih, tetapi "max-autotune" mencari kernel tercepat dan menggunakan CUDA graphs. CUDA graphs mengurangi overhead dengan meluncurkan beberapa operasi GPU melalui satu operasi CPU.

Mengubah tata letak memori menjadi channels_last juga dapat mempercepat model yang banyak menggunakan konvolusi seperti UNets dan VAEs. Lakukan benchmark terlebih dahulu, karena beberapa model justru berjalan lebih lambat.

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.unet = torch.compile(
    pipeline.unet, mode="max-autotune", fullgraph=True
)
pipeline.vae.to(memory_format=torch.channels_last)
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 berjalan lambat pada percobaan pertama, terutama dengan "max-autotune", karena compiler mencari kernel dan kemudian terus menggunakan objek pipeline terkompilasi yang sama (lihat Compile Time Caching dan Caching Configuration untuk cara mengonfigurasi perilaku caching). Memanggil pipeline terkompilasi pada ukuran gambar yang berbeda memicu kompilasi ulang.

Kompilasi bentuk dinamis

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 masalah (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 setara dengan kompilasi grafik penuh dan dapat memangkas waktu kompilasi hingga 8–10x.

Gunakan metode compile_repeated_blocks(), sebuah helper yang membungkus torch.compile, pada denoiser. Panggil pada pipeline.unet atau pipeline.transformer, tergantung pada pipeline.

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"
pipeline.unet.compile_repeated_blocks(fullgraph=True)

Untuk mengaktifkan kompilasi regional untuk model baru, atur _repeated_blocks ke nama kelas dari blok yang diulang (string). Untuk UNet SDXL, gunakan BasicTransformerBlock. Denoiser Transformer mencantumkan nama kelas blok berulang mereka sendiri dengan cara yang sama.

_repeated_blocks = ["BasicTransformerBlock"]

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

Setel fullgraph=True agar torch.compile menghasilkan error saat terjadi graph break, alih-alih secara diam-diam memecah graph, yang akan mengurangi percepatan.

Sinkronisasi GPU

Setiap prediksi denoiser, pipeline memanggil metode step dari scheduler, yang mengindeks sigmas. Menjaga sigmas di GPU dapat memaksa sinkronisasi CPU↔GPU pada setiap langkah. Biaya ini mudah terlewat sampai denoiser dikompilasi. Lebih baik biarkan sigmas tetap di CPU (scheduler Diffusers seperti Euler sudah melakukannya).

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 kustom yang ditargetkan ke model tertentu dan perangkat keras Anda. Postingan Custom kernels for all from Codex and Claude memiliki detail lebih lanjut.

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