Nunchaku Lite

Nunchaku Lite — это backend квантования для загрузки предварительно квантованных контрольных точек в Diffusers. Создайте совместимые контрольные точки с помощью diffuse-compressor. Он квантует и экспортирует трансформер, а затем упаковывает его в пайплайн Diffusers.

Nunchaku Lite основан на исходном движке инференса Nunchaku, библиотеке квантования DeepCompressor и статье SVDQuant.

Установка ядер CUDA

Пакет kernels предоставляет оптимизированные CUDA-ядра, которые загружаются автоматически. Установите его в первую очередь.

pip install -U kernels

Nunchaku Lite загружает свои ядра из репозитория rootonchair/nunchaku-lite-kernels, издатель которого не является доверенным издателем ядер на Hub. При загрузке выполняется код с Hub, поэтому Diffusers требует явного подтверждения, установив DIFFUSERS_TRUST_REMOTE_KERNELS=true. Подробности — в разделе Trusting remote kernels.

Загрузка квантованного конвейера

Загрузите предварительно квантованный пайплайн с помощью from_pretrained(), который читает конфигурацию квантования из config.json.

import torch
from diffusers import DiffusionPipeline

model_id = "rootonchair/ERNIE-Image-Turbo-nunchaku-lite-nvfp4"

pipe = DiffusionPipeline.from_pretrained(
    model_id, dtype=torch.bfloat16,
).to("cuda")  # or "mps", "xpu", "cpu"

prompt = "A modern red armchair in a quiet studio, soft window light, realistic product photography"
image = pipe(
    prompt=prompt,
    height=1024,
    width=1024,
    num_inference_steps=8,
    guidance_scale=1.0,
).images[0]
image.save("ernie-image-turbo-nunchaku-lite.png")

Экспортированный state dict должен точно соответствовать архитектуре целевой модели Diffusers. Например, чекпоинт, квантованный с использованием объединённых QKV-проекций, не загрузится в конфигурацию модели, ожидающую отдельные модули проекций Q, K и V.

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

Nunchaku Lite поддерживает следующие форматы квантованных линейных слоёв.

Используйте nvfp4 на GPU Blackwell. Запуск контрольных точек int4 на Blackwell может быть медленнее, чем nvfp4.

Ядра CUDA в настоящее время поддерживают следующие архитектуры GPU NVIDIA:

• sm_75 (Turing, например RTX 2080)
• sm_80 (Ampere, например A100)
• sm_86 (Ampere, например RTX 3090 и RTX A6000)
• sm_89 (Ada, например RTX 4090)
• sm_120 (Blackwell, например, RTX 5090)

Видеокарты Hopper, такие как sm_90 H100 и H200, в настоящее время не поддерживаются.

Чекпоинты nvfp4 требуют GPU NVIDIA Blackwell или новее. На GPU Blackwell используйте PyTorch >= 2.7 с CUDA >= 12.8. Чекпоинты int4 требуют GPU NVIDIA Turing или новее.

NunchakuLiteQuantizationConfig

Файл config.json должен содержать NunchakuLiteQuantizationConfig. Он определяет тип вычислений compute_dtype во время выполнения и целевые модули для каждого метода квантования Nunchaku Lite.

• compute_dtype: тип данных для буферов с плавающей точкой в квантованных модулях во время выполнения, как правило, torch.bfloat16.
• svdq_w4a4: целевая конфигурация SVDQ W4A4 с параметрами precision, group_size, rank и targets.
• awq_w4a16: целевая конфигурация AWQ W4A16 с параметрами precision, group_size и targets.

Каждая запись в targets должна указывать на линейный слой. Diffusers заменяет каждую цель svdq_w4a4 слоем SVDQ W4A4 и каждую цель awq_w4a16 слоем AWQ W4A16. Пример ниже показывает ожидаемую форму с сокращёнными списками целей.

Перечислите каждый модуль, который вы хотите квантовать, под svdq_w4a4 или awq_w4a16. Модуль может использовать только один метод, поэтому не указывайте один и тот же целевой объект под обоими.

{
  "_class_name": "ErnieImageTransformer2DModel",
  "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"]
    },
    "awq_w4a16": {
      "precision": "int4",
      "group_size": 64,
      "targets": ["final_linear"]
    }
  }
}

Слиянные ядра

Исходный движок Nunchaku получает значительную часть своей скорости благодаря специфичным для модели слиянным путям выполнения. Он объединяет проекции Q, K и V с RMSNorm и RoPE, а для MLP использует слиянное ядро GELU. Nunchaku Lite, в свою очередь, использует стандартную модель Diffusers с универсальными квантованными линейными слоями, поэтому не включает эти слияния.

Приведённые ниже замеры показывают влияние этих слияний на задержку на каждом шаге в исходном движке Nunchaku, измеренное с использованием Flux Schnell на RTX 5090.

Без этих слияний GPU запускает операции проекции, нормализации, ротационного эмбеддинга и MLP как отдельные ядра, что добавляет накладные расходы на запуск. torch.compile может снизить часть этих накладных расходов, хотя не обязательно воспроизводит слиянные ядра исходного движка. Проведите бенчмаркинг скомпилированного конвейера для вашей модели и рабочей нагрузки.

torch.compile

Ядра Nunchaku Lite и квантованные линейные слои совместимы с torch.compile. Скомпилируйте квантованный трансформер после загрузки пайплайна для более быстрого вывода.

pipe.transformer = torch.compile(pipe.transformer, mode="default", fullgraph=True)

Скомпилированный конвейер Nunchaku Lite NVFP4 работает в 1,8 раза быстрее исходного конвейера BF16 (2,271 с → 1,675 с на RTX PRO 6000).

Ресурсы

• diffuse-compressor
• Требования к установке Nunchaku