torchao

torchao cung cấp các dtype hiệu năng cao và các tối ưu hóa dựa trên lượng tử hóa và độ thưa cho việc suy luận và huấn luyện các mô hình PyTorch. Nó hỗ trợ bất kỳ mô hình nào trong bất kỳ phương thức nào, miễn là mô hình đó hỗ trợ tải bằng Accelerate và chứa các lớp torch.nn.Linear.

Hãy đảm bảo rằng Pytorch 2.5+ và torchao đã được cài đặt bằng lệnh bên dưới.

uv pip install -U torch torchao

Mỗi dtype lượng hóa có sẵn dưới dạng một instance riêng của lớp AOBaseConfig. Điều này cung cấp các tùy chọn cấu hình linh hoạt hơn bằng cách phơi bày nhiều đối số khả dụng hơn.

Truyền AOBaseConfig của một dtype lượng tử hóa, chẳng hạn Int4WeightOnlyConfig, vào TorchAoConfig trong from_pretrained().

import torch
from diffusers import DiffusionPipeline, PipelineQuantizationConfig, TorchAoConfig
from torchao.quantization import Int8WeightOnlyConfig

pipeline_quant_config = PipelineQuantizationConfig(
    quant_mapping={"transformer": TorchAoConfig(Int8WeightOnlyConfig(group_size=128, version=2))}
)
pipeline = DiffusionPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    quantization_config=pipeline_quant_config,
    dtype=torch.bfloat16,
    device_map="cuda"  # or "mps", "xpu", "cpu"
)

device_map="cuda" lượng tử hóa từng lớp trên GPU trong khi tải. Cách này nhanh nhưng tạm thời cần thêm bộ nhớ GPU cho trọng số gốc và trọng số đã lượng tử hóa. Nếu mô hình đã sử dụng phần lớn bộ nhớ GPU của bạn, quá trình tải có thể thất bại với lỗi hết bộ nhớ. Trong trường hợp đó, bỏ device_map và chuyển pipeline sang GPU sau khi tải:

pipeline = DiffusionPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    quantization_config=pipeline_quant_config,
    torch_dtype=torch.bfloat16,
)
pipeline.to("cuda")  # or "mps", "xpu", "cpu"

Nếu không có device_map, Diffusers sẽ lượng tử hóa các lớp trên CPU. Cách này chậm hơn nhưng tránh được sự tăng đột biến bộ nhớ GPU tạm thời trong quá trình lượng tử hóa. Để giảm thêm mức sử dụng bộ nhớ GPU, hãy dùng enable_model_cpu_offload() thay thế. Bạn cũng có thể lượng tử hóa các thành phần bổ sung, chẳng hạn như text encoder.

torch.compile

torchao hỗ trợ torch.compile, có thể tăng tốc suy luận chỉ với một dòng mã.

import torch
from diffusers import DiffusionPipeline, PipelineQuantizationConfig, TorchAoConfig
from torchao.quantization import Int4WeightOnlyConfig

pipeline_quant_config = PipelineQuantizationConfig(
    quant_mapping={"transformer": TorchAoConfig(Int4WeightOnlyConfig(group_size=128))}
)
pipeline = DiffusionPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    quantization_config=pipeline_quant_config,
    dtype=torch.bfloat16,
    device_map="cuda"  # or "mps", "xpu", "cpu"
)

pipeline.transformer.compile(transformer, mode="max-autotune", fullgraph=True)

Tham khảo bảng này để xem các chỉ số chuẩn về tốc độ suy luận và mức sử dụng bộ nhớ với Flux và CogVideoX. Các chỉ số chuẩn khác trên nhiều phần cứng khác nhau cũng có sẵn trong kho torchao.

Các phương pháp lượng tử hóa sau huấn luyện FP8 trong torchao hiệu quả trên các GPU có khả năng tính toán tối thiểu là 8.9 (RTX-4090, Hopper, v.v.). FP8 thường cung cấp sự đánh đổi tốt nhất về tốc độ, bộ nhớ và chất lượng khi tạo hình ảnh và video. Chúng tôi khuyến nghị kết hợp FP8 và torch.compile nếu GPU của bạn tương thích.

Các kiểu lượng hóa được hỗ trợ

torchao hỗ trợ lượng tử hóa chỉ trọng số và lượng tử hóa trọng số cùng kích hoạt động cho int8, float3-float8 và uint1-uint7.

