Suy luận LoRA nhanh cho Flux với Diffusers và PEFT

LoRA adapter cung cấp một mức độ tùy biến tuyệt vời cho các mô hình thuộc mọi hình dạng và kích thước. Khi nói đến tạo ảnh, chúng có thể trao quyền cho các mô hình với các kiểu khác nhau, các nhân vật khác nhau và nhiều hơn nữa. Đôi khi, chúng cũng có thể được tận dụng để giảm độ trễ suy luận. Do đó, tầm quan trọng của chúng là tối quan trọng, đặc biệt khi nói đến tùy chỉnh và tinh chỉnh các mô hình.

Trong bài đăng này, chúng tôi sử dụng mô hình Flux.1-Dev để tạo văn bản thành hình ảnh vì sự phổ biến và chấp nhận rộng rãi của nó, đồng thời làm thế nào để tối ưu hóa tốc độ suy luận của nó khi sử dụng LoRA (~2.3x). Nó có hơn 30 nghìn bộ điều hợp được đào tạo với nó (như đã báo cáo trên nền tảng Hugging Face Hub). Do đó, tầm quan trọng của nó đối với cộng đồng là rất lớn.

Lưu ý rằng mặc dù chúng tôi chứng minh khả năng tăng tốc với Flux, nhưng chúng tôi tin rằng công thức của chúng tôi đủ chung để có thể áp dụng cho các mô hình khác.

Nếu bạn nóng lòng muốn bắt đầu với mã, vui lòng xem kho mã đi kèm.

Mục lục

• Những trở ngại trong việc tối ưu hóa suy luận LoRA
• Công thức tối ưu hóa
• Suy luận LoRA được tối ưu hóa trên GPU tiêu dùng
• Kết luận

Những trở ngại trong việc tối ưu hóa suy luận LoRA

Khi phục vụ LoRA, việc hoán đổi nóng (hoán đổi và hoán đổi các LoRA khác nhau) chúng là điều phổ biến. LoRA thay đổi kiến trúc mô hình cơ sở. Ngoài ra, LoRA có thể khác nhau - mỗi LoRA có thể có các thứ hạng khác nhau và các lớp khác nhau mà chúng nhắm mục tiêu để điều chỉnh. Để tính đến các thuộc tính động này của LoRA, chúng ta phải thực hiện các bước cần thiết để đảm bảo các tối ưu hóa mà chúng ta áp dụng là mạnh mẽ.

Ví dụ: chúng ta có thể áp dụng torch.compile trên một mô hình được tải với một LoRA cụ thể để có được khả năng tăng tốc trên độ trễ suy luận. Tuy nhiên, ngay khi chúng ta hoán đổi LoRA với một LoRA khác (có khả năng có cấu hình khác), chúng ta sẽ gặp phải các sự cố biên dịch lại, gây ra tình trạng chậm lại trong quá trình suy luận.

Người ta cũng có thể hợp nhất các tham số LoRA vào các tham số mô hình cơ sở, chạy biên dịch và bỏ hợp nhất các tham số LoRA khi tải các tham số mới. Tuy nhiên, cách tiếp cận này sẽ lại gặp phải vấn đề biên dịch lại bất cứ khi nào quá trình suy luận được chạy, do các thay đổi tiềm ẩn ở cấp độ kiến trúc.

Công thức tối ưu hóa của chúng tôi có tính đến các tình huống nói trên để thực tế nhất có thể. Dưới đây là các thành phần chính trong công thức tối ưu hóa của chúng tôi:

• Flash Attention 3 (FA3)
• torch.compile
• Lượng tử hóa FP8 từ TorchAO
• Sẵn sàng hoán đổi nóng

Lưu ý rằng trong số những điều đã nói ở trên, lượng tử hóa FP8 bị mất mát nhưng thường cung cấp sự đánh đổi về tốc độ-bộ nhớ đáng gờm nhất. Mặc dù chúng tôi đã thử nghiệm công thức chủ yếu bằng GPU NVIDIA, nhưng nó cũng sẽ hoạt động trên GPU AMD.

Công thức tối ưu hóa

