Интеграция 4-битного инференса диффузионных моделей с Nunchaku в библиотеку Diffusers

Большинство бэкендов квантования работают исключительно с весами. Это означает, что они сохраняют параметры модели в пониженной точности, но перед вычислениями распаковывают их обратно до исходного высокого уровня. Такой подход существенно сокращает потребление памяти, однако обычно не ускоряет инференс, а в некоторых случаях даже добавляет небольшую задержку. Метод SVDQuant, который лежит в основе популярного движка Nunchaku, предлагает совершенно иной подход. Он выполняет основные слои трансформера с 4-битными весами и активациями (W4A4), что не только снижает требования к памяти, но и ускоряет цикл деноизинга.

Раньше для использования подобных чекпоинтов требовалась отдельная библиотека инференса. Теперь же в актуальных версиях Diffusers загрузка чекпоинта Nunchaku выполняется так же просто, как и вызов метода from_pretrained(), причем никакой локальной компиляции CUDA не требуется благодаря пакету kernels. Кроме того, сопутствующий инструментарий diffuse-compressor позволяет самостоятельно квантовать новые архитектуры и публиковать их в виде стандартных репозиториев Diffusers.

Для начала работы необходимо установить зависимости. Потребуется свежая версия Diffusers и пакет ядер от Hugging Face:

pip install -U diffusers transformers accelerate kernels bitsandbytes

Затем можно загрузить предварительно квантованный пайплайн точно так же, как и любую другую модель в Diffusers:

import torch
from diffusers import ErnieImagePipeline

pipe = ErnieImagePipeline.from_pretrained(
    "lite-infer/ERNIE-Image-Turbo-nunchaku-lite-nvfp4_r32-bnb4-text-encoder",
    torch_dtype=torch.bfloat16,
).to("cuda")

image = pipe(
    prompt="A cinematic portrait of a red fox in a misty forest at sunrise, detailed fur, volumetric light",
    height=1024,
    width=1024,
    num_inference_steps=8,
    guidance_scale=1.0,
    generator=torch.Generator("cuda").manual_seed(42),
).images[0]
image.save("output.png")

В этом сценарии не требуется использовать собственный класс пайплайна или подключать отдельный движок инференса, а компилировать код на машине пользователя не нужно. Ядра NVFP4 скачиваются с платформы Hugging Face Hub через пакет kernels при первом же обращении к ним. Данный чекпоинт объединяет трансформер Nunchaku NVFP4 с текстовым энкодером bitsandbytes NF4 и позволяет генерировать изображение размером 1024×1024 пикселя примерно за 1,7 секунды на видеокарте RTX 5090. Пиковое потребление памяти при этом составляет около 12 ГБ по сравнению с примерно 24 ГБ для стандартного пайплайна в формате BF16. Подробные сведения о формате чекпоинтов Nunchaku Lite можно найти в официальной документации Diffusers.

Важно учитывать аппаратные ограничения: чекпоинты NVFP4 требуют наличия графического процессора NVIDIA на базе архитектуры Blackwell (видеокарты серии RTX 50, RTX PRO 6000, B200). Для более ранних поколений видеокарт следует использовать варианты с INT4. Подробная информация приведена в таблице аппаратной поддержки ниже.

Особенности метода SVDQuant и архитектура Nunchaku Lite

Стандартное 4-битное квантование представляет серьезную сложность для диффузионных трансформеров, поскольку и веса, и активации содержат большое количество крупных выбросов. Метод SVDQuant решает эту проблему путем переноса выбросов активаций в веса, представления наиболее сложных участков каждой весовой матрицы с помощью небольшой 16-битной низкоранговой ветви и квантования оставшегося остатка до 4 бит. Движок Nunchaku обеспечивает высокую скорость работы за счет использования слитных (fused) ядер для 4-битного пути и низкоранговой ветви.

