torchao

torchao предоставляет высокопроизводительные dtypes и оптимизации на основе квантования и разреженности для инференса и обучения моделей PyTorch. Библиотека поддерживает любые модели в любой модальности, при условии, что они поддерживают загрузку с помощью Accelerate и содержат слои torch.nn.Linear.

Убедитесь, что Pytorch 2.5+ и torchao установлены с помощью команды ниже.

uv pip install -U torch torchao

Каждый dtype квантования доступен как отдельный экземпляр класса AOBaseConfig. Это обеспечивает более гибкие параметры конфигурации за счёт раскрытия большего числа доступных аргументов.

Передайте AOBaseConfig для типа квантования, например Int4WeightOnlyConfig, в TorchAoConfig внутри 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" квантует каждый слой на GPU по мере его загрузки. Это быстро, но временно требует дополнительной памяти GPU для исходных и квантованных весов. Если модель уже занимает большую часть памяти вашего GPU, загрузка может завершиться ошибкой нехватки памяти. В этом случае уберите device_map и переместите конвейер на 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 repository.

Схемы пост-тренировочного квантования FP8 в torchao эффективны для GPU с вычислительной мощностью не ниже 8.9 (RTX-4090, Hopper и т. д.). FP8 часто обеспечивает наилучший баланс между скоростью, потреблением памяти и качеством при генерации изображений и видео. Если ваш GPU совместим, мы рекомендуем использовать FP8 в сочетании с torch.compile.

Поддерживаемые типы квантования

torchao поддерживает квантование только весов, а также квантование весов и динамических активаций для int8, float3-float8 и uint1-uint7.

Квантование только весов хранит веса модели в определённом формате с низким числом бит, но выполняет вычисления в формате с более высокой точностью, например bfloat16. Это снижает требования к памяти, связанные с весами модели, но сохраняет пиковые значения памяти для вычислений активаций.

Динамическое квантование активаций хранит веса модели в формате с низким числом бит, одновременно квантуя активации на лету для экономии дополнительной памяти. Это снижает требования к памяти, связанные с весами модели, а также уменьшает накладные расходы памяти на вычисления активаций. Однако в некоторых случаях это может привести к снижению качества, поэтому рекомендуется тщательно тестировать различные модели.

См. официальную документацию torchao, чтобы лучше понять доступные методы квантования. Полный список вариантов конфигурации доступен здесь.

Некоторые популярные конфигурации квантования приведены ниже:

Сериализация и десериализация квантованных моделей

Чтобы сериализовать квантованную модель в заданном типе данных, сначала загрузите модель с нужным типом данных для квантования, а затем сохраните её с помощью метода 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 weight-only, не могут быть загружены напрямую и могут вызывать UnpicklingError при попытке загрузки моделей, но работают как ожидается при их сохранении. Чтобы обойти эту проблему, можно вручную загрузить state dict в модель. Однако имейте в виду, что для этого необходимо использовать weights_only=False в torch.load, поэтому выполнять это следует только в том случае, если веса получены из надежного источника.

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 поддерживается для PyTorch >= 2.6, как показано в примерах ниже.

Ресурсы

• API квантования TorchAO
• Примеры Diffusers-TorchAO