Trong các bài đăng trên blog trước đây của chúng tôi (bài đăng 1 và bài đăng 2), chúng tôi đã thảo luận về lợi ích của việc sử dụng ba thành phần đầu tiên trong công thức tối ưu hóa của chúng tôi. Áp dụng chúng từng cái một chỉ là một vài dòng mã:

from diffusers import DiffusionPipeline, TorchAoConfig
from diffusers.quantizers import PipelineQuantizationConfig
from utils.fa3_processor import FlashFluxAttnProcessor3_0
import torch

# lượng tử hóa bộ biến đổi Flux với FP8
pipe = DiffusionPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-dev", 
    torch_dtype=torch.bfloat16, 
    quantization_config=PipelineQuantizationConfig(
        quant_mapping={"transformer": TorchAoConfig("float8dq_e4m3_row")}
    )
).to("cuda")

# sử dụng Flash-attention 3
pipe.transformer.set_attn_processor(FlashFluxAttnProcessor3_0())

# sử dụng torch.compile()
pipe.transformer.compile(fullgraph=True, mode="max-autotune")

# thực hiện suy luận
pipe_kwargs = {
    "prompt": "Một con mèo đang cầm một tấm biển có dòng chữ xin chào thế giới",
    "height": 1024,
    "width": 1024,
    "guidance_scale": 3.5,
    "num_inference_steps": 28,
    "max_sequence_length": 512,
}

# lần đầu tiên sẽ chậm hơn, các lần chạy tiếp theo sẽ nhanh hơn
image = pipe(**pipe_kwargs).images[0]

Bộ xử lý FA3 đến từ đây.

Các vấn đề bắt đầu nổi lên khi chúng ta cố gắng hoán đổi và hoán đổi LoRA vào một bộ biến đổi khuếch tán đã được biên dịch (pipe.transformer) mà không kích hoạt quá trình biên dịch lại.

Thông thường, việc tải và dỡ LoRA sẽ yêu cầu biên dịch lại, điều này làm mất đi bất kỳ lợi thế tốc độ nào đạt được từ quá trình biên dịch. May mắn thay, có một cách để tránh sự cần thiết phải biên dịch lại. Bằng cách chuyển hotswap=True, diffusers sẽ giữ nguyên kiến trúc mô hình và chỉ trao đổi trọng số của chính bộ điều hợp LoRA, điều này không cần thiết phải biên dịch lại.

pipe.enable_lora_hotswap(target_rank=max_rank)
pipe.load_lora_weights(<tên-bộ-điều-hợp-lora1>)
# biên dịch *sau khi* tải LoRA đầu tiên
pipe.transformer.compile(mode="max-autotune", fullgraph=True)
image = pipe(**pipe_kwargs).images[0]
# kể từ thời điểm này, hãy tải các LoRA mới bằng `hotswap=True`
pipe.load_lora_weights(<tên-bộ-điều-hợp-lora2>, hotswap=True)
image = pipe(**pipe_kwargs).images[0]

(Để nhắc nhở, lần gọi đầu tiên đến pipe sẽ chậm vì torch.compile là trình biên dịch just-in-time. Tuy nhiên, các lệnh gọi tiếp theo sẽ nhanh hơn đáng kể.)

Điều này thường cho phép hoán đổi LoRA mà không cần biên dịch lại, nhưng có những hạn chế:

• Chúng ta cần cung cấp thứ hạng tối đa trong số tất cả các bộ điều hợp LoRA trước. Do đó, nếu chúng ta có một bộ điều hợp có thứ hạng 16 và một bộ điều hợp khác có 32, chúng ta cần chuyển max_rank=32.
• Các bộ điều hợp LoRA được hoán đổi nóng chỉ có thể nhắm mục tiêu đến cùng một lớp hoặc một tập hợp con các lớp mà LoRA đầu tiên nhắm mục tiêu.
• Việc nhắm mục tiêu bộ mã hóa văn bản chưa được hỗ trợ.

Để biết thêm thông tin về hoán đổi nóng trong Diffusers và các hạn chế của nó, hãy truy cập phần hoán đổi nóng của tài liệu.