Оригинальный движок Nunchaku достигает высокой производительности во многом благодаря специализированным для конкретных моделей путям выполнения, таким как слитные проекции QKV и объединенные ядра GELU/MLP. Подобные оптимизации тесно привязаны к структуре модулей и формату чекпоинтов каждой конкретной архитектуры, поэтому поддержка нового семейства моделей обычно требует индивидуальной интеграции.

Nunchaku Lite представляет собой новый подход к интеграции в библиотеку Diffusers. С его помощью Diffusers может загружать чекпоинты в стиле Nunchaku без привлечения кастомного пайплайна или отдельного внешнего движка. Под капотом Nunchaku Lite заменяет соответствующие модули nn.Linear в стандартной модели Diffusers на линейные слои SVDQ/AWQ, работающие во время выполнения, еще до загрузки самого чекпоинта. Сами ядра CUDA подгружаются из Hub через пакет kernels.

В работе задействованы два семейства ядер:

• svdq_w4a4: 4-битные веса и активации с низкоранговой коррекцией SVDQuant. Этот слой применяется для механизмов внимания (attention) и проекций MLP трансформера, на которые приходится практически вся вычислительная нагрузка. Он доступен в вариантах INT4 и NVFP4.
• awq_w4a16: 4-битные веса с 16-битными активациями. Используется для адаптивной нормализации и проекций модуляции, таких как FLUX adanorm_single / adanorm_zero или модули модуляции Qwen-Image. Эти слои ограничены пропускной способностью памяти и чувствительны к точности, поэтому схема AWQ здесь оптимальна для сохранения качества при одновременной экономии памяти и пространства.

Обратной стороной такого универсального подхода является то, что без специфичных для архитектуры слитных ядер и модулей Nunchaku Lite не может достичь того же прироста производительности, который демонстрирует оригинальный движок Nunchaku. Тем не менее, базовая реализация все равно обеспечивает прирост скорости примерно на 30% при сохранении того же уровня снижения потребления видеопамяти.

Для тех, кто уже работал с библиотеками bitsandbytes или torchao в экосистеме Diffusers, принципы работы будут интуитивно понятны. Репозиторий модели Nunchaku Lite является обычным репозиторием Diffusers. Единственное отличие заключается в наличии блока quantization_config внутри файла config.json конфигурации трансформера:

"quantization_config": {
    "quant_method": "nunchaku_lite",
    "compute_dtype": "bfloat16",
    "svdq_w4a4": {
        "precision": "nvfp4",
        "group_size": 16,
        "rank": 32,
        "targets": [
            "layers.0.self_attention.to_q",
            "layers.0.self_attention.to_k",
            "..."
        ]
    },
    "awq_w4a16": {
        "precision": "int4",
        "group_size": 64,
        "targets": [
            "adaLN_modulation.1",
            "..."
        ]
    }
}

Эта конфигурация сообщает библиотеке Diffusers, какие именно модули были подвергнуты квантованию, какая схема при этом использовалась и какой именно класс слоя времени выполнения Nunchaku Lite следует инициализировать (SVDQW4A4Linear или AWQW4A16Linear). Поскольку квантованная модель сохраняет точную модульную структуру плотной модели, все последующие компоненты (планировщики, хуки загрузки LoRA, механизмы выгрузки, torch.compile) воспринимают ее как стандартную модель Diffusers.

Аппаратная поддержка и совместимость ядер

Nunchaku Lite задействует различные варианты ядер в зависимости от поколения графического процессора и точности чекпоинта:

• svdq_w4a4 (точность: nvfp4) — поддерживаемые GPU: Blackwell (серия RTX 50, RTX PRO 6000, B200)
• svdq_w4a4 (точность: int4) — поддерживаемые GPU: Turing / Ampere / Ada (серии RTX 30 и 40, A100, L40S)
• awq_w4a16 (точность: int4) — поддерживаемые GPU: Turing / Ampere / Ada (серии RTX 30 & 40, A100, L40S)

