SDNQ

SDNQ (SD.Next Quantization) est un backend de quantification sans entraînement qui prend en charge une large gamme de dtypes de poids (de int8 jusqu’au 2 bits, FP8 et d’autres formats flottants à faible nombre de bits), avec correction SVD, rotation de Hadamard et matmul INT8/FP8 quantifiés en option pour une inférence plus rapide. Il fonctionne sur CUDA, ROCm, XPU, MPS et CPU.

Installez la bibliothèque sdnq pour utiliser ce backend.

pip install sdnq

Charger un point de contrôle pré-quantifié

La configuration de quantisation d'un point de contrôle SDNQ pré-quantisé est stockée dans son config.json, il se charge donc comme n'importe quel autre modèle, sans avoir besoin de l'argument quantization_config. La seule exigence est d'avoir sdnq installé.

import torch
from diffusers import DiffusionPipeline

pipe = DiffusionPipeline.from_pretrained(
    "Disty0/Z-Image-Turbo-SDNQ-uint4-svd-r32",
    dtype=torch.bfloat16,
).to("cuda")  # or "mps", "xpu", "cpu"

image = pipe("a cat holding a sign that says hello").images[0]
image.save("output.png")

Définissez DIFFUSERS_SDNQ_TRANSFORMERS=1 (avant d’importer Diffusers) lorsque vous chargez un modèle Transformers quantisé SDNQ de manière isolée, par exemple un encodeur de texte quantisé ou un composant LLM autonome dans un pipeline modulaire. Cela permet à sdnq de s’enregistrer auprès de Transformers afin que le modèle quantisé soit chargé correctement. Vous n’en avez pas besoin lorsqu’un transformer diffusers quantisé SDNQ est chargé en même temps, car le chargement de ce transformer importe déjà sdnq et l’enregistre. Cette option est désactivée par défaut afin que les pipelines n’utilisant pas SDNQ n’importent jamais sdnq.

Quantifier un modèle à la volée

Passez un SDNQConfig à from_pretrained pour quantifier un modèle lors de son chargement. Utilisez PipelineQuantizationConfig pour quantifier des composants spécifiques du pipeline.

import torch
from diffusers import DiffusionPipeline, PipelineQuantizationConfig, SDNQConfig

pipeline_quant_config = PipelineQuantizationConfig(
    quant_mapping={"transformer": SDNQConfig(weights_dtype="uint4", use_svd=True)}
)
pipe = DiffusionPipeline.from_pretrained(
    "Tongyi-MAI/Z-Image-Turbo",
    quantization_config=pipeline_quant_config,
    dtype=torch.bfloat16,
).to("cuda")  # or "mps", "xpu", "cpu"

Ou quantifiez directement un composant unique du modèle.

import torch
from diffusers import ZImageTransformer2DModel, SDNQConfig

transformer = ZImageTransformer2DModel.from_pretrained(
    "Tongyi-MAI/Z-Image-Turbo",
    subfolder="transformer",
    quantization_config=SDNQConfig(weights_dtype="int8"),
    dtype=torch.bfloat16,
)

Tous les arguments sont transmis à sdnq.SDNQConfig. Reportez-vous à la documentation de SDNQ pour la liste complète des dtypes et options pris en charge.

Choisir les couches à quantiser

Certain couches sont plus sensibles à la quantification que d'autres, et l'exclusion des couches sensibles améliore la qualité à un coût mémoire négligeable. SDNQ gère cela nativement : il fournit des listes d'exclusion intégrées (une liste générique ainsi que des listes par modèle pour de nombreuses architectures) qui sont appliquées automatiquement, de sorte que les couches sensibles habituelles restent non quantifiées sans aucune configuration. Pour un contrôle plus fin, il peut également sélectionner dynamiquement les couches en mesurant l'erreur de quantification par couche et en excluant celles qui dépassent un seuil de perte. Consultez le guide de quantification dynamique SDNQ pour plus de détails.

Inférence plus rapide avec des produits matriciels quantifiés

Par défaut, SDNQ déquantise les poids vers le type de calcul avant chaque matmul. L’activation du matmul quantifié exécute à la place le matmul directement dans un type de précision réduite (INT8/FP8/FP16), ce qui est plus rapide sur le matériel disposant d’une installation Triton fonctionnelle (CUDA, ROCm ou XPU).

Activez-le sur un modèle chargé à l’aide de apply_sdnq_options_to_model, qui lit la configuration de quantification existante du modèle et réapplique l’option sur place. Cela fonctionne également pour les points de contrôle pré-quantifiés, où aucune SDNQConfig n’est transmise.

import torch
from diffusers import DiffusionPipeline
from sdnq.common import use_torch_compile as triton_is_available
from sdnq.loader import apply_sdnq_options_to_model

pipe = DiffusionPipeline.from_pretrained(
    "Disty0/Z-Image-Turbo-SDNQ-uint4-svd-r32",
    dtype=torch.bfloat16,
).to("cuda")  # or "mps", "xpu", "cpu"

if triton_is_available and (torch.cuda.is_available() or torch.xpu.is_available()):
    pipe.transformer = apply_sdnq_options_to_model(pipe.transformer, use_quantized_matmul=True)
    pipe.text_encoder = apply_sdnq_options_to_model(pipe.text_encoder, use_quantized_matmul=True)

torch.compile

Les modèles quantifiés SDNQ sont compatibles avec torch.compile pour des accélérations supplémentaires, y compris en combinaison avec le group offloading.

transformer.compile(fullgraph=True)

Enregistrer un modèle quantifié

Les modèles quantifiés peuvent être sérialisés avec save_pretrained et rechargés sans quantization_config.

transformer.save_pretrained("z-image-turbo-sdnq-int8")

Ressources

Update on GitHub