Nunchaku Lite

Nunchaku Lite est un backend de quantification pour charger des points de contrôle pré-quantifiés dans Diffusers. Créez des points de contrôle compatibles avec diffuse-compressor. Il quantifie et exporte un transformer, puis l’empaquète en tant que pipeline Diffusers.

Nunchaku Lite s’appuie sur le moteur d’inférence Nunchaku d’origine, la bibliothèque de quantification DeepCompressor et l’article SVDQuant.

Installer les noyaux CUDA

Le package kernels fournit les noyaux CUDA optimisés, qui sont chargés automatiquement. Installez-le en premier.

pip install -U kernels

Nunchaku Lite charge ses noyaux depuis le dépôt rootonchair/nunchaku-lite-kernels, dont l’éditeur n’est pas un éditeur de noyaux de confiance sur le Hub. Son chargement télécharge et exécute du code depuis le Hub, Diffusers exige donc que vous activiez explicitement cette option en définissant DIFFUSERS_TRUST_REMOTE_KERNELS=true. Consultez Faire confiance aux noyaux distants pour plus de détails.

Charger un pipeline quantifié

Chargez le pipeline préquantisé avec from_pretrained(), qui lit la configuration de quantification depuis 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")

Le dictionnaire d'état exporté doit correspondre exactement à l'architecture du modèle Diffusers cible. Par exemple, un point de contrôle quantifié avec des projections QKV fusionnées ne pourra pas être chargé dans une configuration de modèle qui s'attend à des modules de projection Q, K et V séparés.

Types de quantification pris en charge

Nunchaku Lite prend en charge les formats de couches linéaires quantifiées suivants.

Utilisez nvfp4 sur les GPU Blackwell. L’exécution de points de contrôle int4 sur Blackwell peut être plus lente que nvfp4.

Les noyaux CUDA prennent actuellement en charge les architectures de GPU NVIDIA suivantes :

• sm_75 (Turing, par exemple RTX 2080)
• sm_80 (Ampere, par exemple A100)
• sm_86 (Ampere, par exemple RTX 3090 et RTX A6000)
• sm_89 (Ada, par exemple RTX 4090)
• sm_120 (Blackwell, par exemple RTX 5090)

Les GPU Hopper, tels que sm_90 H100 et H200, ne sont pas pris en charge pour le moment.

Les points de contrôle nvfp4 nécessitent un GPU NVIDIA Blackwell ou plus récent. Sur les GPU Blackwell, utilisez PyTorch >= 2.7 avec CUDA >= 12.8. Les points de contrôle int4 nécessitent un GPU NVIDIA Turing ou plus récent.

NunchakuLiteQuantizationConfig

Le fichier config.json doit inclure un NunchakuLiteQuantizationConfig. Il définit le compute_dtype d’exécution et les modules cibles pour chaque méthode de quantification Nunchaku Lite.

• compute_dtype : type de données d'exécution pour les tampons à virgule flottante dans les modules quantifiés, généralement torch.bfloat16.
• svdq_w4a4 : configuration cible SVDQ W4A4 avec precision, group_size, rank et targets.
• awq_w4a16 : configuration cible AWQ W4A16 avec precision, group_size et targets.

Chaque entrée de targets doit pointer vers une couche linéaire. Diffusers remplace chaque cible svdq_w4a4 par une couche SVDQ W4A4 et chaque cible awq_w4a16 par une couche AWQ W4A16. L’exemple ci-dessous montre la forme attendue avec des listes de cibles raccourcies.

Listez chaque module que vous souhaitez quantiser sous svdq_w4a4 ou awq_w4a16. Un module ne peut utiliser qu’une seule méthode, ne listez donc pas la même cible sous les deux.

{
  "_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"]
    }
  }
}

Noirs fusionnés

Le moteur Nunchaku original tire une grande partie de sa vitesse de chemins d'exécution fusionnés spécifiques au modèle. Il combine les projections Q, K et V avec RMSNorm et RoPE, et utilise un noyau GELU fusionné pour le MLP. Nunchaku Lite, en revanche, utilise le modèle standard de Diffusers avec des couches linéaires quantifiées génériques, et n'inclut donc pas ces fusions.

Les mesures suivantes montrent l’impact de ces fusions sur la latence par étape dans le moteur Nunchaku d’origine, mesuré avec Flux Schnell sur une RTX 5090.

Sans ces fusions, le GPU lance les opérations de projection, de normalisation, d’embedding rotatif et de MLP comme des noyaux distincts, ce qui ajoute un surcoût de lancement. torch.compile peut réduire une partie de ce surcoût, bien qu’il ne reproduise pas nécessairement les noyaux fusionnés du moteur d’origine. Effectuez des benchmarks du pipeline compilé pour votre modèle et votre charge de travail.

torch.compile

Les noyaux Nunchaku Lite et les couches linéaires quantifiées sont compatibles avec torch.compile. Compilez le transformeur quantifié après le chargement du pipeline pour une inférence plus rapide.

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

Le pipeline Nunchaku Lite NVFP4 compilé s'exécute 1,8 fois plus vite que le pipeline BF16 d'origine (2,271 s → 1,675 s sur une RTX PRO 6000).

Ressources

• diffuse-compressor
• Exigences d'installation de Nunchaku