Видеокарты поколений Volta и Hopper в настоящее время 4-битными ядрами не поддерживаются. Модуль квантования проверяет возможности CUDA графического процессора в момент загрузки модели и выдает понятную ошибку вместо того, чтобы генерировать некорректный результат.

Совмещение с другими оптимизациями Diffusers

Nunchaku Lite можно успешно комбинировать с другими инструментами оптимизации памяти и производительности в Diffusers:

• torch.compile. Компиляция трансформера позволяет увеличить сквозной прирост производительности с 1,35x до 1,8x с помощью вызова pipe.transformer.compile(fullgraph=True) или оптимизированного варианта pipe.transformer.compile_repeated_blocks(fullgraph=True) для более быстрой компиляции.
• Квантованные текстовые энкодеры. Трансформер — не единственный компонент с большими требованиями к памяти. Текстовые энкодеры вроде T5 или Qwen3 могут самостоятельно занимать несколько гигабайт. Дополнительное квантование текстового энкодера с помощью алгоритма bitsandbytes NF4 снижает пиковое потребление видеопамяти примерно на 22% по результатам бенчмарков.
• Выгрузка (Offloading). Вспомогательные функции выгрузки в Diffusers, такие как enable_model_cpu_offload() и enable_sequential_cpu_offload(), работают в штатном режиме, если модель необходимо запустить на менее мощной видеокарте.

Все показатели производительности, приведенные ниже, были замерены на видеокарте NVIDIA RTX PRO 6000 (Blackwell) при разрешении 1024×1024 с использованием чекпоинта rootonchair/ERNIE-Image-Turbo-nunchaku-lite-int4-bnb4-text-encoder:

• BF16 baseline: полный пайплайн — 3,00 с, цикл деноизинга — 2,86 с, пиковая VRAM — 31,1 ГБ, ускорение — 1,0x
• Nunchaku Lite NVFP4: полный пайплайн — 2,27 с, цикл деноизинга — 2,13 с, пиковая VRAM — 20,6 ГБ, ускорение — 1,35x
• Nunchaku Lite NVFP4 + torch.compile: полный пайплайн — 1,68 с, цикл деноизинга — 1,53 с, пиковая VRAM — 20,6 ГБ, ускорение — 1,8x
• Nunchaku Lite NVFP4 + NF4 text encoder: полный пайплайн — 2,29 с, цикл деноизинга — 2,13 с, пиковая VRAM — 16,0 ГБ, ускорение — 1,35x

Как показывают данные, Nunchaku сокращает пиковое потребление видеопамяти почти наполовину (до 50%), одновременно повышая скорость работы примерно на 30%. Оставшиеся накладные расходы в основном обусловлены вызовами дополнительных ядер, с чем эффективно справляется компилятор torch.compile, сокращая время работы полного пайплайна до 1,68 секунды, что в 1,8 раза быстрее базового варианта BF16.

Инструментарий diffuse-compressor и квантование новых архитектур

Поддержка Nunchaku Lite в библиотеке Diffusers не привязана к конкретным архитектурам, а сопутствующий инструментарий diffuse-compressor предоставляет полноценный рабочий процесс SVDQuant для моделей Diffusers: калибровку, квантование, упаковку и публикацию. В качестве примера рассмотрим процесс квантования модели FLUX.2 Klein 4B, который охватывает ключевые этапы: анализ структуры модели, калибровку и квантование трансформера, упаковку результата в пайплайн Diffusers, а также проверку и отправку на Hugging Face Hub.

Специальный универсальный сканер обходит модель и определяет целевые объекты: совместимые линейные слои внутри стека повторяющихся блоков трансформера становятся целями для SVDQ W4A4, распознаваемые линейные слои модуляции преобразуются в цели AWQ W4A16, а все остальные элементы остаются в исходной плотной точности.