Lợi ích của quy trình làm việc này trở nên rõ ràng khi chúng ta xem xét độ trễ suy luận mà không sử dụng quá trình biên dịch với hoán đổi nóng.

Những điểm chính:

• Tùy chọn “thông thường + biên dịch” cung cấp khả năng tăng tốc kha khá so với tùy chọn thông thường, nhưng nó gây ra các sự cố biên dịch lại, làm tăng thời gian thực hiện tổng thể. Trong các điểm chuẩn của chúng tôi, chúng tôi không trình bày thời gian biên dịch.
• Khi các vấn đề biên dịch lại được loại bỏ thông qua hoán đổi nóng (còn được gọi là tùy chọn “được tối ưu hóa”), chúng ta đạt được khả năng tăng tốc cao nhất.
• Trong tùy chọn “được tối ưu hóa”, lượng tử hóa FP8 được bật, điều này có thể dẫn đến mất chất lượng. Ngay cả khi không sử dụng FP8, chúng ta vẫn nhận được một lượng tăng tốc kha khá (tùy chọn “không_fp8”).
• Với mục đích trình diễn, chúng tôi sử dụng một nhóm gồm hai LoRA để hoán đổi nóng với quá trình biên dịch. Để có mã đầy đủ, vui lòng tham khảo kho mã đi kèm.

Công thức tối ưu hóa mà chúng tôi đã thảo luận cho đến nay giả định quyền truy cập vào một GPU mạnh mẽ như H100. Tuy nhiên, chúng ta có thể làm gì khi chúng ta bị giới hạn trong việc sử dụng GPU tiêu dùng như RTX 4090? Hãy cùng tìm hiểu.

Suy luận LoRA được tối ưu hóa trên GPU tiêu dùng

Flux.1-Dev (không có bất kỳ LoRA nào), sử dụng kiểu dữ liệu Bfloat16, cần ~33GB bộ nhớ để chạy. Tùy thuộc vào kích thước của mô-đun LoRA và không sử dụng bất kỳ tối ưu hóa nào, mức sử dụng bộ nhớ này có thể tăng hơn nữa. Nhiều GPU tiêu dùng như RTX 4090 chỉ có 24GB. Trong phần còn lại của phần này, chúng ta sẽ coi một máy RTX 4090 là bệ thử nghiệm của chúng ta.

Đầu tiên, để cho phép thực thi Flux.1-Dev từ đầu đến cuối, chúng ta có thể áp dụng CPU offloading, trong đó các thành phần không cần thiết để thực thi tính toán hiện tại được chuyển sang CPU để giải phóng thêm bộ nhớ tăng tốc. Làm như vậy cho phép chúng ta chạy toàn bộ quy trình trong ~22GB trong 35.403 giây trên RTX 4090. Bật quá trình biên dịch có thể giảm độ trễ xuống còn 31.205 giây (tăng tốc 1.12x). Về mã, nó chỉ là một vài dòng:

pipe = DiffusionPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16,
)
pipe.enable_model_cpu_offload()
# Thay vì biên dịch đầy đủ, chúng ta áp dụng biên dịch khu vực
# ở đây để tận dụng `fullgraph=True` và cũng để giảm
# thời gian biên dịch. Bạn có thể tìm thêm chi tiết ở đây:
# https://hf.co/docs/diffusers/main/en/optimization/fp16#regional-compilation
pipe.transformer.compile_repeated_blocks(fullgraph=True)
image = pipe(**pipe_kwargs).images[0]

Lưu ý rằng chúng ta không áp dụng lượng tử hóa FP8 ở đây vì nó không được hỗ trợ với CPU offloading và biên dịch (hỗ trợ chuỗi vấn đề). Do đó, chỉ cần áp dụng lượng tử hóa FP8 cho Flux Transformer là không đủ để giảm thiểu vấn đề cạn kiệt bộ nhớ. Trong trường hợp này, chúng ta đã quyết định xóa nó.

