[Quantization] Add Quanto backend

.github/workflows/nightly_tests.yml

@@ -418,6 +418,8 @@ jobs:
  test_location: "gguf"
  - backend: "torchao"
  test_location: "torchao"
+ - backend: "optimum_quanto"
+ test_location: "quanto"
  runs-on:
  group: aws-g6e-xlarge-plus
  container:

docs/source/en/_toctree.yml

@@ -173,6 +173,8 @@
  title: gguf
  - local: quantization/torchao
  title: torchao
+ - local: quantization/quanto 
+ title: quanto
  title: Quantization Methods
 - sections:
  - local: optimization/fp16

docs/source/en/api/quantization.md

@@ -31,6 +31,11 @@ Learn how to quantize models in the [Quantization](../quantization/overview) gui
 ## GGUFQuantizationConfig
 
 [[autodoc]] GGUFQuantizationConfig
+
+## QuantoConfig
+
+[[autodoc]] QuantoConfig
+
 ## TorchAoConfig
 
 [[autodoc]] TorchAoConfig

docs/source/en/quantization/overview.md

@@ -36,5 +36,6 @@ Diffusers currently supports the following quantization methods.
 - [BitsandBytes](./bitsandbytes)
 - [TorchAO](./torchao)
 - [GGUF](./gguf)