python examples/text_to_image/quantize_hf.py black-forest-labs/FLUX.2-klein-4B \
    --precision int4 --rank 32 --inspect-config

Рекомендуется всегда изучать этот отчет перед запуском квантования. Для модели FLUX.2 Klein 4B ожидаемым результатом является обнаружение 100 целей SVDQ, 3 целей AWQ и 6 внешних плотных линейных слоев при отсутствии пропущенных паттернов или дублирующихся имен. Следующая команда запускает процесс SVDQuant для трансформера и записывает квантованный чекпоинт в файл outputs/checkpoints/svdq-int4_r32-flux-2-klein-4b.safetensors:

python examples/text_to_image/quantize_hf.py black-forest-labs/FLUX.2-klein-4B \
    --precision int4 \
    --output outputs/checkpoints/svdq-int4_r32-flux-2-klein-4b.safetensors

Флаг --precision int4 можно заменить на nvfp4 для сборки весов, оптимизированных под архитектуру Blackwell. Конвертер объединяет квантованный трансформер с остальными компонентами базового пайплайна, записывает компактную конфигурацию nunchaku_lite в файл transformer/config.json и при необходимости выполняет конвертацию текстовых энкодеров в формат NF4:

python examples/convert_nunchaku_lite_diffusers.py \
    --checkpoint outputs/checkpoints/svdq-int4_r32-flux-2-klein-4b.safetensors \
    --model-id black-forest-labs/FLUX.2-klein-4B \
    --bnb4-text-encoder text_encoder \
    --compute-dtype bfloat16 \
    --output-dir outputs/diffusers/FLUX.2-klein-4B-nunchaku-lite-int4-bnb4-text-encoder

Затем можно проверить работоспособность полученного пайплайна:

import torch
from diffusers import DiffusionPipeline

pipe = DiffusionPipeline.from_pretrained(
    "outputs/diffusers/FLUX.2-klein-4B-nunchaku-lite-int4-bnb4-text-encoder",
    device_map="cuda",
)

image = pipe(
    "A glass robot in a greenhouse, cinematic lighting",
    num_inference_steps=4,
    guidance_scale=1.0,
    generator=torch.Generator("cuda").manual_seed(12345),
).images[0]

Убедившись в корректности результатов, можно опубликовать модель на платформе с помощью команды pipe.push_to_hub("your-name/your-model-nunchaku-lite-int4"). После этого другие пользователи смогут загружать ее стандартным методом from_pretrained().

Структурные переписывания и ограничения универсального пути

Стоит отметить, что универсальный путь предполагает возможность квантования архитектуры без структурных изменений. Для достижения максимальной производительности оригинальный движок Nunchaku объединяет группы слоев Diffusers в слитные модули. Универсальный сценарий не способен самостоятельно выявить подобные изменения — например, объединение раздельных проекций Q, К и V в один модуль или разделение слитной проекции на несколько составляющих.

Ярким примером служат проекции QKV в модели FLUX.1-dev. В библиотеке Diffusers для них определены три независимых модуля:

self.to_q = torch.nn.Linear(query_dim, self.inner_dim, bias=bias)
self.to_k = torch.nn.Linear(query_dim, self.inner_dim, bias=bias)
self.to_v = torch.nn.Linear(query_dim, self.inner_dim, bias=bias)

В то же время модуль Nunchaku FLUX объединяет эти слои в единый квантованный модуль to_qkv:

to_qkv = fuse_linears([other.to_q, other.to_k, other.to_v])
self.to_qkv = SVDQW4A4Linear.from_linear(to_qkv, **kwargs)

Подобный объединенный модуль необходим, поскольку слитный оператор Nunchaku обрабатывает проекцию QKV, нормализацию Q/K и вращательные встраивания (rotary embeddings) одновременно. Для сравнения, стандартный путь в Diffusers выполняет эти операции раздельно:

