SDNQ

SDNQ (SD.Next Quantization) — это backend квантования, не требующий обучения, который поддерживает широкий диапазон dtypes весов (от int8 до 2-битных форматов, FP8 и другие низкобитовые плавающие форматы) с опциональной коррекцией SVD, вращением Хадамарда и квантованным matmul INT8/FP8 для более быстрого вывода. Работает на CUDA, ROCm, XPU, MPS и CPU.

Установите библиотеку sdnq, чтобы использовать этот бэкенд.

pip install sdnq

Загрузка предварительно квантованной контрольной точки

Конфигурация квантования предквантованного чекпоинта SDNQ хранится в его config.json, поэтому он загружается так же, как и любая другая модель, без аргумента quantization_config. Единственное требование — наличие установленного sdnq.

import torch
from diffusers import DiffusionPipeline

pipe = DiffusionPipeline.from_pretrained(
    "Disty0/Z-Image-Turbo-SDNQ-uint4-svd-r32",
    dtype=torch.bfloat16,
).to("cuda")  # or "mps", "xpu", "cpu"

image = pipe("a cat holding a sign that says hello").images[0]
image.save("output.png")

Установите DIFFUSERS_SDNQ_TRANSFORMERS=1 (до импорта Diffusers), если вы загружаете SDNQ-квантованную модель Transformers самостоятельно, например, квантованный текстовый энкодер или автономный компонент LLM в модульном конвейере. Это позволяет sdnq зарегистрироваться в Transformers, чтобы квантованная модель загружалась корректно. Эта переменная не требуется, если SDNQ-квантованный трансформер diffusers загружается вместе с ней, поскольку загрузка этого трансформера уже импортирует sdnq и регистрирует его. По умолчанию эта опция отключена, чтобы конвейеры, не использующие SDNQ, никогда не импортировали sdnq.

Квантование модели на лету

Передайте SDNQConfig в from_pretrained, чтобы квантовать модель при загрузке. Используйте PipelineQuantizationConfig, чтобы квантовать конкретные компоненты конвейера.

import torch
from diffusers import DiffusionPipeline, PipelineQuantizationConfig, SDNQConfig

pipeline_quant_config = PipelineQuantizationConfig(
    quant_mapping={"transformer": SDNQConfig(weights_dtype="uint4", use_svd=True)}
)
pipe = DiffusionPipeline.from_pretrained(
    "Tongyi-MAI/Z-Image-Turbo",
    quantization_config=pipeline_quant_config,
    dtype=torch.bfloat16,
).to("cuda")  # or "mps", "xpu", "cpu"

Или выполните квантование отдельного компонента модели напрямую.

import torch
from diffusers import ZImageTransformer2DModel, SDNQConfig

transformer = ZImageTransformer2DModel.from_pretrained(
    "Tongyi-MAI/Z-Image-Turbo",
    subfolder="transformer",
    quantization_config=SDNQConfig(weights_dtype="int8"),
    dtype=torch.bfloat16,
)

Все аргументы передаются в sdnq.SDNQConfig. Полный список поддерживаемых dtypes и опций см. в документации SDNQ.

Выбор слоев для квантования

Некоторые слои чувствительны к квантованию сильнее, чем другие, и пропуск чувствительных слоёв повышает качество при незначительных затратах памяти. SDNQ решает эту задачу из коробки: в нём есть встроенные списки пропусков (универсальный список и списки для конкретных моделей для многих архитектур), которые применяются автоматически, поэтому обычно чувствительные слои остаются неквантованными без какой-либо настройки. Для более точного управления SDNQ также может выбирать слои динамически, измеряя ошибку квантования для каждого слоя и пропуская те, где ошибка превышает порог потерь. Подробности см. в руководстве по динамическому квантованию SDNQ.

Более быстрый вывод с использованием квантованного matmul

По умолчанию SDNQ выполняет деквантизацию весов обратно в тип данных вычислений перед каждым matmul. Включение квантованного matmul позволяет выполнять умножение матриц непосредственно в низкопрецизионном типе данных (INT8/FP8/FP16), что быстрее на оборудовании с рабочей установкой Triton (CUDA, ROCm или XPU).

Включите его для загруженной модели с помощью apply_sdnq_options_to_model, который считывает существующую конфигурацию квантования модели и повторно применяет опцию на месте. Это также работает для предварительно заквантованных контрольных точек, для которых не передаётся SDNQConfig.

import torch
from diffusers import DiffusionPipeline
from sdnq.common import use_torch_compile as triton_is_available
from sdnq.loader import apply_sdnq_options_to_model

pipe = DiffusionPipeline.from_pretrained(
    "Disty0/Z-Image-Turbo-SDNQ-uint4-svd-r32",
    dtype=torch.bfloat16,
).to("cuda")  # or "mps", "xpu", "cpu"

if triton_is_available and (torch.cuda.is_available() or torch.xpu.is_available()):
    pipe.transformer = apply_sdnq_options_to_model(pipe.transformer, use_quantized_matmul=True)
    pipe.text_encoder = apply_sdnq_options_to_model(pipe.text_encoder, use_quantized_matmul=True)

torch.compile

Модели с квантованием SDNQ совместимы с torch.compile для дополнительного ускорения, в том числе в сочетании с group offloading.

transformer.compile(fullgraph=True)

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

Квантованные модели можно сериализовать с помощью save_pretrained и загрузить обратно без quantization_config.

transformer.save_pretrained("z-image-turbo-sdnq-int8")

Ресурсы

Update on GitHub