Add GemLite quantization backend support for pre-quantized checkpoints

Author: gabe-engineers

.github/workflows/nightly_tests.yml

@@ -348,6 +348,9 @@ jobs:
           - backend: "gguf"
             test_location: "gguf"
             additional_deps: ["peft", "kernels"]
+          - backend: "gemlite"
+            test_location: "gemlite"
+            additional_deps: []
           - backend: "torchao"
             test_location: "torchao"
             additional_deps: []

docs/source/en/_toctree.yml

@@ -174,6 +174,8 @@
     title: bitsandbytes
   - local: quantization/gguf
     title: gguf
+  - local: quantization/gemlite
+    title: GemLite
   - local: quantization/nunchaku
     title: Nunchaku Lite
   - local: quantization/torchao

docs/source/en/api/quantization.md

@@ -30,6 +30,10 @@ Quantization techniques reduce memory and computational costs by representing we
 
 [[autodoc]] quantizers.quantization_config.GGUFQuantizationConfig
 
+## GemLiteConfig
+
+[[autodoc]] quantizers.quantization_config.GemLiteConfig
+
 ## NunchakuLiteQuantizationConfig
 
 [[autodoc]] quantizers.quantization_config.NunchakuLiteQuantizationConfig

docs/source/en/quantization/gemlite.md

