Hugging Face's logo

Diffusers documentation

torchao

Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

torchao

torchao 提供基于量化和稀疏性的高性能数据类型及优化方案，用于 PyTorch 模型的推理和训练。只要模型支持通过 Accelerate 加载，并且包含 torch.nn.Linear 层，即可支持任意模态下的任意模型。

请确保使用以下命令安装了 Pytorch 2.5+ 和 torchao。

uv pip install -U torch torchao

每种量化数据类型都作为 AOBaseConfig 类的一个独立实例提供。通过暴露更多可用参数，这提供了更灵活的配置选项。

在 from_pretrained() 中，将量化数据类型的 AOBaseConfig（例如 Int4WeightOnlyConfig）传递给 TorchAoConfig。

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" 会在加载时对每一层进行 GPU 量化。这种方式速度很快，但会临时占用额外的 GPU 内存来存储原始权重和量化后的权重。如果模型已经占用了大部分 GPU 内存，加载过程可能会因内存不足而失败。在这种情况下，请去掉 device_map，并在加载完成后将 pipeline 移至 GPU：

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"

如果不使用 device_map，Diffusers 会在 CPU 上对层进行量化。这种方式速度较慢，但可避免量化过程中 GPU 内存的临时峰值。若要进一步降低 GPU 内存占用，请改用 enable_model_cpu_offload()。此外，你还可以对文本编码器等其他组件进行量化。

torch.compile

torchao 支持 torch.compile，只需一行代码即可加速推理。

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)

有关 Flux 和 CogVideoX 的推理速度和内存使用基准测试，请参阅此表格。torchao 仓库中还提供了各种硬件上的更多基准测试。

torchao 中的 FP8 后训练量化方案适用于计算能力至少为 8.9 的 GPU（如 RTX-4090、Hopper 等）。在生成图像和视频时，FP8 通常能在速度、内存和质量之间提供最佳的平衡。如果您的 GPU 兼容，我们建议将 FP8 与 torch.compile 结合使用。

支持的量化类型

torchao 支持针对 int8、float3-float8 以及 uint1-uint7 的仅权重量化，以及权重与动态激活值量化。

仅权重量化以特定的低比特数据类型存储模型权重，但使用更高精度的数据类型（例如 bfloat16）执行计算。这降低了模型权重带来的内存需求，但保留了激活值计算时的内存峰值。

动态激活量化将模型权重以低比特数据类型存储，同时在对激活值进行即时量化以节省额外内存。这降低了模型权重带来的内存需求，同时也减少了激活值计算带来的内存开销。不过，这有时可能会带来质量上的权衡，因此建议对不同的模型进行充分测试。

请参阅 torchao 官方文档，以更好地了解可用的量化方法。完整的配置选项列表可在 此处 找到。

以下是一些常见的量化配置示例：

类别配置类
整数量化Int4WeightOnlyConfig, Int8WeightOnlyConfig, Int8DynamicActivationInt8WeightConfig
8 位浮点量化Float8WeightOnlyConfig, Float8DynamicActivationFloat8WeightConfig
无符号整数量化IntxWeightOnlyConfig

量化模型的序列化与反序列化

要在指定的 dtype 中序列化量化模型，请先以所需的量化 dtype 加载模型，然后使用 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)

要加载序列化的量化模型，请使用 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")

如果您使用的是 torch<=2.6.0，某些量化方法（例如 uint4 仅权重量化）无法直接加载，在尝试加载模型时可能会引发 UnpicklingError，但在保存时表现正常。为了绕过此问题，可以手动将 state dict 加载到模型中。但请注意，这要求在 torch.load 中使用 weights_only=False，因此仅当权重来自可信来源时才应执行此操作。

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)

如下示例所示，AutoModel API 支持 PyTorch >= 2.6。

资源

Update on GitHub

←Nunchaku Lite quanto→