Lượng tử hóa chỉ trọng số lưu trữ các trọng số của mô hình trong một kiểu dữ liệu bit thấp cụ thể nhưng thực hiện tính toán bằng một kiểu dữ liệu có độ chính xác cao hơn, chẳng hạn như bfloat16. Điều này làm giảm yêu cầu bộ nhớ từ các trọng số của mô hình nhưng vẫn giữ nguyên các đỉnh bộ nhớ cho việc tính toán kích hoạt.

Lượng tử hóa kích hoạt động (Dynamic activation quantization) lưu trữ trọng số mô hình trong một kiểu dữ liệu (dtype) số bit thấp, đồng thời lượng tử hóa các kích hoạt (activations) theo thời gian thực để tiết kiệm thêm bộ nhớ. Cách làm này làm giảm yêu cầu bộ nhớ từ trọng số mô hình, đồng thời cũng làm giảm chi phí bộ nhớ từ các phép tính kích hoạt. Tuy nhiên, điều này đôi khi có thể dẫn đến đánh đổi về chất lượng, do đó khuyến nghị kiểm tra kỹ lưỡng các mô hình khác nhau.

Tham khảo tài liệu chính thức của torchao để hiểu rõ hơn về các phương pháp lượng tử hóa có sẵn. Danh sách đầy đủ các tùy chọn cấu hình có sẵn tại đây.

Một số cấu hình lượng tử hóa phổ biến được sử dụng như sau:

Chuẩn hóa và khử chuẩn hóa các mô hình lượng tử hóa

Để serialize một model đã được lượng tử hóa trong một dtype cho trước, trước tiên hãy tải model với dtype lượng tử hóa mong muốn, sau đó lưu nó bằng phương thức save_pretrained().

import torch
from diffusers import AutoModel, TorchAoConfig
from torchao.quantization import Int8WeightOnlyConfig

quantization_config = TorchAoConfig(Int8WeightOnlyConfig())
transformer = AutoModel.from_pretrained(
    "black-forest-labs/Flux.1-Dev",
    subfolder="transformer",
    quantization_config=quantization_config,
    dtype=torch.bfloat16,
)
transformer.save_pretrained("/path/to/flux_int8wo", safe_serialization=False)

Để tải một mô hình lượng tử hóa đã được serialize, hãy sử dụng phương thức from_pretrained().

import torch
from diffusers import FluxPipeline, AutoModel

transformer = AutoModel.from_pretrained("/path/to/flux_int8wo", dtype=torch.bfloat16, use_safetensors=False)
pipe = FluxPipeline.from_pretrained("black-forest-labs/Flux.1-Dev", transformer=transformer, dtype=torch.bfloat16)
pipe.to("cuda")  # or "mps", "xpu", "cpu"

prompt = "A cat holding a sign that says hello world"
image = pipe(prompt, num_inference_steps=30, guidance_scale=7.0).images[0]
image.save("output.png")

Nếu bạn đang sử dụng torch<=2.6.0, một số phương pháp lượng tử hóa, chẳng hạn như uint4 weight-only, không thể được tải trực tiếp và có thể gây ra UnpicklingError khi cố gắng tải các mô hình, nhưng hoạt động như mong đợi khi lưu chúng. Để khắc phục vấn đề này, bạn có thể tải state dict thủ công vào mô hình. Tuy nhiên, hãy lưu ý rằng điều này yêu cầu sử dụng weights_only=False trong torch.load, do đó chỉ nên chạy nếu các trọng số được lấy từ một nguồn đáng tin cậy.

import torch
from accelerate import init_empty_weights
from diffusers import FluxPipeline, AutoModel, TorchAoConfig
from torchao.quantization import IntxWeightOnlyConfig

# Serialize the model
transformer = AutoModel.from_pretrained(
    "black-forest-labs/Flux.1-Dev",
    subfolder="transformer",
    quantization_config=TorchAoConfig(IntxWeightOnlyConfig(dtype=torch.uint4)),
    dtype=torch.bfloat16,
)
transformer.save_pretrained("/path/to/flux_uint4wo", safe_serialization=False, max_shard_size="50GB")
# ...

# Load the model
state_dict = torch.load("/path/to/flux_uint4wo/diffusion_pytorch_model.bin", weights_only=False, map_location="cpu")
with init_empty_weights():
    transformer = AutoModel.from_config("/path/to/flux_uint4wo/config.json")
transformer.load_state_dict(state_dict, strict=True, assign=True)

API AutoModel được hỗ trợ cho PyTorch >= 2.6 như minh họa trong các ví dụ bên dưới.

Tài nguyên

• API Lượng tử hóa TorchAO
• Ví dụ Diffusers-TorchAO