+- [Quanto](./quanto.md)
 
 [This resource](https://huggingface.co/docs/transformers/main/en/quantization/overview#when-to-use-what) provides a good overview of the pros and cons of different quantization techniques.

docs/source/en/quantization/quanto.md

@@ -0,0 +1,148 @@
+<!--Copyright 2025 The HuggingFace Team. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with
+the License. You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
+an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
+specific language governing permissions and limitations under the License.
+
+-->
+
+# Quanto
+
+[Quanto](https://github.com/huggingface/optimum-quanto) is a PyTorch quantization backend for [Optimum](https://huggingface.co/docs/optimum/en/index). It has been designed with versatility and simplicity in mind:
+
+- All features are available in eager mode (works with non-traceable models)
+- Supports quantization aware training
+- Quantized models are compatible with `torch.compile`
+- Quantized models are Device agnostic (e.g CUDA,XPU,MPS,CPU)
+
+In order to use the Quanto backend, you will first need to install `optimum-quanto>=0.2.6` and `accelerate`
+
+```shell
+pip install optimum-quanto accelerate
+```
+
+Now you can quantize a model by passing the `QuantoConfig` object to the `from_pretrained()` method. Although the Quanto library does allow quantizing `nn.Conv2d` and `nn.LayerNorm` modules, currently, Diffusers only supports quantizing the weights in the `nn.Linear` layers of a model. The following snippet demonstrates how to apply `float8` quantization with Quanto. 
+
+```python
+import torch
+from diffusers import FluxTransformer2DModel, QuantoConfig
+
+model_id = "black-forest-labs/FLUX.1-dev"
+quantization_config = QuantoConfig(weights_dtype="float8")
+transformer = FluxTransformer2DModel.from_pretrained(
+ model_id,
+ subfolder="transformer",
+ quantization_config=quantization_config,
+ torch_dtype=torch.bfloat16,
+)
+
+pipe = FluxPipeline.from_pretrained(model_id, transformer=transformer, torch_dtype=torch_dtype)
+pipe.to("cuda")
+
+prompt = "A cat holding a sign that says hello world"
+image = pipe(
+ prompt, num_inference_steps=50, guidance_scale=4.5, max_sequence_length=512
+).images[0]
+image.save("output.png")
+```
+
+## Skipping Quantization on specific modules
+
+It is possible to skip applying quantization on certain modules using the `modules_to_not_convert` argument in the `QuantoConfig`. Please ensure that the modules passed in to this argument match the keys of the modules in the `state_dict` 
+
+```python
+import torch
+from diffusers import FluxTransformer2DModel, QuantoConfig
+
+model_id = "black-forest-labs/FLUX.1-dev"
+quantization_config = QuantoConfig(weights_dtype="float8", modules_to_not_convert=["proj_out"])
+transformer = FluxTransformer2DModel.from_pretrained(
+ model_id,
+ subfolder="transformer",
+ quantization_config=quantization_config,
+ torch_dtype=torch.bfloat16,
+)
+```
+
+## Using `from_single_file` with the Quanto Backend
+
+`QuantoConfig` is compatible with `~FromOriginalModelMixin.from_single_file`. 
+
+```python
+import torch
+from diffusers import FluxTransformer2DModel, QuantoConfig
+
+ckpt_path = "https://huggingface.co/black-forest-labs/FLUX.1-dev/blob/main/flux1-dev.safetensors"
+quantization_config = QuantoConfig(weights_dtype="float8")
+transformer = FluxTransformer2DModel.from_single_file(ckpt_path, quantization_config=quantization_config, torch_dtype=torch.bfloat16)
+```
+
+## Saving Quantized models
+
+Diffusers supports serializing Quanto models using the `~ModelMixin.save_pretrained` method.
+
+The serialization and loading requirements are different for models quantized directly with the Quanto library and models quantized
+with Diffusers using Quanto as the backend. It is currently not possible to load models quantized directly with Quanto into Diffusers using `~ModelMixin.from_pretrained`
+
+```python
+import torch
+from diffusers import FluxTransformer2DModel, QuantoConfig
+
+model_id = "black-forest-labs/FLUX.1-dev"
+quantization_config = QuantoConfig(weights_dtype="float8")
+transformer = FluxTransformer2DModel.from_pretrained(
+ model_id,
+ subfolder="transformer",
+ quantization_config=quantization_config,
+ torch_dtype=torch.bfloat16,
+)
+# save quantized model to reuse
+transformer.save_pretrained("<your quantized model save path>")
+
+# you can reload your quantized model with
+model = FluxTransformer2DModel.from_pretrained("<your quantized model save path>")
+```
+
+## Using `torch.compile` with Quanto
+
+Currently the Quanto backend supports `torch.compile` for the following quantization types:
+
+- `int8` weights 
+
+```python
+import torch
+from diffusers import FluxPipeline, FluxTransformer2DModel, QuantoConfig
+
+model_id = "black-forest-labs/FLUX.1-dev"
+quantization_config = QuantoConfig(weights_dtype="int8")
+transformer = FluxTransformer2DModel.from_pretrained(
+ model_id,
+ subfolder="transformer",
+ quantization_config=quantization_config,
+ torch_dtype=torch.bfloat16,
+)
+transformer = torch.compile(transformer, mode="max-autotune", fullgraph=True)
+
+pipe = FluxPipeline.from_pretrained(
+ model_id, transformer=transformer, torch_dtype=torch_dtype
+)
+pipe.to("cuda")
+images = pipe("A cat holding a sign that says hello").images[0]
+images.save("flux-quanto-compile.png")
+```
+
+## Supported Quantization Types
+
+### Weights
+
+- float8
+- int8
+- int4
+- int2
+
+

setup.py

@@ -128,6 +128,10 @@
  "GitPython<3.1.19",
  "scipy",
  "onnx",
+ "optimum_quanto>=0.2.6",
+ "gguf>=0.10.0",
+ "torchao>=0.7.0",
+ "bitsandbytes>=0.43.3",
  "regex!=2019.12.17",
  "requests",
  "tensorboard",
@@ -235,6 +239,11 @@ def run(self):
 )
 extras["torch"] = deps_list("torch", "accelerate")
 
+extras["bitsandbytes"] = deps_list("bitsandbytes", "accelerate")
+extras["gguf"] = deps_list("gguf", "accelerate")
+extras["optimum_quanto"] = deps_list("optimum_quanto", "accelerate")
+extras["torchao"] = deps_list("torchao", "accelerate")
+
 if os.name == "nt": # windows
  extras["flax"] = [] # jax is not supported on windows
 else:

src/diffusers/__init__.py

@@ -2,6 +2,15 @@
 
 from typing import TYPE_CHECKING
 
+from diffusers.quantizers import quantization_config
+from diffusers.utils import dummy_gguf_objects
+from diffusers.utils.import_utils import (
+ is_bitsandbytes_available,
+ is_gguf_available,
+ is_optimum_quanto_version,
+ is_torchao_available,
+)
+
 from .utils import (
  DIFFUSERS_SLOW_IMPORT,
  OptionalDependencyNotAvailable,
@@ -11,6 +20,7 @@
  is_librosa_available,
  is_note_seq_available,
  is_onnx_available,
+ is_optimum_quanto_available,
  is_scipy_available,
  is_sentencepiece_available,
  is_torch_available,
@@ -32,7 +42,7 @@
  "loaders": ["FromOriginalModelMixin"],
  "models": [],
  "pipelines": [],
- "quantizers.quantization_config": ["BitsAndBytesConfig", "GGUFQuantizationConfig", "TorchAoConfig"],
+ "quantizers.quantization_config": [],
  "schedulers": [],
  "utils": [
  "OptionalDependencyNotAvailable",
@@ -54,6 +64,55 @@
  ],
 }
 
+try:
+ if not is_bitsandbytes_available():
+ raise OptionalDependencyNotAvailable()
+except OptionalDependencyNotAvailable:
+ from .utils import dummy_bitsandbytes_objects
+
+ _import_structure["utils.dummy_bitsandbytes_objects"] = [
+ name for name in dir(dummy_bitsandbytes_objects) if not name.startswith("_")
+ ]
+else:
+ _import_structure["quantizers.quantization_config"].append("BitsAndBytesConfig")
+
+try:
+ if not is_gguf_available():
+ raise OptionalDependencyNotAvailable()
+except OptionalDependencyNotAvailable:
+ from .utils import dummy_gguf_objects
+
+ _import_structure["utils.dummy_gguf_objects"] = [
+ name for name in dir(dummy_gguf_objects) if not name.startswith("_")
+ ]
+else:
+ _import_structure["quantizers.quantization_config"].append("GGUFQuantizationConfig")
+
+try:
+ if not is_torchao_available():
+ raise OptionalDependencyNotAvailable()
+except OptionalDependencyNotAvailable:
+ from .utils import dummy_torchao_objects
+
+ _import_structure["utils.dummy_torchao_objects"] = [
+ name for name in dir(dummy_torchao_objects) if not name.startswith("_")
+ ]
+else:
+ _import_structure["quantizers.quantization_config"].append("TorchAoConfig")
+
+try:
+ if not is_optimum_quanto_available():
+ raise OptionalDependencyNotAvailable()
+except OptionalDependencyNotAvailable:
+ from .utils import dummy_optimum_quanto_objects
+
+ _import_structure["utils.dummy_optimum_quanto_objects"] = [
+ name for name in dir(dummy_optimum_quanto_objects) if not name.startswith("_")
+ ]
+else:
+ _import_structure["quantizers.quantization_config"].append("QuantoConfig")
+
+
 try:
  if not is_onnx_available():
  raise OptionalDependencyNotAvailable()
@@ -598,7 +657,38 @@
 
 if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
  from .configuration_utils import ConfigMixin
- from .quantizers.quantization_config import BitsAndBytesConfig, GGUFQuantizationConfig, TorchAoConfig
+
+ try:
+ if not is_bitsandbytes_available():
+ raise OptionalDependencyNotAvailable()
+ except OptionalDependencyNotAvailable:
+ from .utils.dummy_bitsandbytes_objects import *
+ else:
+ from .quantizers.quantization_config import BitsAndBytesConfig
+
+ try:
+ if not is_gguf_available():
+ raise OptionalDependencyNotAvailable()
+ except OptionalDependencyNotAvailable:
+ from .utils.dummy_gguf_objects import *
+ else:
+ from .quantizers.quantization_config import GGUFQuantizationConfig
+
+ try:
+ if not is_torchao_available():
+ raise OptionalDependencyNotAvailable()
+ except OptionalDependencyNotAvailable:
+ from .utils.dummy_torchao_objects import *
+ else:
+ from .quantizers.quantization_config import TorchAoConfig
+
+ try:
+ if not is_optimum_quanto_available():
+ raise OptionalDependencyNotAvailable()
+ except OptionalDependencyNotAvailable:
+ from .utils.dummy_optimum_quanto_objects import *
+ else:
+ from .quantizers.quantization_config import QuantoConfig
 
  try:
  if not is_onnx_available():

src/diffusers/dependency_versions_table.py

@@ -35,6 +35,10 @@
  "GitPython": "GitPython<3.1.19",
  "scipy": "scipy",
  "onnx": "onnx",
+ "optimum_quanto": "optimum_quanto>=0.2.6",
+ "gguf": "gguf>=0.10.0",
+ "torchao": "torchao>=0.7.0",
+ "bitsandbytes": "bitsandbytes>=0.43.3",
  "regex": "regex!=2019.12.17",
  "requests": "requests",
  "tensorboard": "tensorboard",

src/diffusers/models/model_loading_utils.py

@@ -245,6 +245,9 @@ def load_model_dict_into_meta(
  ):
  param = param.to(torch.float32)
  set_module_kwargs["dtype"] = torch.float32
+ # For quantizers have save weights using torch.float8_e4m3fn
+ elif hf_quantizer is not None and param.dtype == getattr(torch, "float8_e4m3fn", None):
+ pass
  else:
  param = param.to(dtype)
  set_module_kwargs["dtype"] = dtype
@@ -292,7 +295,9 @@ def load_model_dict_into_meta(
  elif is_quantized and (
  hf_quantizer.check_if_quantized_param(model, param, param_name, state_dict, param_device=param_device)
  ):
- hf_quantizer.create_quantized_param(model, param, param_name, param_device, state_dict, unexpected_keys)
+ hf_quantizer.create_quantized_param(
+ model, param, param_name, param_device, state_dict, unexpected_keys, dtype=dtype
+ )
  else:
  set_module_tensor_to_device(model, param_name, param_device, value=param, **set_module_kwargs)
 

src/diffusers/quantizers/auto.py

@@ -26,22 +26,26 @@
  GGUFQuantizationConfig,
  QuantizationConfigMixin,
  QuantizationMethod,
+ QuantoConfig,
  TorchAoConfig,
 )
+from .quanto import QuantoQuantizer
 from .torchao import TorchAoHfQuantizer
 
 
 AUTO_QUANTIZER_MAPPING = {
  "bitsandbytes_4bit": BnB4BitDiffusersQuantizer,
  "bitsandbytes_8bit": BnB8BitDiffusersQuantizer,
  "gguf": GGUFQuantizer,
+ "quanto": QuantoQuantizer,
  "torchao": TorchAoHfQuantizer,
 }
 
 AUTO_QUANTIZATION_CONFIG_MAPPING = {
  "bitsandbytes_4bit": BitsAndBytesConfig,
  "bitsandbytes_8bit": BitsAndBytesConfig,
  "gguf": GGUFQuantizationConfig,
+ "quanto": QuantoConfig,
  "torchao": TorchAoConfig,
 }