# Nunchaku Lite

Nunchaku Lite — это backend квантования для загрузки предварительно квантованных контрольных точек в Diffusers. Создайте совместимые контрольные точки с помощью [diffuse-compressor](https://github.com/rootonchair/diffuse-compressor). Он квантует и экспортирует трансформер, а затем упаковывает его в пайплайн Diffusers.

Nunchaku Lite основан на исходном движке инференса [Nunchaku](https://github.com/nunchaku-ai/nunchaku), библиотеке квантования [DeepCompressor](https://github.com/nunchaku-ai/deepcompressor) и статье [SVDQuant](https://arxiv.org/abs/2411.05007).

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

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

```bash
pip install -U kernels
```

Nunchaku Lite загружает свои ядра из репозитория [`rootonchair/nunchaku-lite-kernels`](https://huggingface.co/rootonchair/nunchaku-lite-kernels), издатель которого не является доверенным издателем ядер на Hub. При загрузке выполняется код с Hub, поэтому Diffusers требует явного подтверждения, установив `DIFFUSERS_TRUST_REMOTE_KERNELS=true`. Подробности — в разделе [Trusting remote kernels](../optimization/attention_backends#trusting-remote-kernels).

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

Загрузите предварительно квантованный пайплайн с помощью [from_pretrained()](/docs/diffusers/v0.43.0/ru/api/pipelines/overview#diffusers.DiffusionPipeline.from_pretrained), который читает конфигурацию квантования из `config.json`.

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

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

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

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

> [!TIP]
> Используйте `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)

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

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

| Метод | Точность | Размер группы | Примечания |
|---|---:|---:|---|
| `svdq_w4a4` | `nvfp4` | 16 | Использует ядра NVFP4 runtime с коррекцией низкого ранга SVDQ. |
| `svdq_w4a4` | `int4` | 64 | Использует ядра INT4 W4A4 с поправкой низкого ранга SVDQ. |
| `awq_w4a16` | `int4` | 64 | Использует ядра AWQ-стиля только для весов в формате INT4. |

## NunchakuLiteQuantizationConfig

Файл `config.json` должен содержать [NunchakuLiteQuantizationConfig](/docs/diffusers/v0.43.0/ru/api/quantization#diffusers.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`. Модуль может использовать только один метод, поэтому не указывайте один и тот же целевой объект под обоими.

```json
{
  "_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](https://github.com/nunchaku-ai/nunchaku) получает значительную часть своей скорости благодаря специфичным для модели слиянным путям выполнения. Он объединяет проекции Q, K и V с RMSNorm и RoPE, а для MLP использует слиянное ядро GELU. Nunchaku Lite, в свою очередь, использует стандартную модель Diffusers с универсальными квантованными линейными слоями, поэтому не включает эти слияния.

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

| Оптимизация | Задержка при выключенном | Задержка при включенном | Ускорение |
| --- | --- | --- | --- |
| Фьюжн QKV+norm+RoPE (комбинированный) | 1.057 s | 0.834 s | **1.27×** |
| — группировка Q/K/V в один matmul, по отдельности | 1.059 s | 1.034 s | 1.025× |
| — Фьюжн эпилога RMSNorm+RoPE поверх группировки | 1.034 s | 0.836 s | 1.236× |
| Слияние GELU-MLP | 0.766 s | 0.722 s | **1.06×** |

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

## torch.compile

Ядра Nunchaku Lite и квантованные линейные слои совместимы с [`torch.compile`](../optimization/fp16#torchcompile). Скомпилируйте квантованный трансформер после загрузки пайплайна для более быстрого вывода.

```python
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](https://github.com/rootonchair/diffuse-compressor)
- [Требования к установке Nunchaku](https://nunchaku.tech/docs/nunchaku/installation/installation.html)