# Nunchaku Lite

Nunchaku Lite é um backend de quantização para carregar checkpoints pré-quantizados em Diffusers. Crie checkpoints compatíveis com [diffuse-compressor](https://github.com/rootonchair/diffuse-compressor). Ele quantiza e exporta um transformer, em seguida, o empacota como um pipeline do Diffusers.

Nunchaku Lite é construído sobre o motor de inferência original [Nunchaku](https://github.com/nunchaku-ai/nunchaku), a biblioteca de quantização [DeepCompressor](https://github.com/nunchaku-ai/deepcompressor) e o [artigo SVDQuant](https://arxiv.org/abs/2411.05007).

## Instale os kernels CUDA

O pacote kernels fornece os kernels CUDA otimizados, que são carregados automaticamente. Instale-o primeiro.

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

Nunchaku Lite carrega seus kernels do repositório [`rootonchair/nunchaku-lite-kernels`](https://huggingface.co/rootonchair/nunchaku-lite-kernels), cujo publicador não é um publicador de kernels confiável no Hub. Carregá-lo baixa e executa código do Hub, portanto, o Diffusers exige que você opte explicitamente por isso definindo `DIFFUSERS_TRUST_REMOTE_KERNELS=true`. Veja [Confiança em kernels remotos](../optimization/attention_backends#trusting-remote-kernels) para mais detalhes.

## Carregar um pipeline quantizado

Carregue o pipeline pré-quantizado com [from_pretrained()](/docs/diffusers/v0.43.0/pt/api/pipelines/overview#diffusers.DiffusionPipeline.from_pretrained), que lê a configuração de quantização de `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]
> O state dict exportado deve corresponder exatamente à arquitetura do modelo Diffusers de destino. Por exemplo, um checkpoint quantizado com projeções QKV fundidas não será carregado em uma configuração de modelo que espera módulos de projeção Q, K e V separados.

## Tipos de quantização suportados

O Nunchaku Lite suporta os seguintes formatos de camadas lineares quantizadas.

> [!TIP]
> Use `nvfp4` em GPUs Blackwell. Executar checkpoints `int4` no Blackwell pode ser mais lento do que `nvfp4`.

Os kernels CUDA atualmente suportam as seguintes arquiteturas de GPU da NVIDIA:

- `sm_75` (Turing, por exemplo, RTX 2080)
- `sm_80` (Ampere, por exemplo, A100)
- `sm_86` (Ampere, por exemplo, RTX 3090 e RTX A6000)
- `sm_89` (Ada, por exemplo, RTX 4090)
- `sm_120` (Blackwell, por exemplo, RTX 5090)

> [!NOTE]
> As GPUs Hopper, como `sm_90` H100 e H200, não são suportadas no momento.

Os checkpoints `nvfp4` exigem uma GPU NVIDIA Blackwell ou mais recente. Em GPUs Blackwell, use PyTorch >= 2.7 com CUDA >= 12.8. Os checkpoints `int4` exigem uma GPU NVIDIA Turing ou mais recente.

| Método | Precisão | Tamanho do grupo | Observações |
|---|---:|---:|---|
| `svdq_w4a4` | `nvfp4` | 16 | Usa kernels de runtime NVFP4 com correção de baixo rank SVDQ. |
| `svdq_w4a4` | `int4` | 64 | Usa kernels INT4 W4A4 com correção de baixo posto SVDQ. |
| `awq_w4a16` | `int4` | 64 | Usa kernels AWQ-style apenas para pesos em INT4. |

## NunchakuLiteQuantizationConfig

O arquivo `config.json` deve incluir um [NunchakuLiteQuantizationConfig](/docs/diffusers/v0.43.0/pt/api/quantization#diffusers.NunchakuLiteQuantizationConfig). Ele define o `compute_dtype` de runtime e os módulos de destino para cada método de quantização Nunchaku Lite.

- `compute_dtype`: dtype de runtime para buffers de ponto flutuante em módulos quantizados, tipicamente `torch.bfloat16`.
- `svdq_w4a4`: configuração de destino SVDQ W4A4 com `precision`, `group_size`, `rank` e `targets`.
- `awq_w4a16`: configuração de destino AWQ W4A16 com `precision`, `group_size` e `targets`.

Cada entrada em `targets` deve apontar para uma camada linear. O Diffusers substitui cada alvo `svdq_w4a4` por uma camada SVDQ W4A4 e cada alvo `awq_w4a16` por uma camada AWQ W4A16. O exemplo abaixo mostra a forma esperada com listas de alvos encurtadas.

Liste cada módulo que deseja quantizar sob `svdq_w4a4` ou `awq_w4a16`. Um módulo pode usar apenas um método, portanto, não liste o mesmo destino em ambos.

```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"]
    }
  }
}
```

## Kernels fundidos

O motor original [Nunchaku](https://github.com/nunchaku-ai/nunchaku) obtém grande parte de sua velocidade de caminhos de execução fundidos específicos do modelo. Ele combina as projeções Q, K e V com RMSNorm e RoPE, e usa um kernel GELU fundido para o MLP. O Nunchaku Lite, por outro lado, usa o modelo padrão do Diffusers com camadas lineares quantizadas genéricas, portanto, não inclui essas fusões.

As medições a seguir mostram o impacto na latência por etapa dessas fusões no motor Nunchaku original, medido com Flux Schnell em uma RTX 5090.

| Otimização | Latência com ele DESLIGADO | Latência com ele ATIVADO | Aceleração |
| --- | --- | --- | --- |
| Fusão QKV+norm+RoPE (combinada) | 1.057 s | 0.834 s | **1.27×** |
| — agrupando Q/K/V em um único matmul, isoladamente | 1.059 s | 1.034 s | 1.025× |
| — fusão do epílogo RMSNorm+RoPE, além do agrupamento | 1.034 s | 0.836 s | 1.236× |
| Fusão GELU-MLP | 0.766 s | 0.722 s | **1.06×** |

Sem essas fusões, a GPU inicia as operações de projeção, normalização, incorporação rotativa e MLP como kernels separados, o que adiciona sobrecarga de lançamento. `torch.compile` pode reduzir parte dessa sobrecarga, embora não necessariamente reproduza os kernels fundidos do engine original. Realize o benchmark do pipeline compilado para o seu modelo e carga de trabalho.

## torch.compile

Os kernels do Nunchaku Lite e as camadas lineares quantizadas são compatíveis com [`torch.compile`](../optimization/fp16#torchcompile). Compile o transformer quantizado após carregar o pipeline para uma inferência mais rápida.

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

O pipeline Nunchaku Lite NVFP4 compilado executa 1,8x mais rápido que o pipeline BF16 original (2,271s → 1,675s em uma RTX PRO 6000).

## Recursos

- [diffuse-compressor](https://github.com/rootonchair/diffuse-compressor)
- [Requisitos de instalação do Nunchaku](https://nunchaku.tech/docs/nunchaku/installation/installation.html)