torchao

TorchAO 是 PyTorch 的架構最佳化庫。它提供用於推理和訓練的高效能資料型別、最佳化技術和核心，並支援與原生 PyTorch 功能（如 torch.compile、FullyShardedDataParallel (FSDP) 等）的組合。

在開始之前，請確保已安裝 Pytorch 2.5+ 和 TorchAO。

pip install -U torch torchao

透過將 TorchAoConfig 傳遞給 from_pretrained() 來量化模型（您也可以載入預量化模型）。這適用於任何模態中的任何模型，只要它支援使用 Accelerate 載入幷包含 torch.nn.Linear 層。

下面的示例僅將權重量化為 int8。

import torch
from diffusers import FluxPipeline, AutoModel, TorchAoConfig

model_id = "black-forest-labs/FLUX.1-dev"
dtype = torch.bfloat16

quantization_config = TorchAoConfig("int8wo")
transformer = AutoModel.from_pretrained(
    model_id,
    subfolder="transformer",
    quantization_config=quantization_config,
    torch_dtype=dtype,
)
pipe = FluxPipeline.from_pretrained(
    model_id,
    transformer=transformer,
    torch_dtype=dtype,
)
pipe.to("cuda")

# Without quantization: ~31.447 GB
# With quantization: ~20.40 GB
print(f"Pipeline memory usage: {torch.cuda.max_memory_reserved() / 1024**3:.3f} GB")

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

TorchAO 與 torch.compile 完全相容，這使其與其他量化方法區分開來。這使得只需一行程式碼即可輕鬆加速推理。

# In the above code, add the following after initializing the transformer
transformer = torch.compile(transformer, mode="max-autotune", fullgraph=True)

有關 Flux 和 CogVideoX 上的速度和記憶體基準測試，請參閱此處的表格。您還可以找到各種硬體的一些 torchao 基準測試資料。

torchao 中的 FP8 訓練後量化方案對計算能力至少為 8.9 的 GPU（RTX-4090、Hopper 等）有效。FP8 在生成影像和影片時通常能提供最佳的速度、記憶體和質量權衡。如果您的 GPU 相容，我們建議結合使用 FP8 和 torch.compile。

torchao 還透過 autoquant 支援自動量化 API。自動量化透過比較每種技術在所選輸入型別和形狀上的效能，來確定適用於模型的最佳量化策略。目前，這可以直接用於底層建模元件。Diffusers 將來也會公開自動量化配置選項。

TorchAoConfig 類接受三個引數

quant_type：一個字串值，表示以下量化型別之一。
modules_to_not_convert：一個模組完整/部分名稱列表，對於這些模組不應執行量化。例如，若要不對 FluxTransformer2DModel 的第一個塊執行任何量化，可以指定：modules_to_not_convert=["single_transformer_blocks.0"]。
kwargs：一個關鍵字引數字典，用於傳遞給根據 quant_type 呼叫的底層量化方法。

支援的量化型別

torchao 支援 int8、float3-float8 和 uint1-uint7 的僅權重（weight-only）量化以及權重和動態啟用（dynamic-activation）量化。

僅權重量化將模型權重儲存在特定的低位資料型別中，但使用更高精度的資料型別（如 bfloat16）執行計算。這降低了模型權重的記憶體需求，但保留了啟用計算的記憶體峰值。

動態啟用量化將模型權重儲存在低位資料型別中，同時即時量化啟用以節省額外記憶體。這降低了模型權重的記憶體需求，同時也降低了啟用計算的記憶體開銷。然而，這有時可能會以質量為代價，因此建議徹底測試不同的模型。

支援的量化方法如下：

類別	完整函式名稱	簡稱
整數量化	`int4_weight_only`, `int8_dynamic_activation_int4_weight`, `int8_weight_only`, `int8_dynamic_activation_int8_weight`	`int4wo`, `int4dq`, `int8wo`, `int8dq`
浮點 8 位量化	`float8_weight_only`, `float8_dynamic_activation_float8_weight`, `float8_static_activation_float8_weight`	`float8wo`, `float8wo_e5m2`, `float8wo_e4m3`, `float8dq`, `float8dq_e4m3`, `float8dq_e4m3_tensor`, `float8dq_e4m3_row`
浮點 X 位量化	`fpx_weight_only`	`fpX_eAwB`，其中 `X` 是位數 (1-7)，`A` 是指數位，`B` 是尾數位。約束：`X == A + B + 1`
無符號整數量化	`uintx_weight_only`	`uint1wo`, `uint2wo`, `uint3wo`, `uint4wo`, `uint5wo`, `uint6wo`, `uint7wo`

一些量化方法是別名（例如，int8wo 是 int8_weight_only 的常用簡稱）。這允許直接使用 torchao 文件中描述的量化方法，同時方便記憶它們的簡稱。

請參閱官方 torchao 文件，以更好地瞭解可用的量化方法和詳盡的配置選項列表。

序列化和反序列化量化模型

要以給定資料型別序列化量化模型，首先使用所需量化資料型別載入模型，然後使用 save_pretrained() 方法儲存它。

import torch
from diffusers import AutoModel, TorchAoConfig

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

要載入序列化的量化模型，請使用 from_pretrained() 方法。

import torch
from diffusers import FluxPipeline, AutoModel

transformer = AutoModel.from_pretrained("/path/to/flux_int8wo", torch_dtype=torch.bfloat16, use_safetensors=False)
pipe = FluxPipeline.from_pretrained("black-forest-labs/Flux.1-Dev", transformer=transformer, torch_dtype=torch.bfloat16)
pipe.to("cuda")

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

如果您使用的是 torch<=2.6.0，某些量化方法（如 uint4wo）無法直接載入，並且在嘗試載入模型時可能會導致 UnpicklingError，但在儲存它們時可以正常工作。為了解決這個問題，可以手動將狀態字典載入到模型中。但是請注意，這需要在 torch.load 中使用 weights_only=False，因此只有在權重來自可信來源時才應執行此操作。

import torch
from accelerate import init_empty_weights
from diffusers import FluxPipeline, AutoModel, TorchAoConfig

# Serialize the model
transformer = AutoModel.from_pretrained(
    "black-forest-labs/Flux.1-Dev",
    subfolder="transformer",
    quantization_config=TorchAoConfig("uint4wo"),
    torch_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)

對於 PyTorch >= 2.6，支援 AutoModel API，如下面的示例所示。

資源

< > 在 GitHub 上更新