Diffusers 文件
加速推理
並獲得增強的文件體驗
開始使用
加速推理
擴散模型在推理時速度較慢,因為生成是一個迭代過程,其中噪聲在一定數量的“步”中逐漸細化為影像或影片。為了加快此過程,您可以嘗試使用不同的排程器,降低模型權重的精度以加快計算,使用更節省記憶體的注意力機制等。
結合使用這些技術,推理速度將比單獨使用任何一種技術都要快。
本指南將介紹如何加速推理。
模型資料型別
模型權重的精度和資料型別會影響推理速度,因為更高的精度需要更多的記憶體來載入和更多的時間來執行計算。PyTorch 預設以 float32 或全精度載入模型權重,因此更改資料型別是快速獲得更快推理的簡單方法。
bfloat16 類似於 float16,但對數值誤差更魯棒。bfloat16 的硬體支援各不相同,但大多數現代 GPU 都能夠支援 bfloat16。
import torch
from diffusers import StableDiffusionXLPipeline
pipeline = StableDiffusionXLPipeline.from_pretrained(
"stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.bfloat16
).to("cuda")
prompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"
pipeline(prompt, num_inference_steps=30).images[0]
縮放點積注意力
記憶體高效注意力機制優化了推理速度*和*記憶體使用!
縮放點積注意力(SDPA)實現了多種注意力後端,包括FlashAttention、xFormers和一個原生的 C++ 實現。它會自動為您的硬體選擇最優的後端。
如果您使用 PyTorch >= 2.0,SDPA 預設啟用,無需對程式碼進行額外更改。不過,如果您想選擇自己的注意力後端,也可以嘗試其他後端。下面的示例使用 torch.nn.attention.sdpa_kernel 上下文管理器來啟用高效注意力。
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", torch_dtype=torch.bfloat16
).to("cuda")
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 透過將 PyTorch 程式碼和操作編譯成最佳化過的核心來加速推理。Diffusers 通常會編譯計算密集型模型,如 UNet、Transformer 或 VAE。
啟用以下編譯器設定以獲得最大速度(更多選項請參閱完整列表)。
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
載入並編譯 UNet 和 VAE。您可以選擇多種不同的模式,但 `"max-autotune"` 透過編譯到 CUDA 圖來最佳化最快速度。CUDA 圖透過單個 CPU 操作啟動多個 GPU 操作,有效降低了開銷。
使用 PyTorch 2.3.1,您可以控制 torch.compile 的快取行為。這對於像 `"max-autotune"` 這樣的編譯模式特別有益,它會對多個編譯標誌進行網格搜尋以找到最佳配置。在torch.compile 中的編譯時間快取教程中瞭解更多資訊。
將記憶體佈局更改為channels_last也可以最佳化記憶體和推理速度。
pipeline = StableDiffusionXLPipeline.from_pretrained(
"stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16
).to("cuda")
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]
第一次編譯速度較慢,但一旦編譯完成,速度會顯著加快。嘗試僅在相同型別的推理操作上使用編譯後的流水線。在不同影像尺寸上呼叫編譯後的流水線會重新觸發編譯,這既慢又低效。
區域編譯
區域編譯透過僅編譯模型的特定重複區域(或塊)而不是整個模型來減少冷啟動編譯時間。編譯器會重用其他塊的快取和編譯程式碼。
Accelerate 提供了 compile_regions 方法,用於自動順序編譯 `nn.Module` 的重複塊。模型的其餘部分則單獨編譯。
# 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", torch_dtype=torch.float16
).to("cuda")
pipeline.unet = compile_regions(pipeline.unet, mode="reduce-overhead", fullgraph=True)
圖斷裂
重要的是在 torch.compile 中指定 `fullgraph=True`,以確保底層模型中沒有圖斷裂。這允許您利用 torch.compile 的優勢,而不會有任何效能下降。對於 UNet 和 VAE,這會改變您訪問返回變數的方式。
- 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]
GPU 同步
在去噪器做出預測後,排程器會每次呼叫 `step()` 函式,並對 `sigmas` 變數進行索引。當放置在 GPU 上時,由於 CPU 和 GPU 之間的通訊同步,它會引入延遲。當去噪器已被編譯時,這會變得更加明顯。
一般來說,`sigmas` 應該保留在 CPU 上,以避免通訊同步和延遲。
基準測試
請參閱diffusers/benchmarks資料集,檢視編譯後流水線的推理延遲和記憶體使用資料。
diffusers-torchao 儲存庫還包含 Flux 和 CogVideoX 編譯版本的基準測試結果。
動態量化
動態量化透過降低精度以實現更快的數學運算來提高推理速度。這種特殊型別的量化會根據執行時的資料而不是使用固定縮放因子來確定如何縮放啟用。因此,縮放因子與資料更準確地對齊。
以下示例使用 torchao 庫對 UNet 和 VAE 應用 動態 int8 量化。
請參閱我們的torchao文件,瞭解有關如何使用 Diffusers torchao 整合的更多資訊。
配置編譯器標籤以獲得最大速度。
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
使用 dynamic_quant_filter_fn 過濾掉 UNet 和 VAE 中一些不從動態量化中受益的線性層。
pipeline = StableDiffusionXLPipeline.from_pretrained(
"stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.bfloat16
).to("cuda")
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]
融合投影矩陣
fuse_qkv_projections 方法是實驗性的,目前僅限於 Stable Diffusion 流水線。請檢視此 PR 以瞭解如何在其他流水線中啟用它。
在注意力塊中,輸入被投影到三個子空間,由投影矩陣 Q、K 和 V 表示。這些投影通常是分開計算的,但您可以將它們水平合併為一個矩陣,並一步執行投影。這增加了輸入投影的矩陣乘法的大小,也提高了量化的影響。
pipeline.fuse_qkv_projections()