Hugging Face Diffusers Quantization Backends

Hugging Face Diffusers now integrates several quantization backends to make large diffusion models, such as FLUX.1-dev, more accessible by reducing their memory and compute requirements. These backends allow users to shrink models without drastically compromising image quality, with 8-bit quantization often producing results nearly indistinguishable from high-precision BF16 models.

FLUX.1-dev Model Architecture and Memory Baseline

To understand the impact of quantization, it is necessary to establish the memory baseline for the FLUX.1-dev model in BF16 precision, which requires approximately 31.447 GB of memory. The model consists of four primary components:

  • Transformer (MMDiT): The core generative multimodal diffusion transformer, requiring 23.8 GB.
  • Text Encoder 2 (T5): Used for nuanced comprehension and text rendering, requiring 9.52 GB.
  • Text Encoder 1 (CLIP): Used for initial text understanding, requiring 246 MB.
  • VAE: Translates images between pixel and latent space, requiring 168 MB.

Quantization efforts primarily target the transformer and the T5 text encoder to achieve the most significant memory savings.

Supported Quantization Backends

bitsandbytes (BnB)

bitsandbytes provides 8-bit and 4-bit quantization. For FLUX.1-dev, 4-bit quantization (using NF4) reduces memory after loading to 12.584 GB and peak memory to 17.281 GB, while maintaining an inference time of 12 seconds on an NVIDIA H100 80GB GPU. 8-bit quantization increases memory usage (19.273 GB loading / 24.432 GB peak) and slows inference to 27 seconds.

torchao

torchao is a PyTorch-native library for architecture optimization. It supports several weight-only formats:

  • int4_weight_only: Lowest memory footprint (10.635 GB loading / 14.654 GB peak) but slowest inference (109 seconds).
  • int8_weight_only: Moderate memory (17.020 GB loading / 21.482 GB peak) with fast inference (15 seconds).
  • float8_weight_only: Similar memory and speed to int8 (17.016 GB loading / 21.488 GB peak; 15 seconds).

Quanto

Integrated via the optimum library, Quanto supports various precisions:

  • INT4: 12.254 GB loading / 16.139 GB peak memory; 109 seconds inference.
  • INT8: 17.330 GB loading / 21.814 GB peak memory; 15 seconds inference.
  • FP8: 16.395 GB loading / 20.898 GB peak memory; 16 seconds inference.

GGUF

Diffusers supports the GGUF file format, allowing the use of pre-quantized models from the llama.cpp community via from_single_file. Benchmarks for FLUX.1-dev show:

  • Q2_k: 13.264 GB loading / 17.752 GB peak memory; 26 seconds inference.
  • Q4_1: 16.838 GB loading / 21.326 GB peak memory; 23 seconds inference.
  • Q8_0: 21.502 GB loading / 25.973 GB peak memory; 15 seconds inference.

FP8 Layerwise Casting

enable_layerwise_casting is a memory optimization that stores weights in FP8 (e4m3) but dynamically casts them to a higher compute precision (e.g., BF16) during calculation. This reduces memory after loading to 23.682 GB and peak memory to 28.451 GB, with an inference time of 13 seconds.

Combining Quantization with Memory Optimizations

Quantization backends can be paired with other Diffusers optimization techniques to further reduce VRAM usage:

  • Model CPU Offloading (enable_model_cpu_offload): Moves entire components between CPU and GPU. Combining BnB 4-bit with this method reduces peak memory to 12.383 GB.
  • Group Offloading (enable_group_offload): Moves groups of internal layers to the CPU. Combining FP8 layerwise casting with group offloading reduces memory after loading to 9.264 GB and peak memory to 14.232 GB.
  • torch.compile: Accelerates execution via PyTorch 2.x. While it does not lower memory, it significantly speeds up inference. For example, torchao int4_weight_only inference time drops from 109 seconds to 6 seconds after compilation (though compilation takes ~285 seconds).

Backend Selection Guide

Depending on the hardware and performance goals, different backends are recommended:

  • NVIDIA Users (Easiest): Use bitsandbytes 4/8-bit.
  • Inference Speed: Use torchao, GGUF, or bitsandbytes combined with torch.compile().
  • Hardware Flexibility (CPU/MPS) or FP8: Use Quanto.
  • Hopper/Ada Architecture: Use FP8 Layerwise Casting.
  • Existing GGUF Models: Use GGUF loading via from_single_file.

Sources