Tăng tốc suy luận

Mô hình Diffusion chạy chậm ở giai đoạn suy luận vì quá trình tạo ra là một quy trình lặp, trong đó nhiễu được tinh chỉnh dần thành hình ảnh hoặc video qua một số “bước” nhất định. Để tăng tốc quá trình này, bạn có thể thử nghiệm với các schedulers khác nhau, giảm độ chính xác của trọng số mô hình để tính toán nhanh hơn, sử dụng các cơ chế chú ý tiết kiệm bộ nhớ hơn, và nhiều phương pháp khác.

Kết hợp và sử dụng các kỹ thuật này cùng nhau để tăng tốc độ suy luận nhanh hơn so với việc chỉ sử dụng một kỹ thuật đơn lẻ.

Hướng dẫn này sẽ trình bày cách tăng tốc suy luận.

Loại dữ liệu của mô hình

Độ chính xác và kiểu dữ liệu của trọng số mô hình ảnh hưởng đến tốc độ suy luận vì độ chính xác cao hơn đòi hỏi nhiều bộ nhớ hơn để nạp và nhiều thời gian hơn để thực hiện các phép tính. PyTorch nạp trọng số mô hình ở float32 hoặc độ chính xác đầy đủ theo mặc định, do đó việc thay đổi kiểu dữ liệu là một cách đơn giản để nhanh chóng đạt được suy luận nhanh hơn.

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]

Chú ý tích vô hướng có tỷ lệ

Chú ý tiết kiệm bộ nhớ được tối ưu hóa cho tốc độ suy luận và lượng bộ nhớ sử dụng!

Scaled dot product attention (SDPA) triển khai một số backend attention, FlashAttention, xFormers, và một bản triển khai C++ native. Nó tự động chọn backend tối ưu nhất cho phần cứng của bạn.

SDPA được bật mặc định nếu bạn đang sử dụng PyTorch >= 2.0 và không cần thay đổi thêm gì trong mã của bạn. Tuy nhiên, bạn có thể thử nghiệm với các backend attention khác nếu muốn tự chọn. Ví dụ bên dưới sử dụng context manager torch.nn.attention.sdpa_kernel để bật attention hiệu quả.

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 tăng tốc độ suy luận bằng cách biên dịch mã và các thao tác PyTorch thành các kernel được tối ưu hóa. Diffusers thường biên dịch các model có cường độ tính toán cao hơn như UNet, transformer hoặc VAE.

Bật các cài đặt trình biên dịch sau để đạt tốc độ tối đa (xem danh sách đầy đủ để biết thêm các tùy chọn).

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

Tải và biên dịch UNet và VAE. Có một số chế độ khác nhau để bạn lựa chọn, nhưng "max-autotune" tối ưu cho tốc độ nhanh nhất bằng cách biên dịch thành một CUDA graph. CUDA graphs thực sự giảm thiểu overhead bằng cách khởi chạy nhiều thao tác GPU thông qua một thao tác CPU duy nhất.

Với PyTorch 2.3.1, bạn có thể kiểm soát hành vi cache của torch.compile. Điều này đặc biệt hữu ích cho các chế độ biên dịch như "max-autotune", vốn thực hiện tìm kiếm lưới trên một số cờ biên dịch để tìm cấu hình tối ưu. Tìm hiểu thêm trong hướng dẫn Compile Time Caching in torch.compile.

Việc thay đổi bố trí bộ nhớ thành channels_last cũng tối ưu hóa bộ nhớ và tốc độ suy luận.

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]

Việc biên dịch diễn ra chậm ở lần đầu tiên, nhưng sau khi được biên dịch, tốc độ sẽ nhanh hơn đáng kể. Hãy cố gắng chỉ sử dụng pipeline đã biên dịch cho cùng một loại thao tác suy luận. Việc gọi pipeline đã biên dịch với kích thước ảnh khác sẽ kích hoạt lại quá trình biên dịch, gây chậm và kém hiệu quả.

Biên dịch hình dạng động

