Đưa kỹ thuật suy luận khuếch tán 4-bit Nunchaku vào thư viện Diffusers

Quay lại các bài viết

Các mô hình diffusion transformer lớn có thể tạo ra những hình ảnh (hoặc thậm chí video, đoạn âm thanh và giờ là cả văn bản) tuyệt đẹp, nhưng việc tải một mô hình text-to-image hiện đại ở độ chính xác BF16 thường yêu cầu 20-30 GB VRAM, điều này khiến các mô hình này nằm ngoài tầm với của hầu hết GPU phổ thông. Lượng tử hóa (quantization) là một giải pháp mạnh mẽ cho vấn đề này, và Diffusers đã tích hợp sẵn một số backend lượng tử hóa như bitsandbytes, GGUF, torchao và Quanto, những thứ mà chúng tôi đã đề cập trong bài "Khám phá các Backend Lượng tử hóa trong Diffusers".

Hầu hết các backend này chỉ hỗ trợ lượng tử hóa trọng số (weight-only). Điều này có nghĩa là chúng lưu trữ các trọng số ở độ chính xác thấp và giải lượng tử hóa chúng về độ chính xác cao tại thời điểm tính toán. Cách này giúp giảm đáng kể mức sử dụng bộ nhớ, nhưng thường không làm tăng tốc độ suy luận (inference) và thậm chí có thể gây thêm một chút độ trễ.

SVDQuant, phương pháp lượng tử hóa đứng sau công cụ suy luận Nunchaku phổ biến, lại có một cách tiếp cận khác. Nó chạy các lớp transformer chính với trọng số và kích hoạt 4-bit (W4A4), giúp giảm bộ nhớ đồng thời tăng tốc vòng lặp khử nhiễu (denoising loop). Các chi tiết sẽ được đề cập bên dưới, nhưng cho đến nay, việc sử dụng các checkpoint này đòi hỏi một thư viện suy luận riêng biệt.

Với Diffusers hiện tại, việc tải một checkpoint Nunchaku đơn giản như gọi hàm from_pretrained, không cần biên dịch CUDA cục bộ nhờ vào gói kernels. Ngoài ra, bộ công cụ đi kèm diffuse-compressor cho phép bạn tự lượng tử hóa các kiến trúc mới và xuất bản chúng dưới dạng các kho lưu trữ Diffusers thông thường.

Mục lục

Bắt đầu với Nunchaku Lite

Đầu tiên, hãy cài đặt các yêu cầu. Bạn cần một phiên bản Diffusers mới nhất và gói Hugging Face kernels:

Sau đó, tải một pipeline đã được lượng tử hóa sẵn giống như bất kỳ mô hình Diffusers nào khác:

Không cần lớp pipeline tùy chỉnh hay công cụ suy luận riêng biệt, và cũng không có gì cần phải biên dịch cục bộ. Các kernel NVFP4 sẽ được tải xuống từ Hub thông qua trang kernels của Nunchaku Lite trong lần đầu tiên sử dụng. Checkpoint này kết hợp transformer Nunchaku NVFP4 với bộ mã hóa văn bản (text encoder) bitsandbytes NF4, và tạo ra một hình ảnh 1024x1024 trong khoảng 1,7 giây trên RTX 5090 với mức sử dụng bộ nhớ đỉnh điểm khoảng 12 GB, so với khoảng 24 GB của pipeline BF16. Bạn có thể tìm thêm chi tiết về định dạng checkpoint Nunchaku Lite trong tài liệu chính thức của Diffusers.

Các checkpoint NVFP4 yêu cầu GPU NVIDIA Blackwell (RTX 50 series, RTX PRO 6000, B200). Đối với các thế hệ cũ hơn, hãy sử dụng các biến thể INT4. Xem bảng hỗ trợ phần cứng bên dưới để biết chi tiết.

Bối cảnh: SVDQuant và Nunchaku

SVDQuant là phương pháp lượng tử hóa đứng sau Nunchaku, công cụ suy luận CUDA tham chiếu của nó. Lượng tử hóa 4-bit tiêu chuẩn rất khó thực hiện đối với các diffusion transformer vì cả trọng số và kích hoạt đều chứa các giá trị ngoại lai (outliers) lớn. SVDQuant xử lý vấn đề này bằng cách chuyển các giá trị ngoại lai của kích hoạt vào trọng số, biểu diễn phần khó nhất của mỗi ma trận trọng số bằng một nhánh hạng thấp (low-rank) 16-bit nhỏ, và lượng tử hóa phần dư còn lại thành 4-bit. Nunchaku giúp quá trình này trở nên nhanh chóng với các kernel hợp nhất (fused kernels) cho đường dẫn 4-bit và nhánh hạng thấp.

Giới thiệu Nunchaku Lite