@@ -0,0 +1,57 @@
+<!--Copyright 2026 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.
+
+-->
+
+# GemLite
+
+[GemLite](https://github.com/mobiusml/gemlite) is a quantization backend for loading prequantized checkpoints in
+Diffusers. It replaces supported `torch.nn.Linear` layers with GemLite layers, which run packed low-bit weights
+directly with GemLite kernels.
+
+## Install GemLite
+
+GemLite requires version 0.6.0 or later.
+
+```bash
+pip install -U "gemlite>=0.6.0"
+```
+
+## Load a quantized pipeline
+
+GemLite only supports loading prequantized checkpoints. The quantization configuration is stored in the checkpoint's
+`config.json` and read automatically by [`~DiffusionPipeline.from_pretrained`].
+
+```python
+import torch
+from diffusers import DiffusionPipeline
+
+model_id = "gabe-engineers/bonsai-image-ternary-4B-gemlite-2bit-unpacked-encoder"
+
+pipe = DiffusionPipeline.from_pretrained(
+    model_id,
+    dtype=torch.float16,
+    device_map="cuda",
+)
+
+image = pipe(
+    prompt="A bonsai tree in a quiet ceramic studio, soft morning light",
+    height=1024,
+    width=1024,
+    num_inference_steps=4,
+    guidance_scale=1.0,
+).images[0]
+image.save("bonsai-gemlite.png")
+```
+
+> [!TIP]
+> `dtype` must match the `compute_dtype` in the checkpoint's GemLite quantization configuration. This
+> checkpoint uses `torch.float16`.

docs/source/en/quantization/overview.md

@@ -145,6 +145,7 @@ The following backends support loading prequantized checkpoints out of the box.
 |---|---|
 | [bitsandbytes](./bitsandbytes) | Config is saved in `config.json`; no extra arguments needed. |
 | [GGUF](./gguf) | Uses `from_single_file` with Model classes; pipeline-level loading is not supported. |
+| [GemLite](./gemlite) | Loads prequantized checkpoints; requires the `gemlite` library. |
 | [AutoRound](./autoround) | Only loading is supported; quantize first with the AutoRound CLI or Python API. |
 | [Nunchaku Lite](./nunchaku) | Config is saved in `config.json`; requires the `kernels` package. |
 | [ModelOpt](./modelopt) | Supports both quantizing on the fly and loading prequantized models. |

src/diffusers/__init__.py

@@ -9,6 +9,7 @@
     is_accelerate_available,
     is_auto_round_available,
     is_bitsandbytes_available,
+    is_gemlite_available,
     is_gguf_available,
     is_librosa_available,
     is_note_seq_available,
@@ -46,6 +47,7 @@
     "schedulers": [],
     "utils": [
         "OptionalDependencyNotAvailable",
+        "is_gemlite_available",
         "is_inflect_available",
         "is_invisible_watermark_available",
         "is_librosa_available",
@@ -85,6 +87,18 @@
 else:
     _import_structure["quantizers.quantization_config"].append("GGUFQuantizationConfig")
 
+try:
+    if not is_torch_available() and not is_accelerate_available() and not is_gemlite_available():
+        raise OptionalDependencyNotAvailable()
+except OptionalDependencyNotAvailable:
+    from .utils import dummy_gemlite_objects
+
+    _import_structure["utils.dummy_gemlite_objects"] = [
+        name for name in dir(dummy_gemlite_objects) if not name.startswith("_")
+    ]
+else:
+    _import_structure["quantizers.quantization_config"].append("GemLiteConfig")
+
 try:
     if not is_torch_available() and not is_accelerate_available() and not is_torchao_available():
         raise OptionalDependencyNotAvailable()
@@ -947,6 +961,14 @@
     else:
         from .quantizers.quantization_config import GGUFQuantizationConfig
 
+    try:
+        if not is_gemlite_available():
+            raise OptionalDependencyNotAvailable()
+    except OptionalDependencyNotAvailable:
+        from .utils.dummy_gemlite_objects import *
+    else:
+        from .quantizers.quantization_config import GemLiteConfig
+
     try:
         if not is_torchao_available():
             raise OptionalDependencyNotAvailable()

src/diffusers/models/modeling_utils.py

@@ -1411,6 +1411,11 @@ def from_pretrained(cls, pretrained_model_name_or_path: str | os.PathLike | None
             "error_msgs": error_msgs,
         }
 
+        # Quantizers may need to finalize module state before dispatch reads `model.state_dict()`.
+        if hf_quantizer is not None:
+            hf_quantizer.postprocess_model(model)
+            model.hf_quantizer = hf_quantizer
+
         # Dispatch model with hooks on all devices if necessary
         if device_map is not None:
             device_map_kwargs = {
@@ -1420,10 +1425,6 @@ def from_pretrained(cls, pretrained_model_name_or_path: str | os.PathLike | None
             }
             dispatch_model(model, **device_map_kwargs)
 
-        if hf_quantizer is not None:
-            hf_quantizer.postprocess_model(model)
-            model.hf_quantizer = hf_quantizer
-
         if (
             torch_dtype is not None
             and torch_dtype == getattr(torch, "float8_e4m3fn", None)

src/diffusers/quantizers/auto.py

@@ -20,12 +20,14 @@
 
 from .autoround import AutoRoundQuantizer
 from .bitsandbytes import BnB4BitDiffusersQuantizer, BnB8BitDiffusersQuantizer
+from .gemlite import GemLiteQuantizer
 from .gguf import GGUFQuantizer
 from .modelopt import NVIDIAModelOptQuantizer
 from .nunchaku import NunchakuLiteQuantizer
 from .quantization_config import (
     AutoRoundConfig,
     BitsAndBytesConfig,
+    GemLiteConfig,
     GGUFQuantizationConfig,
     NunchakuLiteQuantizationConfig,
     NVIDIAModelOptConfig,
@@ -41,6 +43,7 @@
 AUTO_QUANTIZER_MAPPING = {
     "bitsandbytes_4bit": BnB4BitDiffusersQuantizer,
     "bitsandbytes_8bit": BnB8BitDiffusersQuantizer,
+    "gemlite": GemLiteQuantizer,
     "gguf": GGUFQuantizer,
     "quanto": QuantoQuantizer,
     "torchao": TorchAoHfQuantizer,
@@ -52,6 +55,7 @@
 AUTO_QUANTIZATION_CONFIG_MAPPING = {
     "bitsandbytes_4bit": BitsAndBytesConfig,
     "bitsandbytes_8bit": BitsAndBytesConfig,
+    "gemlite": GemLiteConfig,
     "gguf": GGUFQuantizationConfig,
     "quanto": QuantoConfig,
     "torchao": TorchAoConfig,

src/diffusers/quantizers/gemlite/__init__.py

@@ -0,0 +1 @@
+from .gemlite_quantizer import GemLiteQuantizer