Quanto

Бэкенд Quanto устарел и будет удалён в версии 1.0.0. Рекомендуется перейти на один из других поддерживаемых бэкендов квантования, например bitsandbytes или torchao.

Quanto — это бэкенд квантования для PyTorch в Optimum. Он создан с упором на универсальность и простоту:

Все функции доступны в режиме eager (работает с моделями, которые нельзя проследить)
Поддерживает обучение с учётом квантования
Квантованные модели совместимы с torch.compile
Квантованные модели не зависят от устройства (например, CUDA, XPU, MPS, CPU)

Чтобы использовать бэкенд Quanto, сначала необходимо установить optimum-quanto>=0.2.6 и accelerate

pip install optimum-quanto accelerate

Теперь вы можете квантовать модель, передав объект QuantoConfig в метод from_pretrained(). Хотя библиотека Quanto позволяет квантовать модули nn.Conv2d и nn.LayerNorm, в настоящее время Diffusers поддерживает квантование только весов в слоях nn.Linear модели. Приведённый ниже фрагмент кода демонстрирует, как применить квантование float8 с помощью Quanto.

import torch
from diffusers import FluxTransformer2DModel, QuantoConfig

model_id = "black-forest-labs/FLUX.1-dev"
quantization_config = QuantoConfig(weights_dtype="float8")
transformer = FluxTransformer2DModel.from_pretrained(
      model_id,
      subfolder="transformer",
      quantization_config=quantization_config,
      dtype=torch.bfloat16,
)

pipe = FluxPipeline.from_pretrained(model_id, transformer=transformer, dtype=dtype)
pipe.to("cuda")  # or "mps", "xpu", "cpu"

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

Пропуск квантования для конкретных модулей

С помощью аргумента modules_to_not_convert в QuantoConfig можно пропустить применение квантования для определённых модулей. Убедитесь, что модули, переданные в этот аргумент, соответствуют ключам модулей в state_dict.

import torch
from diffusers import FluxTransformer2DModel, QuantoConfig

model_id = "black-forest-labs/FLUX.1-dev"
quantization_config = QuantoConfig(weights_dtype="float8", modules_to_not_convert=["proj_out"])
transformer = FluxTransformer2DModel.from_pretrained(
      model_id,
      subfolder="transformer",
      quantization_config=quantization_config,
      dtype=torch.bfloat16,
)

Использование from_single_file с бэкендом Quanto

QuantoConfig совместим с ~FromOriginalModelMixin.from_single_file.

import torch
from diffusers import FluxTransformer2DModel, QuantoConfig

ckpt_path = "https://huggingface.co/black-forest-labs/FLUX.1-dev/blob/main/flux1-dev.safetensors"
quantization_config = QuantoConfig(weights_dtype="float8")
transformer = FluxTransformer2DModel.from_single_file(ckpt_path, quantization_config=quantization_config, dtype=torch.bfloat16)

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

Diffusers поддерживает сериализацию моделей Quanto с помощью метода ~ModelMixin.save_pretrained.

Требования к сериализации и загрузке различаются для моделей, квантованных непосредственно с помощью библиотеки Quanto, и моделей, квантованных с помощью Diffusers с использованием Quanto в качестве бэкенда. В настоящее время невозможно загрузить модели, квантованные непосредственно с помощью Quanto, в Diffusers с использованием ~ModelMixin.from_pretrained

import torch
from diffusers import FluxTransformer2DModel, QuantoConfig

model_id = "black-forest-labs/FLUX.1-dev"
quantization_config = QuantoConfig(weights_dtype="float8")
transformer = FluxTransformer2DModel.from_pretrained(
      model_id,
      subfolder="transformer",
      quantization_config=quantization_config,
      dtype=torch.bfloat16,
)
# save quantized model to reuse
transformer.save_pretrained("<your quantized model save path>")

# you can reload your quantized model with
model = FluxTransformer2DModel.from_pretrained("<your quantized model save path>")

Использование torch.compile с Quanto

В настоящее время бэкенд Quanto поддерживает torch.compile для следующих типов квантования:

int8-веса

import torch
from diffusers import FluxPipeline, FluxTransformer2DModel, QuantoConfig

model_id = "black-forest-labs/FLUX.1-dev"
quantization_config = QuantoConfig(weights_dtype="int8")
transformer = FluxTransformer2DModel.from_pretrained(
    model_id,
    subfolder="transformer",
    quantization_config=quantization_config,
    dtype=torch.bfloat16,
)
transformer = torch.compile(transformer, mode="max-autotune", fullgraph=True)

pipe = FluxPipeline.from_pretrained(
    model_id, transformer=transformer, dtype=dtype
)
pipe.to("cuda")  # or "mps", "xpu", "cpu"
images = pipe("A cat holding a sign that says hello").images[0]
images.save("flux-quanto-compile.png")

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

Веса

float8
int8
int4
int2

Update on GitHub