SDNQ

SDNQ (SD.Next Quantization) é um backend de quantização sem treinamento que suporta uma ampla variedade de dtypes de pesos (int8 até 2-bit, FP8 e outros formatos de ponto flutuante de baixo bit) com correção SVD opcional, rotação de Hadamard e matmul quantizado INT8/FP8 para inferência mais rápida. Ele roda em CUDA, ROCm, XPU, MPS e CPU.

Instale a biblioteca sdnq para usar este backend.

pip install sdnq

Carregar um checkpoint pré-quantizado

A configuração de quantização de um checkpoint SDNQ pré-quantizado é armazenada em seu config.json, portanto ele é carregado como qualquer outro modelo, sem a necessidade do argumento quantization_config. O único requisito é ter o sdnq instalado.

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

Defina DIFFUSERS_SDNQ_TRANSFORMERS=1 (antes de importar o Diffusers) ao carregar um modelo Transformers quantizado com SDNQ de forma isolada, como um codificador de texto quantizado ou um componente LLM independente em um pipeline modular. Isso permite que o sdnq se registre no Transformers para que o modelo quantizado seja carregado corretamente. Você não precisa disso quando um transformer do diffusers quantizado com SDNQ é carregado junto, pois o carregamento desse transformer já importa o sdnq e o registra. Essa opção está desativada por padrão para que pipelines que não usam SDNQ nunca importem o sdnq.

Quantize um modelo em tempo de execução

Passe um SDNQConfig para from_pretrained para quantizar um modelo durante o carregamento. Use PipelineQuantizationConfig para quantizar componentes específicos do 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 quantize um único componente do modelo diretamente.

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,
)

Todos os argumentos são encaminhados para sdnq.SDNQConfig. Consulte a documentação do SDNQ para a lista completa de dtypes e opções suportadas.

Escolhendo quais camadas quantizar

Algumas camadas são mais sensíveis à quantização do que outras, e ignorar as sensíveis melhora a qualidade com um pequeno custo de memória. O SDNQ trata isso nativamente: ele inclui listas de ignorar embutidas (uma genérica e listas por modelo para muitas arquiteturas) que são aplicadas automaticamente, de modo que as camadas sensíveis habituais permanecem não quantizadas sem qualquer configuração. Para um controle mais fino, ele também pode selecionar camadas dinamicamente, medindo o erro de quantização por camada e ignorando as que excedem um limiar de perda. Consulte o guia de quantização dinâmica do SDNQ para mais detalhes.

Inferência mais rápida com matmul quantizado

Por padrão, o SDNQ desquantiza os pesos de volta para o dtype de computação antes de cada matmul. Habilitar o matmul quantizado executa o matmul diretamente em um dtype de baixa precisão (INT8/FP8/FP16), o que é mais rápido em hardware com uma instalação funcional do Triton (CUDA, ROCm ou XPU).

Ative-o em um modelo carregado com apply_sdnq_options_to_model, que lê a configuração de quantização existente do modelo e reaplica a opção no local. Isso também funciona para checkpoints pré-quantizados, nos quais nenhuma SDNQConfig é passada.

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

Modelos quantizados SDNQ são compatíveis com torch.compile para ganhos adicionais de velocidade, inclusive em conjunto com group offloading.

transformer.compile(fullgraph=True)

Salvar um modelo quantizado

Modelos quantizados podem ser serializados com save_pretrained e recarregados sem um quantization_config.

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

Recursos

Update on GitHub