Do đó, để tận dụng sơ đồ lượng tử hóa FP8, chúng ta cần tìm cách thực hiện nó mà không cần CPU offloading. Đối với Flux.1-Dev, nếu chúng ta áp dụng thêm lượng tử hóa cho bộ mã hóa văn bản T5, chúng ta sẽ có thể tải và chạy toàn bộ quy trình trong 24GB. Dưới đây là so sánh kết quả có và không có bộ mã hóa văn bản T5 được lượng tử hóa (lượng tử hóa NF4 từ bitsandbytes).

Như chúng ta có thể nhận thấy trong hình trên, việc lượng tử hóa bộ mã hóa văn bản T5 không gây ra quá nhiều mất mát về chất lượng. Kết hợp bộ mã hóa văn bản T5 đã lượng tử hóa và Flux Transformer đã lượng tử hóa FP8 với torch.compile mang lại cho chúng ta kết quả khá hợp lý - 9.668 giây từ 32.27 giây (tăng tốc ~3.3x lớn) mà không làm giảm chất lượng đáng kể.

Có thể tạo hình ảnh với 24 GB VRAM ngay cả khi không lượng tử hóa bộ mã hóa văn bản T5, nhưng điều đó sẽ làm cho quy trình tạo của chúng ta trở nên phức tạp hơn một chút.

Chúng ta hiện có một cách để chạy toàn bộ quy trình Flux.1-Dev với lượng tử hóa FP8 trên RTX 4090. Chúng ta có thể áp dụng công thức tối ưu hóa đã thiết lập trước đó để tối ưu hóa suy luận LoRA trên cùng một phần cứng. Vì FA3 không được hỗ trợ trên RTX 4090, nên chúng ta sẽ tuân theo công thức tối ưu hóa sau với lượng tử hóa T5 mới được thêm vào:

• Lượng tử hóa FP8
• torch.compile
• Sẵn sàng hoán đổi nóng
• Lượng tử hóa T5 (với NF4)

Trong bảng dưới đây, chúng ta hiển thị các số độ trễ suy luận với các kết hợp khác nhau của các thành phần trên được áp dụng.

Ghi chú nhanh:

• Quá trình biên dịch cung cấp khả năng tăng tốc 2x lớn so với đường cơ sở.
• Các tùy chọn khác mang lại lỗi OOM ngay cả khi đã bật offloading.

Kết luận

Bài đăng này đã phác thảo một công thức tối ưu hóa để suy luận LoRA nhanh với Flux, chứng minh khả năng tăng tốc đáng kể. Cách tiếp cận của chúng tôi kết hợp Flash Attention 3, torch.compile và lượng tử hóa FP8 đồng thời đảm bảo khả năng hoán đổi nóng mà không gặp sự cố khi biên dịch lại. Trên các GPU cao cấp như H100, thiết lập được tối ưu hóa này cung cấp khả năng tăng tốc 2.23x so với đường cơ sở.

Đối với GPU tiêu dùng, cụ thể là RTX 4090, chúng tôi đã giải quyết các giới hạn về bộ nhớ bằng cách giới thiệu lượng tử hóa bộ mã hóa văn bản T5 (NF4) và tận dụng quá trình biên dịch khu vực. Công thức toàn diện này đạt được khả năng tăng tốc đáng kể 2.04x, làm cho suy luận LoRA trên Flux khả thi và hiệu quả ngay cả với VRAM hạn chế. Thông tin chi tiết chính là bằng cách quản lý cẩn thận quá trình biên dịch và lượng tử hóa, lợi ích của LoRA có thể được nhận ra đầy đủ trên các cấu hình phần cứng khác nhau.

Hy vọng rằng các công thức từ bài đăng này sẽ truyền cảm hứng cho bạn để tối ưu hóa các trường hợp sử dụng dựa trên LoRA của bạn, hưởng lợi từ suy luận nhanh chóng.

Tài nguyên

Dưới đây là danh sách các tài nguyên quan trọng mà chúng tôi đã trích dẫn trong suốt bài đăng này:

• Giới thiệu Flux Fast: Làm cho Flux hoạt động trên H100
• torch.compile và Diffusers: Hướng dẫn thực hành để đạt hiệu suất cao nhất
• Hướng dẫn LoRA trong Diffusers

Link bài viết gốc

• Tags:
• Ai
• July 23, 2025
• Huggingface.co