Hãy luôn sử dụng phiên bản nightly của PyTorch để có hỗ trợ tốt hơn.

torch.compile theo dõi hình dạng đầu vào và các điều kiện, và nếu chúng khác đi, nó sẽ biên dịch lại mô hình. Ví dụ, nếu một mô hình được biên dịch trên ảnh có độ phân giải 1024x1024 và được sử dụng trên ảnh có độ phân giải khác, quá trình biên dịch lại sẽ được kích hoạt.

Để tránh biên dịch lại, hãy thêm dynamic=True để cố gắng tạo ra một kernel động hơn, giúp tránh việc biên dịch lại khi các điều kiện thay đổi.

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

Việc chỉ định use_duck_shape=False hướng dẫn trình biên dịch xem có nên sử dụng cùng một biến biểu tượng để đại diện cho các kích thước đầu vào giống nhau hay không. Để biết thêm chi tiết, hãy xem lời bình này.

Không phải tất cả các model đều có thể hưởng lợi từ biên dịch động ngay từ đầu và có thể cần được điều chỉnh. Tham khảo PR này đã cải thiện triển khai AuraFlowPipeline để tận dụng biên dịch động.

Hãy mở một issue nếu biên dịch động không hoạt động như mong đợi với một mô hình Diffusers.

Biên dịch theo vùng

Biên dịch theo vùng giảm độ trễ khởi động lạnh bằng cách chỉ biên dịch khối nhỏ và lặp lại thường xuyên của một mô hình - thường là một lớp transformer - và cho phép tái sử dụng các sản phẩm đã biên dịch cho mọi lần xuất hiện tiếp theo. Đối với nhiều kiến trúc khuếch tán, điều này mang lại tốc độ chạy thời gian thực tương đương với biên dịch toàn đồ thị và giảm thời gian biên dịch từ 8 đến 10 lần.

Sử dụng phương thức compile_repeated_blocks(), một hàm trợ giúp bọc torch.compile, trên bất kỳ thành phần nào, chẳng hạn như mô hình transformer, như minh họa bên dưới.

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

Để bật biên dịch theo vùng cho một mô hình mới, hãy thêm thuộc tính _repeated_blocks vào lớp mô hình, chứa tên các lớp (dưới dạng chuỗi) của các khối mà bạn muốn biên dịch.

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

Để xem thêm các ví dụ biên dịch theo khu vực, hãy xem PR tham chiếu.

Ngoài ra, Accelerate còn có phương thức compile_regions tự động chọn các khối ứng viên trong mô hình để biên dịch. Phần đồ thị còn lại được biên dịch riêng biệt. Điều này hữu ích cho các thí nghiệm nhanh vì có ít tùy chọn hơn để bạn thiết lập khối nào cần biên dịch hoặc điều chỉnh cờ biên dịch.

# 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() được thiết kế để rõ ràng một cách có chủ đích. Hãy liệt kê các khối cần lặp lại trong _repeated_blocks và hàm trợ giúp chỉ biên dịch các khối đó. Nó mang lại hành vi có thể dự đoán được và dễ suy luận về việc tái sử dụng cache chỉ trong một dòng mã.

Điểm ngắt đồ thị

Việc chỉ định fullgraph=True trong torch.compile là rất quan trọng để đảm bảo không có sự ngắt mạch (graph breaks) trong mô hình cơ bản. Điều này cho phép bạn tận dụng torch.compile mà không bị suy giảm hiệu suất. Đối với UNet và VAE, điều này thay đổi cách bạn truy cập các biến trả về.

- 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]

Đồng bộ GPU

Hàm step() được gọi trên scheduler sau mỗi lần denoiser đưa ra dự đoán, và biến sigmas được truy cập theo chỉ số. Khi đặt trên GPU, nó gây ra độ trễ do đồng bộ hóa giao tiếp giữa CPU và GPU. Hiện tượng này trở nên rõ rệt hơn khi denoiser đã được biên dịch.

Nhìn chung, sigmas nên giữ trên CPU để tránh đồng bộ hóa giao tiếp và độ trễ.

