Nunchaku Lite в Diffusers: 4-битная генерация на потребительских GPU

Генерация изображений с помощью больших диффузионных трансформеров (Diffusion Transformers, DiT) переживает настоящий бум. Современные модели способны создавать фотореалистичные снимки, видео и даже аудио, но цена этого качества — колоссальные вычислительные ресурсы. Загрузка современной модели текстового синтеза в формате BF16 (half-precision) часто требует от 20 до 30 ГБ видеопамяти (VRAM). Для большинства пользователей с потребительскими видеокартами, такими как RTX 3090 или 4090, это становится серьезным барьером, ограничивающим доступ к передовым алгоритмам.

Квантование (quantization) стало ключевым решением этой проблемы. Оно позволяет значительно сократить объем памяти, необходимый для хранения весов и активаций нейросети, без критической потери качества. В экосистеме Hugging Face Diffusers уже интегрированы различные бэкенды квантования, такие как bitsandbytes, GGUF, torchao и Quanto. Однако большинство из них работают по принципу «только веса» (weight-only): веса хранятся в низком разрешении, но при вычислениях они расквантовываются обратно в высокое. Это экономит память, но не всегда ускоряет процесс генерации, а иногда даже добавляет небольшую задержку из-за накладных расходов на преобразование.

Здесь на сцену выходит Nunchaku Lite — новая интеграция, основанная на методе SVDQuant. В отличие от классических подходов, Nunchaku Lite использует 4-битные веса и активации (W4A4) для основных слоев трансформера. Это не только снижает потребление VRAM, но и реально ускоряет цикл денормализации (denoising loop). В этой статье мы подробно разберем, как работает эта технология, как ее установить, какие аппаратные требования она предъявляет и как запустить свои собственные квантованные модели прямо в Diffusers.

01Фундамент: SVDQuant и архитектура Nunchaku

Чтобы понять преимущества Nunchaku Lite, нужно заглянуть в основу метода SVDQuant. Стандартное 4-битное квантование сложно применить к диффузионным трансформерам из-за наличия крупных выбросов (outliers) как в весах, так и в активациях. Если просто усечь значения, качество изображения резко упадет.

SVDQuant решает эту проблему элегантным способом: выбросы активаций «перемещаются» в веса. Самая сложная часть каждой матрицы весов представляется с помощью небольшого ветвления с низкой ранговой матрицей (low-rank branch) в 16-битном формате, а оставшийся остаток квантуется до 4 бит. Это позволяет сохранить точность там, где она критична, и использовать низкую точность там, где это безопасно.

Движок Nunchaku делает этот процесс быстрым за счет использования «слитых» (fused) ядер CUDA. Он объединяет проекцию низкой ранговой матрицы с ядром квантования и проекцию обратной связи с ядром вычислений 4-бит. Это устраняет накладные расходы на доступ к памяти, которые обычно возникают при работе с 16-битными ветвями. Результатом является модель, которая работает быстрее и потребляет меньше памяти, чем стандартные BF16 аналоги.

02Что такое Nunchaku Lite и чем он отличается?

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

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

В основе лежат два семейства ядер, загружаемых из репозитория Hugging Face kernels:

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

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

03Нативная загрузка в Diffusers

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

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

Пример кода запуска

Для начала работы необходимо установить последние версии библиотек:

pip install -U diffusers transformers accelerate kernels bitsandbytes

Затем загрузите предварительно квантованный пайплайн, например, ERNIE-Image-Turbo:

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=4,
    guidance_scale=1.0,
    generator=torch.Generator("cuda").manual_seed(42),
).images[0]

image.save("output.png")

Этот чекпоинт сочетает трансформер Nunchaku NVFP4 с текстовым кодировщиком bitsandbytes NF4. На видеокарте RTX 5090 он генерирует изображение размером 1024x1024 примерно за 1,7 секунды, используя пиковую память около 12 ГБ. Для сравнения, стандартный пайплайн BF16 потребляет около 24 ГБ VRAM.

04Аппаратные требования и поддержка GPU

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

• NVFP4 (svdq_w4a4): Требует NVIDIA Blackwell (серия RTX 50, RTX PRO 6000, B200). Это нативная поддержка для новейших архитектур.
• INT4 (svdq_w4a4 и awq_w4a16): Поддерживается на архитектурах Turing, Ampere и Ada (RTX 30 и 40 серии, A100, L40S).

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

05Получение максимальной скорости и экономии памяти

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

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

