# torchao

[torchao](https://github.com/pytorch/ao) 提供基于量化和稀疏性的高性能数据类型及优化方案，用于 PyTorch 模型的推理和训练。只要模型支持通过 [Accelerate](https://hf.co/docs/accelerate/index) 加载，并且包含 `torch.nn.Linear` 层，即可支持任意模态下的任意模型。

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

```bash
uv pip install -U torch torchao
```

每种量化数据类型都作为 [AOBaseConfig](https://docs.pytorch.org/ao/main/api_reference/generated/torchao.core.config.AOBaseConfig.html) 类的一个独立实例提供。通过暴露更多可用参数，这提供了更灵活的配置选项。

在 [from_pretrained()](/docs/diffusers/v0.43.0/zh/api/models/overview#diffusers.ModelMixin.from_pretrained) 中，将量化数据类型的 `AOBaseConfig`（例如 [Int4WeightOnlyConfig](https://docs.pytorch.org/ao/main/api_reference/generated/torchao.quantization.Int4WeightOnlyConfig.html)）传递给 [TorchAoConfig](/docs/diffusers/v0.43.0/zh/api/quantization#diffusers.TorchAoConfig)。

```py
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：

```py
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()](/docs/diffusers/v0.43.0/zh/api/pipelines/overview#diffusers.DiffusionPipeline.enable_model_cpu_offload)。此外，你还可以对文本编码器等其他组件进行量化。

## torch.compile

torchao 支持 [torch.compile](../optimization/fp16#torchcompile)，只需一行代码即可加速推理。

```python
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 的推理速度和内存使用基准测试，请参阅此[表格](https://github.com/huggingface/diffusers/pull/10009#issue-2688781450)。torchao [仓库](https://github.com/pytorch/ao/tree/main/torchao/quantization#benchmarks)中还提供了各种硬件上的更多基准测试。

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

## 支持的量化类型

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

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

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

请参阅 [torchao 官方文档](https://docs.pytorch.org/ao/stable/index.html)，以更好地了解可用的量化方法。完整的配置选项列表可在 [此处](https://docs.pytorch.org/ao/main/workflows/inference.html#inference-workflows) 找到。

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

| **类别** | **配置类** |
|---|---|
| **整数量化** | [`Int4WeightOnlyConfig`](https://docs.pytorch.org/ao/stable/api_reference/generated/torchao.quantization.Int4WeightOnlyConfig.html), [`Int8WeightOnlyConfig`](https://docs.pytorch.org/ao/stable/api_reference/generated/torchao.quantization.Int8WeightOnlyConfig.html), [`Int8DynamicActivationInt8WeightConfig`](https://docs.pytorch.org/ao/stable/api_reference/generated/torchao.quantization.Int8DynamicActivationInt8WeightConfig.html) |
| **8 位浮点量化** | [`Float8WeightOnlyConfig`](https://docs.pytorch.org/ao/stable/api_reference/generated/torchao.quantization.Float8WeightOnlyConfig.html), [`Float8DynamicActivationFloat8WeightConfig`](https://docs.pytorch.org/ao/stable/api_reference/generated/torchao.quantization.Float8DynamicActivationFloat8WeightConfig.html) |
| **无符号整数量化** | [`IntxWeightOnlyConfig`](https://docs.pytorch.org/ao/stable/api_reference/generated/torchao.quantization.IntxWeightOnlyConfig.html) |

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

要在指定的 dtype 中序列化量化模型，请先以所需的量化 dtype 加载模型，然后使用 [save_pretrained()](/docs/diffusers/v0.43.0/zh/api/models/overview#diffusers.ModelMixin.save_pretrained) 方法将其保存。

```python
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()](/docs/diffusers/v0.43.0/zh/api/models/overview#diffusers.ModelMixin.from_pretrained) 方法。

```python
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`，因此仅当权重来自可信来源时才应执行此操作。

```python
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)
```

> [!TIP]
> 如下示例所示，[AutoModel](/docs/diffusers/v0.43.0/zh/api/models/auto_model#diffusers.AutoModel) API 支持 PyTorch >= 2.6。

## 资源

- [TorchAO 量化 API](https://docs.pytorch.org/ao/stable/index.html)
- [Diffusers-TorchAO 示例](https://github.com/sayakpaul/diffusers-torchao)