Hugging Face 的 LOGO

Diffusers 文档

torchao

Hugging Face's logo
加入 Hugging Face 社区

并获得增强的文档体验

开始使用

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 上更新

←gguf quanto→