Nunchaku Lite

Nunchaku Lite é um backend de quantização para carregar checkpoints pré-quantizados em Diffusers. Crie checkpoints compatíveis com 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, a biblioteca de quantização DeepCompressor e o artigo SVDQuant.

Instale os kernels CUDA

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

pip install -U kernels

Nunchaku Lite carrega seus kernels do repositório 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 para mais detalhes.

Carregar um pipeline quantizado

Carregue o pipeline pré-quantizado com from_pretrained(), que lê a configuração de quantização de 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")

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.

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)

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.

NunchakuLiteQuantizationConfig

O arquivo config.json deve incluir um 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.

{
  "_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 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.

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 reproduza necessariamente 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. Compile o transformer quantizado após carregar o pipeline para uma inferência mais rápida.

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
• Requisitos de instalação do Nunchaku