Quanto quantization cuts memory for Transformer diffusion pipelines

TL;DR

Quantization with Hugging Face Quanto reduces the GPU memory needed for Transformer‑based diffusion pipelines (e.g., PixArt‑Sigma, Stable Diffusion 3, Aura Flow) from ~12 GB to as low as ~5 GB with only minor latency impact and negligible quality loss.

Introduction – Why memory matters for diffusion Transformers

Transformer backbones have become the dominant architecture for high‑resolution text‑to‑image diffusion models, scaling from 0.6 B to 8 B parameters. Larger models increase GPU memory consumption dramatically; a full Stable Diffusion 3 inference in FP16 occupies 18.8 GB. This memory barrier limits adoption on consumer GPUs and hampers rapid experimentation. The post demonstrates that Quanto’s quantization utilities, integrated in the Diffusers library, can dramatically shrink the memory footprint while preserving visual quality.

Quantization basics with Quanto

Quanto is a PyTorch‑based quantization toolkit that lives inside Hugging Face Optimum. It supports weight‑only quantization to several low‑precision formats (FP8, INT8, INT4) and can be applied to any Diffusers module.

from optimum.quanto import freeze, qfloat8, quantize
from diffusers import PixArtSigmaPipeline
import torch

pipeline = PixArtSigmaPipeline.from_pretrained(
    "PixArt-alpha/PixArt-Sigma-XL-2-1024-MS", torch_dtype=torch.float16
).to("cuda")

# Quantize the diffusion transformer to FP8
quantize(pipeline.transformer, weights=qfloat8)
freeze(pipeline.transformer)

The same quantize/freeze calls work for text encoders or any other sub‑module.

Memory and latency results – FP8 weight quantization

Diffusion transformer only

Batch size Memory (GB) Latency (s)
1 11.55 (FP8) vs 12.09 (FP16) 1.54 vs 1.20
4 11.55 (FP8) vs 12.09 (FP16) 5.11 vs 4.48

The FP8 weights cut memory by ~0.5 GB with a modest latency increase.

Adding text‑encoder quantization

Batch size Quantize TE? Memory (GB) Latency (s)
1 No 11.55 1.54
1 Yes 5.36 1.60
4 No 11.55 5.11
4 Yes 5.36 5.14

Quantizing both the diffusion transformer and the text encoder halves the memory usage while keeping latency essentially unchanged.

Generality across models

The authors evaluated three pipelines:

  • PixArt‑Sigma (0.61 B parameters)
  • Stable Diffusion 3 (medium) (2.03 B parameters, three text encoders)
  • Aura Flow (6.84 B parameters)

For PixArt‑Sigma and Aura Flow, quantizing the text encoder always yielded large memory savings. Stable Diffusion 3 required selective quantization because the second text encoder (the middle CLIP variant) degrades quality when quantized. Recommended strategies are:

  • Quantize only the first CLIP encoder, or
  • Quantize only the third T5 encoder, or
  • Quantize both the first and third encoders.

A representative table for SD‑3 (batch 1, diffusion transformer always FP8) shows memory ranging from 8.20 GB (all three encoders quantized) to 16.40 GB (no quantization).

Additional findings

bfloat16 vs fp16 on H100

On NVIDIA H100 GPUs, using bfloat16 with INT8 or FP8 weights improves latency:

Precision Quantization Memory (GB) Latency (s)
FP16 INT8 5.363 1.538
BF16 INT8 5.364 1.454
FP16 FP8 5.363 1.601
BF16 FP8 5.363 1.495

INT8 (qint8) and fused QKV projections

INT8 weights are faster than FP8, especially when the attention QKV projections are fused (fuse_qkv_projections()). For PixArt‑Sigma (batch 1) latency dropped from 1.538 s (INT8, no fuse) to 1.504 s (INT8 with fused QKV).

INT4 (qint4) for aggressive compression

When combined with bfloat16 on H100, INT4 reduces memory dramatically (e.g., PixArt‑Sigma from 9.38 GB to 3.06 GB) but latency rises (≈7.6 s) because computation still occurs in bfloat16. Quality loss is noticeable; the authors recommend excluding the final projection layer (proj_out) from quantization to mitigate degradation.

Saving and loading quantized Diffusers models

Quanto provides model classes that can be persisted and re‑loaded:

from diffusers import PixArtTransformer2DModel
from optimum.quanto import QuantizedPixArtTransformer2DModel, qfloat8

model = PixArtTransformer2DModel.from_pretrained(
    "PixArt-alpha/PixArt-Sigma-XL-2-1024-MS", subfolder="transformer"
)
qmodel = QuantizedPixArtTransformer2DModel.quantize(model, weights=qfloat8)
qmodel.save_pretrained("pixart-sigma-fp8")  # 587 MB checkpoint

Loading follows the same pattern and the quantized transformer can be injected into a DiffusionPipeline.

Practical tips for developers

  • Mix‑and‑match quantization types per module (e.g., FP8 for the text encoder, INT8 for the diffusion transformer) to balance memory and speed.
  • Combine Quanto quantization with Diffusers’ existing memory‑saving utilities such as enable_model_cpu_offload() for further reductions.
  • For INT4 deployments, always exclude the final projection layer (exclude="proj_out") to preserve image fidelity.

Conclusion

Quantizing Transformer‑based diffusion pipelines with Hugging Face Quanto cuts GPU memory requirements by up to 70 % (from ~12 GB to ~5 GB) while keeping latency low and visual quality largely intact. The technique works across multiple state‑of‑the‑art models, and the resulting checkpoints are dramatically smaller (e.g., 587 MB vs 2.44 GB). Developers can readily integrate these workflows into existing Diffusers pipelines and further combine them with other memory‑optimisation strategies.

Acknowledgments

Thanks to Pedro Cuenca for extensive reviews of the post.

Sources