Công cụ Nunchaku gốc đạt được tốc độ cao phần lớn nhờ vào các đường dẫn thực thi hợp nhất dành riêng cho từng mô hình, chẳng hạn như các phép chiếu QKV hợp nhất và các kernel GELU/MLP hợp nhất. Những tối ưu hóa đó gắn liền với bố cục module và định dạng checkpoint của từng kiến trúc, vì vậy việc hỗ trợ một dòng mô hình mới thường đòi hỏi công việc tích hợp riêng cho mô hình đó.

Nunchaku Lite là con đường tích hợp mới trong Diffusers. Với nó, Diffusers có thể tải các checkpoint theo phong cách Nunchaku mà không cần pipeline tùy chỉnh hay công cụ suy luận riêng biệt. Về cơ bản, Nunchaku Lite vá các module nn.Linear liên quan của một mô hình Diffusers tiêu chuẩn bằng các lớp tuyến tính SVDQ/AWQ thời gian thực trước khi checkpoint được tải. Các kernel CUDA được lấy từ Hub thông qua gói kernels. Hai họ kernel được sử dụng:

Sự đánh đổi ở đây là, nếu không có các module và kernel hợp nhất dành riêng cho kiến trúc, Nunchaku Lite không thể đạt được tốc độ như công cụ Nunchaku gốc. Tuy nhiên, cách triển khai cơ bản này vẫn mang lại tốc độ tăng khoảng 30% trong khi vẫn giữ nguyên mức giảm VRAM.

Tải nguyên bản trong Diffusers

Nếu bạn đã từng sử dụng bitsandbytes hoặc torchao trong Diffusers, cơ chế này sẽ rất quen thuộc. Một kho lưu trữ mô hình Nunchaku Lite là một kho lưu trữ Diffusers thông thường. Điểm đặc biệt duy nhất là khối quantization_config bên trong file config.json của transformer:

Cấu hình này cho Diffusers biết những module nào đã được lượng tử hóa, lược đồ nào được sử dụng và lớp runtime Nunchaku Lite nào cần khởi tạo (SVDQW4A4Linear hoặc AWQW4A16Linear).

Vì mô hình đã lượng tử hóa giữ nguyên cấu trúc module chính xác của mô hình dày đặc (dense), mọi thứ ở hạ nguồn (bộ lập lịch, hook tải LoRA, offloading, torch.compile) đều nhìn thấy một mô hình Diffusers bình thường.

Hỗ trợ phần cứng

Nunchaku Lite sử dụng các biến thể kernel khác nhau tùy thuộc vào thế hệ GPU và độ chính xác của checkpoint:

Các GPU Volta và Hopper hiện không được hỗ trợ bởi các kernel 4-bit. Bộ lượng tử hóa sẽ xác thực khả năng CUDA của GPU tại thời điểm tải và đưa ra lỗi rõ ràng thay vì tạo ra kết quả không chính xác.

Đạt được tốc độ cao hơn và bộ nhớ thấp hơn

Nunchaku Lite có thể được kết hợp với các tối ưu hóa bộ nhớ và tốc độ khác của Diffusers.

torch.compile. Việc biên dịch transformer giúp cải thiện tốc độ end-to-end từ 1,35x lên 1,8x:

Bộ mã hóa văn bản đã lượng tử hóa. Transformer không phải là thành phần duy nhất chiếm nhiều bộ nhớ. Các bộ mã hóa văn bản như T5 hoặc Qwen3 có thể chiếm vài gigabyte. Việc lượng tử hóa thêm bộ mã hóa văn bản bằng bitsandbytes NF4 giúp giảm VRAM đỉnh điểm khoảng 22% trong thử nghiệm của chúng tôi.

Offloading. Các trình hỗ trợ offloading của Diffusers như enable_model_cpu_offload và enable_sequential_cpu_offload vẫn hoạt động như bình thường nếu bạn cần đưa pipeline vào một GPU nhỏ hơn.

Các điểm chuẩn (Benchmarks)

Tất cả các con số dưới đây được đo trên NVIDIA RTX PRO 6000 (Blackwell) ở độ phân giải 1024x1024 sử dụng rootonchair/ERNIE-Image-Turbo-nunchaku-lite-int4-bnb4-text-encoder.

Độ trễ và bộ nhớ end-to-end

Như đã trình bày ở trên, Nunchaku giảm VRAM đỉnh điểm tới 50% trong khi vẫn cải thiện độ trễ khoảng 30%. Độ trễ còn lại chủ yếu đến từ việc khởi chạy thêm các kernel, điều mà torch.compile có thể giảm thiểu, đưa toàn bộ pipeline xuống còn 1,68 giây, nhanh hơn 1,8 lần so với mức cơ sở BF16.

Chất lượng hình ảnh

Tự lượng tử hóa mô hình của riêng bạn