Компиляция трансформера с помощью torch.compile значительно улучшает общее ускорение. В то время как базовое Nunchaku Lite дает ускорение 1.35x, добавление torch.compile повышает этот показатель до 1.8x. Это достигается за счет снижения накладных расходов на запуск дополнительных ядер.

pipe.transformer.compile(fullgraph=True)
# Или для более быстрой компиляции:
pipe.transformer.compile_repeated_blocks(fullgraph=True)

2. Квантование текстовых кодировщиков

Трансформер — не единственный компонент, потребляющий много памяти. Текстовые кодировщики, такие как T5 или Qwen3, могут занимать несколько гигабайт сами по себе. Дополнительное квантование текстового кодировщика с помощью bitsandbytes NF4 снижает пиковое потребление VRAM примерно на 22% в наших тестах.

3. Выгрузка на CPU (Offloading)

Если вам нужно уместить пайплайн на видеокарту с меньшим объемом памяти, стандартные помощники Diffusers, такие как enable_model_cpu_offload() и enable_sequential_cpu_offload(), работают как обычно и совместимы с Nunchaku Lite.

06Бенчмарки: цифры против BF16

Все приведенные ниже цифры были измерены на NVIDIA RTX PRO 6000 (Blackwell) при разрешении 1024x1024 с использованием модели rootonchair/ERNIE-Image-Turbo-nunchaku-lite-int4-bnb4-text-encoder.

Как видно из таблицы, Nunchaku Lite снижает пиковое потребление VRAM до 50% и одновременно улучшает задержку примерно на 30%. Оставшаяся задержка в основном связана с дополнительными запусками ядер, которые torch.compile успешно mitigates, снижая полную задержку пайплайна до 1.68 секунды (1.8x быстрее базового BF16).

07Квантование собственных моделей

Поддержка Nunchaku Lite в Diffusers архитектурно-независима. Инструментарий diffuse-compressor предоставляет полный рабочий процесс SVDQuant для моделей Diffusers: калибровка, квантование, упаковка и публикация.

Давайте рассмотрим пример квантования модели FLUX.2 Klein 4B. Этот процесс включает несколько шагов:

1. Инспекция модели

Генерический сканер проходит по модели и определяет, что будет квантовано. Совместимые линейные слои внутри стека трансформера становятся целями 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 плотных внешних линейных слоев.

2. Запуск квантования

Эта команда запускает SVDQuant на трансформере и записывает квантованный чекпоинт:

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.

3. Упаковка пайплайна Diffusers

Конвертер объединяет квантованный трансформер с другими компонентами базового пайплайна, записывает компактную конфигурацию 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

4. Проверка и публикация

После загрузки и проверки качества изображения вы можете отправить модель в Hub:

pipe.push_to_hub("your-name/your-model-nunchaku-lite-int4")

08Готовые чекпоинты

Чтобы начать работу прямо сейчас, вы можете использовать следующие репозитории:

• 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.
• OzzyGT/Krea_2_Turbo_nunchaku_lite_nvfp4: NVFP4 чекпоинт Krea 2 Turbo.
• lite-infer: коллекция других чекпоинтов Nunchaku Lite.

09Что это значит на практике

Интеграция Nunchaku Lite в Hugging Face Diffusers знаменует собой важный шаг к демократизации доступа к мощным диффузионным моделям. Раньше для работы с 4-битными моделями требовалось знание специфичных библиотек, ручная компиляция ядер и глубокое понимание архитектуры инференса. Теперь же пользователь может загрузить квантованную модель так же просто, как и обычную, используя стандартный интерфейс Diffusers.

Для разработчиков и энтузиастов ИИ это означает возможность запускать сложные модели, такие как FLUX или ERNIE, на видеокартах среднего класса (RTX 30/40 серии) с приемлемой скоростью и качеством. Экономия до 50% VRAM позволяет либо запускать модели на более дешевом оборудовании, либо использовать более высокие разрешения и более длинные циклы денормализации для улучшения детализации. А возможность кастомного квантования своих моделей через diffuse-compressor открывает путь к созданию собственных оптимизированных решений, готовых к публикации и использованию сообществом.

Хотя Nunchaku Lite не достигает абсолютного максимума скорости, который дает специализированный движок с ручными оптимизациями, он предлагает идеальный баланс между производительностью, совместимостью и простотой использования. Это делает его отличным выбором для большинства сценариев применения, где важна доступность и гибкость.

Источник: Hugging Face ↗