diffusers/main/vi/optimization/fp16.md

Độ chính xác và biên dịch

Độ chính xác thấp hơn và biên dịch là hai cách để tăng tốc suy luận của Diffusers. Nạp trọng số ở bfloat16 hoặc float16, sau đó biên dịch bộ khử nhiễu (UNet hoặc transformer) bằng torch.compile hoặc biên dịch theo vùng.

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. Diffusers nạp trọng số mô hình ở float32 khi bạn bỏ qua dtype, 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.

bfloat16 tương tự float16 nhưng bền vững hơn trước các lỗi số. Hỗ trợ phần cứng cho bfloat16 khác nhau, nhưng hầu hết các GPU hiện đại đều có khả năng hỗ trợ 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 tương tự như bfloat16 nhưng có thể dễ gặp lỗi số hơn.

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]

Chế độ TensorFloat-32 (tf32) được hỗ trợ trên các GPU NVIDIA Ampere và mới hơn, và nó thực hiện các phép toán tích chập và nhân ma trận ở tf32. Việc lưu trữ và các phép toán khác được giữ ở float32. Nó tăng tốc các phép toán chạy ở float32, do đó hữu ích nhất khi pipeline hoặc một phần của nó vẫn ở float32.

PyTorch chỉ bật chế độ tf32 cho các phép tích chập theo mặc định và bạn sẽ cần bật nó một cách rõ ràng cho các phép nhân ma trận.

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]

Tham khảo tài liệu huấn luyện độ chính xác hỗn hợp để biết thêm chi tiết.

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

Scaled dot product attention (SDPA) là cơ chế attention mặc định trên PyTorch 2.0 trở lên, thông qua AttnProcessor2_0 hoặc backend native của attention dispatcher. PyTorch sẽ chọn một kernel SDPA phù hợp với phần cứng của bạn. Để biết thêm về FlashAttention, SageAttention, xFormers, Hub kernels và các backend khác, hãy xem Attention backends.

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 tối ưu. Bạn thường biên dịch denoiser, thành phần chiếm phần lớn thời gian chạy, cụ thể là pipeline.unet hoặc pipeline.transformer, và đôi khi 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 denoiser cùng VAE. Sử dụng pipeline.unet trên các pipeline UNet như Stable Diffusion XL, hoặc pipeline.transformer trên các pipeline kiểu Flux và DiT. Có một số chế độ khác nhau để bạn lựa chọn, nhưng "max-autotune" sẽ tìm kiếm các kernel nhanh nhất và sử dụng CUDA graphs. CUDA graphs 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.

Việc thay đổi bố trí bộ nhớ thành channels_last cũng có thể tăng tốc các mô hình nặng về phép卷积 như UNets và VAEs. Hãy chạy benchmark trước, vì một số mô hình có thể chạy chậm hơ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.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]

Việc biên dịch diễn ra chậm ở lần đầu tiên, đặc biệt khi sử dụng "max-autotune", vì trình biên dịch sẽ tìm kiếm các kernel và sau đó tiếp tục sử dụng cùng một đối tượng pipeline đã biên dịch (xem Compile Time Caching và Caching Configuration để biết cách cấu hình hành vi lưu cache). Việc gọi pipeline đã biên dịch với kích thước ảnh khác sẽ kích hoạt quá trình biên dịch lại.

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

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 tương đương với biên dịch toàn đồ thị và có thể giảm thời gian biên dịch lên đến 8–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 denoiser. Gọi nó trên pipeline.unet hoặc pipeline.transformer, tùy thuộc vào 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)

Để bật biên dịch theo vùng cho một mô hình mới, hãy đặt _repeated_blocks thành tên các lớp của các khối lặp lại (chuỗi). Đối với UNet của SDXL, hãy sử dụng BasicTransformerBlock. Các bộ khử nhiễu Transformer cũng liệt kê tên lớp khối lặp lại của chúng theo cách tương tự.

_repeated_blocks = ["BasicTransformerBlock"]

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ị

Đặt fullgraph=True để torch.compile báo lỗi khi gặp graph break thay vì âm thầm tách graph, điều này làm giảm tốc độ tăng tốc.

Đồng bộ GPU

Sau mỗi lần dự đoán của denoiser, pipeline gọi phương thức step của scheduler, phương thức này sẽ truy cập vào sigmas. Việc giữ sigmas trên GPU có thể buộc đồng bộ hóa CPU↔GPU ở mỗi bước. Chi phí này dễ bị bỏ qua cho đến khi denoiser được biên dịch. Nên để sigmas trên CPU (các scheduler của Diffusers như Euler đã làm điều này).

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. Bài đăng Các kernel tùy chỉnh cho tất cả từ Codex và Claude có thêm chi tiết.

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.