query = attn.to_q(hidden_states)
key = attn.to_k(hidden_states)
value = attn.to_v(hidden_states)

query = query.unflatten(-1, (attn.heads, -1))
key = key.unflatten(-1, (attn.heads, -1))
value = value.unflatten(-1, (attn.heads, -1))

query = attn.norm_q(query)
key = attn.norm_k(key)

if image_rotary_emb is not None:
    query = apply_rotary_emb(query, image_rotary_emb, sequence_dim=1)
    key = apply_rotary_emb(key, image_rotary_emb, sequence_dim=1)

В то же время путь Nunchaku передает сгруппированную проекцию, модули нормализации и вращательные встраивания в один слитный оператор:

qkv = fused_qkv_norm_rottary(
    hidden_states,
    attn.to_qkv,
    attn.norm_q,
    attn.norm_k,
    image_rotary_emb
)

Это и есть та самая структурная модификация, которую универсальный механизм выполнить автоматически не может. В Diffusers имеются три целевых модуля с префиксами параметров to_q, to_k и to_v, в то время как Nunchaku использует один объединенный модуль to_qkv. Конфигурация целей, специфичная для конкретной модели, или специальный адаптер должны указывать, что параметры Q, K и V должны быть конкатенированы вдоль выходной размерности именно в таком порядке и загружены в to_qkv.

Подобные структурные преобразования описываются специальным файлом конфигурации целевых объектов для конкретной модели на этапе квантования и обрабатываются небольшим runtime-адаптером при загрузке чекпоинта. Скрипт квантования FLUX.2 Klein 4B служит наглядным примером целевой конфигурации для создания чекпоинта со структурными изменениями, а репозиторий rootonchair/nunchaku-lite предоставляет необходимые адаптеры времени выполнения для загрузки сгруппированных тензоров QKV, разделения слитных проекций и выполнения прочих объединенных операций. Полное руководство по процессу доступно в инструкции по добавлению новой модели.

Доступные ресурсы и репозитории

Для тех, кто хочет опробовать технологию незамедлительно, подготовлен ряд готовых репозиториев:

• rootonchair/ERNIE-Image-Turbo-nunchaku-lite-int4-bnb4-text-encoder: модель INT4 ERNIE-Image-Turbo с текстовым энкодером bitsandbytes NF4.
• rootonchair/ERNIE-Image-Turbo-nunchaku-lite-nvfp4-bnb4-text-encoder: модель NVFP4 ERNIE-Image-Turbo с текстовым энкодером bitsandbytes NF4.
• OzzyGT/Krea_2_Turbo_nunchaku_lite_nvfp4: чекпоинт NVFP4 Krea 2 Turbo.
• lite-infer: дополнительные чекпоинты и коллекции Nunchaku Lite.

Ядра SVDQuant от Nunchaku представляют собой один из наиболее эффективных способов продуктивного запуска диффузионных трансформеров на потребительском «железе», и теперь они поддерживаются в Diffusers на нативном уровне. Предварительно квантованные чекпоинты загружаются вызовом from_pretrained(), а инструментарий diffuse-compressor позволяет квантовать новые архитектуры без необходимости ожидания официальной поддержки на уровне движка. Квантование как весов, так и активаций по схеме W4A4 снижает требования к памяти и уменьшает задержку деноизинга, сохраняя качество генерируемых изображений максимально близким к оригиналу в формате BF16.

Авторы разработки выражают признательность мейнтейнерам Diffusers за рецензии и поддержку в процессе интеграции, а также команде MIT HAN Lab / Nunchaku за оригинальные наработки по методу SVDQuant. Отдельная благодарность выражается Марку Сану (Marc Sun) за отзывы о материале блога, Альваро Сомозе (Álvaro Somoza) за тестирование nunchaku-lite и ценные комментарии, а также компании SilverAI за поддержку проекта и предоставленную среду для проведения большей части разработческих тестов.

Источник: huggingface.co