Tham khảo bài viết trên blog torch.compile and Diffusers: A Hands-On Guide to Peak Performance để tối đa hóa hiệu suất với torch.compile cho các mô hình khuếch tán.

Benchmarks

Tham khảo bộ dữ liệu diffusers/benchmarks để xem dữ liệu về độ trễ suy luận và mức sử dụng bộ nhớ của các pipeline đã biên dịch.

Kho diffusers-torchao cũng chứa các kết quả benchmark cho các phiên bản đã biên dịch của Flux và CogVideoX.

Kernels

Kernels là một thư viện để xây dựng, phân phối và nạp các kernel tính toán tối ưu trên Hub. Nó hỗ trợ các kernel attention và các kernel CUDA tùy chỉnh cho các phép toán như RMSNorm, GEGLU, RoPE và AdaLN.

Hướng dẫn Tích hợp Pipeline Diffusers cho thấy cách tích hợp một kernel với kỹ năng add cuda-kernels. Kỹ năng này cho phép một agent, chẳng hạn như Claude hoặc Codex, viết các kernel tùy chỉnh nhắm đến một mô hình cụ thể và phần cứng của bạn.

Cài đặt kỹ năng add cuda-kernels để hướng dẫn một agent cách viết kernel. Bài đăng trên blog Custom kernels for all from Codex and Claude đề cập chi tiết hơn về vấn đề này.

Ví dụ, một kernel RMSNorm tùy chỉnh (được tạo bởi kỹ năng add cuda-kernels) kết hợp với torch.compile giúp tăng tốc độ tạo video LTX-Video lên 1,43 lần trên H100.

Lượng tử hóa động

Dynamic quantization cải thiện tốc độ suy luận bằng cách giảm độ chính xác để cho phép các phép toán nhanh hơn. Loại lượng tử hóa cụ thể này xác định cách mở rộng các kích hoạt dựa trên dữ liệu tại thời điểm chạy thay vì sử dụng một hệ số mở rộng cố định. Kết quả là, hệ số mở được căn chỉnh chính xác hơn với dữ liệu.

Ví dụ bên dưới áp dụng lượng tử hóa int8 động cho UNet và VAE bằng thư viện torchao.

Tham khảo tài liệu torchao của chúng tôi để tìm hiểu thêm về cách sử dụng tích hợp torchao của Diffusers.

Cấu hình các thẻ trình biên dịch để đạt tốc độ tối đa.

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

Loại bỏ một số lớp tuyến tính trong UNet và VAE không được hưởng lợi từ lượng tử hóa động bằng 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]

Các ma trận chiếu đã hợp nhất

Phương thức fuse_qkv_projections mang tính thử nghiệm và hỗ trợ chủ yếu giới hạn ở các pipeline Stable Diffusion. Xem PR này để tìm hiểu cách bật tính năng này cho các pipeline khác.

Một đầu vào được chiếu vào ba không gian con, được biểu diễn bởi các ma trận chiếu Q, K và V, trong một khối chú ý. Các phép chiếu này thường được tính toán riêng biệt, nhưng bạn có thể kết hợp ngang chúng thành một ma trận duy nhất và thực hiện phép chiếu trong một bước duy nhất. Điều này làm tăng kích thước của các phép nhân ma trận của các phép chiếu đầu vào và cũng cải thiện tác động của lượng tử hóa.

pipeline.fuse_qkv_projections()

Tài nguyên

• Đọc bài đăng trên blog Presenting Flux Fast: Making Flux go brrr on H100s để tìm hiểu thêm về cách bạn có thể kết hợp tất cả các tối ưu hóa này với TorchInductor và AOTInductor để đạt được tốc độ nhanh hơn ~2,5 lần bằng cách sử dụng các công thức từ flux-fast.

  Các công thức này hỗ trợ phần cứng AMD và Flux.1 Kontext Dev.

• Đọc bài đăng trên blog torch.compile và Diffusers: Hướng dẫn thực hành để đạt hiệu suất tối đa để tối đa hóa hiệu suất khi sử dụng torch.compile.