From f6e40fb8988f83770176854025dddd2f0fd45806 Mon Sep 17 00:00:00 2001 From: Gabriel Date: Tue, 21 Jul 2026 12:54:36 -0500 Subject: [PATCH 01/10] Add GemLite quantization backend support for pre-quantized checkpoints --- .github/workflows/nightly_tests.yml | 3 + docs/source/en/_toctree.yml | 2 + docs/source/en/api/quantization.md | 4 + docs/source/en/quantization/gemlite.md | 57 ++ docs/source/en/quantization/overview.md | 1 + src/diffusers/__init__.py | 22 + src/diffusers/models/modeling_utils.py | 9 +- src/diffusers/quantizers/auto.py | 4 + src/diffusers/quantizers/gemlite/__init__.py | 1 + .../quantizers/gemlite/gemlite_quantizer.py | 316 ++++++++ .../quantizers/quantization_config.py | 148 ++++ src/diffusers/utils/__init__.py | 2 + src/diffusers/utils/dummy_gemlite_objects.py | 17 + src/diffusers/utils/import_utils.py | 27 + src/diffusers/utils/testing_utils.py | 20 + tests/quantization/gemlite/__init__.py | 0 tests/quantization/gemlite/test_gemlite.py | 727 ++++++++++++++++++ tests/testing_utils.py | 21 + 18 files changed, 1377 insertions(+), 4 deletions(-) create mode 100644 docs/source/en/quantization/gemlite.md create mode 100644 src/diffusers/quantizers/gemlite/__init__.py create mode 100644 src/diffusers/quantizers/gemlite/gemlite_quantizer.py create mode 100644 src/diffusers/utils/dummy_gemlite_objects.py create mode 100644 tests/quantization/gemlite/__init__.py create mode 100644 tests/quantization/gemlite/test_gemlite.py diff --git a/.github/workflows/nightly_tests.yml b/.github/workflows/nightly_tests.yml index 678d106a5fc7..4a2ab03d94a8 100644 --- a/.github/workflows/nightly_tests.yml +++ b/.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: [] diff --git a/docs/source/en/_toctree.yml b/docs/source/en/_toctree.yml index 98773eab8866..e74c411848f5 100644 --- a/docs/source/en/_toctree.yml +++ b/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 diff --git a/docs/source/en/api/quantization.md b/docs/source/en/api/quantization.md index 406b992a6ca9..86e87258036d 100644 --- a/docs/source/en/api/quantization.md +++ b/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 diff --git a/docs/source/en/quantization/gemlite.md b/docs/source/en/quantization/gemlite.md new file mode 100644 index 000000000000..78ed6cfe5334 --- /dev/null +++ b/docs/source/en/quantization/gemlite.md @@ -0,0 +1,57 @@ + + +# 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`. diff --git a/docs/source/en/quantization/overview.md b/docs/source/en/quantization/overview.md index ff5657917e78..2498a2a18b87 100644 --- a/docs/source/en/quantization/overview.md +++ b/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. | diff --git a/src/diffusers/__init__.py b/src/diffusers/__init__.py index e72de0924c07..32de5643ee44 100644 --- a/src/diffusers/__init__.py +++ b/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() diff --git a/src/diffusers/models/modeling_utils.py b/src/diffusers/models/modeling_utils.py index b8d4c048c3cd..6e292c4ae747 100644 --- a/src/diffusers/models/modeling_utils.py +++ b/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) diff --git a/src/diffusers/quantizers/auto.py b/src/diffusers/quantizers/auto.py index 03f2347792ea..fbda84448042 100644 --- a/src/diffusers/quantizers/auto.py +++ b/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, diff --git a/src/diffusers/quantizers/gemlite/__init__.py b/src/diffusers/quantizers/gemlite/__init__.py new file mode 100644 index 000000000000..fb19f4c38ce8 --- /dev/null +++ b/src/diffusers/quantizers/gemlite/__init__.py @@ -0,0 +1 @@ +from .gemlite_quantizer import GemLiteQuantizer diff --git a/src/diffusers/quantizers/gemlite/gemlite_quantizer.py b/src/diffusers/quantizers/gemlite/gemlite_quantizer.py new file mode 100644 index 000000000000..b1fe16ec5e1f --- /dev/null +++ b/src/diffusers/quantizers/gemlite/gemlite_quantizer.py @@ -0,0 +1,316 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +from ...utils import get_module_from_name, is_gemlite_available, is_gemlite_version, is_torch_available, logging +from ..base import DiffusersQuantizer + + +if TYPE_CHECKING: + from ...models.modeling_utils import ModelMixin + from ..quantization_config import GemLiteConfig + +if is_torch_available(): + import torch + import torch.nn as nn + + +logger = logging.get_logger(__name__) + + +GEMLITE_STATE_NAMES = ("W_q", "bias", "scales", "zeros", "metadata", "orig_shape", "meta_scale") + +_GEMLITE_MIN_VERSION = "0.6.0" + + +def _is_in_skip_modules(name: str, modules_to_not_convert: list[str]) -> bool: + return any((key + "." in name) or (key == name) for key in modules_to_not_convert) + + +def _normalize_torch_device(device: Any) -> "torch.device": + if isinstance(device, int): + return torch.device(f"cuda:{device}") + return torch.device(device) + + +def _replace_with_gemlite_linear( + model: "ModelMixin", modules_to_not_convert: list[str], quantization_config: "GemLiteConfig" +) -> int: + """ + Replace eligible `nn.Linear` modules in `model` with GemLite modules whose serialized tensors have meta shapes. + + Returns the number of replaced modules. Modules in `modules_to_not_convert` are left unchanged. + """ + from gemlite.core import DType, GemLiteLinearTriton + from gemlite.dtypes import DTYPE_TO_TORCH, PACKING_BITWIDTH_TO_TORCH_DTYPE + + gemlite_dtypes = { + "fp16": DType.FP16, + "float16": DType.FP16, + "bf16": DType.BF16, + "bfloat16": DType.BF16, + "fp32": DType.FP32, + "float32": DType.FP32, + } + input_dtype = gemlite_dtypes[quantization_config.input_dtype] + output_dtype = gemlite_dtypes[quantization_config.output_dtype] + scales_gemlite_dtype = gemlite_dtypes[quantization_config.scales_dtype] + scales_dtype = DTYPE_TO_TORCH[scales_gemlite_dtype.value] + zeros_dtype = DTYPE_TO_TORCH[gemlite_dtypes[quantization_config.zeros_dtype].value] + packed_dtype = PACKING_BITWIDTH_TO_TORCH_DTYPE[quantization_config.packing_bitwidth] + + elements_per_sample = quantization_config.packing_bitwidth // quantization_config.bits + quantized_fqns = ( + set(quantization_config.quantized_fqns) if quantization_config.quantized_fqns is not None else None + ) + + def initialize_serialized_tensors(gemlite_linear: "nn.Module", linear: "nn.Linear") -> None: + gemlite_linear.elements_per_sample = elements_per_sample + gemlite_linear.meta_dtype = scales_gemlite_dtype + gemlite_linear.channel_scale_mode = 0 + gemlite_linear.W_group_mode = 0 + gemlite_linear.data_contiguous = True + gemlite_linear.W_q = torch.empty( + linear.in_features // elements_per_sample, + linear.out_features, + dtype=packed_dtype, + device=linear.weight.device, + ) + gemlite_linear.scales = torch.empty( + linear.in_features // quantization_config.group_size, + linear.out_features, + dtype=scales_dtype, + device=linear.weight.device, + ) + gemlite_linear.zeros = torch.empty( + linear.in_features // quantization_config.group_size, + linear.out_features, + dtype=zeros_dtype, + device=linear.weight.device, + ) + gemlite_linear.metadata = torch.empty( + len(gemlite_linear.get_meta_args()), dtype=torch.int32, device=linear.weight.device + ) + gemlite_linear.orig_shape = torch.empty(2, dtype=torch.int32, device=linear.weight.device) + gemlite_linear.meta_scale = torch.empty((), dtype=torch.float32, device=linear.weight.device) + + def replace(module: "nn.Module", prefix: str = "") -> int: + replaced = 0 + for name, child in module.named_children(): + child_name = f"{prefix}.{name}" if prefix else name + should_replace = ( + isinstance(child, nn.Linear) + and (quantized_fqns is None or child_name in quantized_fqns) + and not _is_in_skip_modules(child_name, modules_to_not_convert) + ) + if should_replace: + gemlite_linear = GemLiteLinearTriton( + W_nbits=quantization_config.bits, + group_size=quantization_config.group_size, + in_features=child.in_features, + out_features=child.out_features, + input_dtype=input_dtype, + output_dtype=output_dtype, + ).to(child.weight.device) + # Match the checkpoint's serialized tensor shapes and dtypes so Accelerate sizes the device map correctly. + initialize_serialized_tensors(gemlite_linear, child) + if child.bias is not None: + gemlite_linear.bias = nn.Parameter(torch.empty_like(child.bias), requires_grad=False) + gemlite_linear._gemlite_loaded_param_names = set() + setattr(module, name, gemlite_linear) + replaced += 1 + else: + replaced += replace(child, child_name) + return replaced + + return replace(model) + + +class GemLiteQuantizer(DiffusersQuantizer): + """ + Diffusers quantizer for GemLite. + + GemLite provides Triton kernels for fast low-bit matrix multiplication. This quantizer wires those kernels into + Diffusers by replacing eligible `torch.nn.Linear` layers with `GemLiteLinearTriton` modules, allowing quantized + weights to run directly through GemLite kernels instead of being materialized back to full precision. + + This quantizer only loads pre-quantized checkpoints. It replaces `torch.nn.Linear` modules with + `GemLiteLinearTriton` modules before weight loading, then restores the serialized GemLite state (`W_q`, `scales`, + `zeros`, `metadata`, `orig_shape`, `meta_scale`, ...) through the low-memory loader. The quantization config + provides the serialized layout, so the replacement tensors have the checkpoint's exact shapes and dtypes before + Accelerate estimates module sizes for `device_map="auto"`. + + Modules listed in `modules_to_not_convert` are skipped and left in their original dtype. + """ + + use_keep_in_fp32_modules = True + requires_calibration = False + required_packages = ["gemlite"] + + def __init__(self, quantization_config, **kwargs): + super().__init__(quantization_config, **kwargs) + if not self.pre_quantized: + raise ValueError("GemLite quantization in Diffusers only supports loading pre-quantized checkpoints.") + + self.compute_dtype = quantization_config.compute_dtype + self.modules_to_not_convert = quantization_config.modules_to_not_convert or [] + if not isinstance(self.modules_to_not_convert, list): + self.modules_to_not_convert = [self.modules_to_not_convert] + + def update_torch_dtype(self, torch_dtype: "torch.dtype | None") -> "torch.dtype": + if torch_dtype is None: + return self.compute_dtype + if torch_dtype != self.compute_dtype: + raise ValueError( + "`torch_dtype` passed to `from_pretrained` must match `GemLiteConfig.compute_dtype`. " + f"Got {torch_dtype} and {self.compute_dtype}, respectively." + ) + return torch_dtype + + def get_special_dtypes_update(self, model, torch_dtype: "torch.dtype") -> dict[str, "torch.dtype"]: + special_dtypes = super().get_special_dtypes_update(model, torch_dtype) + + from gemlite.core import GemLiteLinearTriton + + for module_name, module in model.named_modules(): + if not isinstance(module, GemLiteLinearTriton): + continue + for tensor_name, tensor in module.named_parameters(recurse=False): + name = f"{module_name}.{tensor_name}" if module_name else tensor_name + special_dtypes[name] = tensor.dtype + for tensor_name, tensor in module.named_buffers(recurse=False): + name = f"{module_name}.{tensor_name}" if module_name else tensor_name + special_dtypes[name] = tensor.dtype + + return special_dtypes + + @property + def supports_parallel_loading(self) -> bool: + return False + + def validate_environment(self, *args, **kwargs): + device_map = kwargs.get("device_map") + if isinstance(device_map, dict) and "disk" in device_map.values(): + raise ValueError("GemLite quantization does not support disk offloading.") + + if not is_gemlite_available(): + raise ImportError( + "Using GemLite quantization requires the gemlite library. Please install it with `pip install gemlite`." + ) + if is_gemlite_version("<", _GEMLITE_MIN_VERSION): + raise ImportError( + f"Using GemLite quantization requires gemlite>={_GEMLITE_MIN_VERSION}. " + "Please upgrade with `pip install -U gemlite`." + ) + try: + __import__("gemlite.core") + except Exception as error: + raise ImportError("GemLite is installed but its core linear module could not be imported.") from error + + def check_if_quantized_param( + self, + model: "ModelMixin", + param_value: "torch.Tensor", + param_name: str, + state_dict: dict[str, Any], + **kwargs, + ) -> bool: + module, tensor_name = get_module_from_name(model, param_name) + from gemlite.core import GemLiteLinearTriton + + return tensor_name in GEMLITE_STATE_NAMES and isinstance(module, GemLiteLinearTriton) + + def create_quantized_param( + self, + model: "ModelMixin", + param_value: "torch.Tensor", + param_name: str, + target_device: "torch.device", + state_dict: dict[str, Any] | None = None, + unexpected_keys: list[str] | None = None, + **kwargs, + ): + module, tensor_name = get_module_from_name(model, param_name) + target_device = _normalize_torch_device(target_device) + if tensor_name not in GEMLITE_STATE_NAMES: + raise ValueError(f"`{param_name}` is not a GemLite serialized tensor.") + + value = state_dict[param_name] if state_dict is not None else param_value + value = value.to(target_device) + + setattr(module, tensor_name, torch.nn.Parameter(value, requires_grad=False)) + module._gemlite_loaded_param_names.add(tensor_name) + + def _process_model_before_weight_loading( + self, + model: "ModelMixin", + device_map, + keep_in_fp32_modules: list[str] = [], + **kwargs, + ): + """ + Replace `nn.Linear` modules with GemLite modules before the checkpoint tensors are loaded. The serialized + layout in the config creates correctly shaped GemLite tensors. Modules in `keep_in_fp32_modules` are excluded. + """ + self.modules_to_not_convert.extend(keep_in_fp32_modules) + self.modules_to_not_convert = [module for module in self.modules_to_not_convert if module is not None] + + _replace_with_gemlite_linear(model, self.modules_to_not_convert, self.quantization_config) + model.config.quantization_config = self.quantization_config + + def _process_model_after_weight_loading(self, model: "ModelMixin", **kwargs): + from gemlite.core import GemLiteLinearTriton + + has_gemlite_linear = False + for module in model.modules(): + if not isinstance(module, GemLiteLinearTriton): + continue + has_gemlite_linear = True + state_dict = { + name: getattr(module, name) + for name in GEMLITE_STATE_NAMES + if name in module._gemlite_loaded_param_names + } + module.load_state_dict(state_dict) + + if not has_gemlite_linear: + logger.warning("No linear modules are using GemLite.") + return model + + def get_state_dict_and_metadata( + self, state_dict: dict[str, Any], safe_serialization: bool = False + ) -> tuple[dict[str, Any], dict[str, Any]]: + # This hook runs from `ModelMixin.save_pretrained`, before the state dict is written. SafeTensors requires + # contiguous tensors, so return a self-consistent GemLite state dict by materializing `W_q` and updating its + # paired `data_contiguous` metadata together. The generic save path's later `.contiguous()` is then a no-op. + state_dict = state_dict.copy() + for name, w_q in state_dict.items(): + module_name, _, tensor_name = name.rpartition(".") + if tensor_name != "W_q": + continue + + metadata_name = f"{module_name}.metadata" if module_name else "metadata" + orig_shape_name = f"{module_name}.orig_shape" if module_name else "orig_shape" + if metadata_name not in state_dict or orig_shape_name not in state_dict: + continue + + metadata = state_dict[metadata_name] + data_contiguous = bool(metadata[-1].item()) + if not data_contiguous: + state_dict[name] = w_q.contiguous() + state_dict[metadata_name] = metadata.clone() + state_dict[metadata_name][-1] = 1 + + return state_dict, {} + + @property + def is_serializable(self): + return True + + @property + def is_trainable(self) -> bool: + return False + + @property + def is_compileable(self) -> bool: + return True diff --git a/src/diffusers/quantizers/quantization_config.py b/src/diffusers/quantizers/quantization_config.py index 72a0227d51df..521caecaa52e 100644 --- a/src/diffusers/quantizers/quantization_config.py +++ b/src/diffusers/quantizers/quantization_config.py @@ -45,6 +45,7 @@ class QuantizationMethod(str, Enum): BITS_AND_BYTES = "bitsandbytes" GGUF = "gguf" + GEMLITE = "gemlite" NUNCHAKU_LITE = "nunchaku_lite" TORCHAO = "torchao" QUANTO = "quanto" @@ -430,6 +431,153 @@ def __init__(self, compute_dtype: "torch.dtype" | None = None): self.compute_dtype = torch.float32 +@dataclass +class GemLiteConfig(QuantizationConfigMixin): + """Configuration class for GemLite quantization. + + GemLite provides Triton kernels for fast low-bit matrix multiplication. In Diffusers, this config replaces + supported `torch.nn.Linear` modules with GemLite linear layers so quantized weights can be executed with + GemLite kernels directly, without first materializing the weights back to full precision. + + GemLite in Diffusers only loads pre-quantized checkpoints whose linear layers were serialized as + `GemLiteLinearTriton` modules (`W_q`, `scales`, `zeros`, `metadata`, ... tensors). The config describes the + serialized tensor layout with `bits`, `group_size`, `packing_bitwidth`, `input_dtype`, `output_dtype`, + `scales_dtype`, and `zeros_dtype`, allowing Diffusers to construct correctly sized meta tensors before Accelerate + assigns modules to devices with `device_map="auto"`. + + Args: + compute_dtype (`torch.dtype`, *optional*, defaults to `torch.float16`): + The dtype used to load non-quantized modules. If `torch_dtype` is passed to `from_pretrained`, it must + match `compute_dtype`. + modules_to_not_convert (`list[str]` or `str`, *optional*): + The list of modules to not replace with GemLite linear layers. + format (`str`, *optional*): + Name of the serialized GemLite format, preserved as checkpoint metadata. + bits (`int`): + Number of bits per quantized weight. Must be one of 1, 2, 4, 8, or 16. + group_size (`int`): + Positive number of weights represented by each scale and zero-point entry in a pre-quantized checkpoint. + packing_bitwidth (`int`): + Bitwidth of each packed `W_q` storage element in a pre-quantized checkpoint. Must be 8, 16, 32, or 64 and + divisible by `bits`. + solver (`str`, *optional*): + Quantization solver used to produce the checkpoint, preserved as checkpoint metadata. + input_dtype (`str`): + GemLite input dtype stored in the serialized metadata of a pre-quantized checkpoint. + output_dtype (`str`): + GemLite output dtype stored in the serialized metadata of a pre-quantized checkpoint. + scales_dtype (`str`): + Storage dtype of the serialized `scales` tensor in a pre-quantized checkpoint. + zeros_dtype (`str`): + Storage dtype of the serialized `zeros` tensor in a pre-quantized checkpoint. + quantized_fqns (`list[str]`, *optional*): + Fully-qualified names of the `nn.Linear` modules replaced by GemLite linears. When omitted, every linear + module except those in `modules_to_not_convert` is replaced. + """ + + _SUPPORTED_SERIALIZED_DTYPES = ("fp16", "float16", "bf16", "bfloat16", "fp32", "float32") + _SUPPORTED_BITS = (1, 2, 4, 8, 16) + _SUPPORTED_PACKING_BITWIDTHS = (8, 16, 32, 64) + + def __init__( + self, + compute_dtype: "torch.dtype" | str | None = None, + modules_to_not_convert=None, + format: str | None = None, + bits: int | None = None, + group_size: int | None = None, + packing_bitwidth: int | None = None, + solver: str | None = None, + input_dtype: str | None = None, + output_dtype: str | None = None, + scales_dtype: str | None = None, + zeros_dtype: str | None = None, + quantized_fqns: list[str] | None = None, + **kwargs, + ): + self.quant_method = QuantizationMethod.GEMLITE + self.compute_dtype = compute_dtype + self.modules_to_not_convert = modules_to_not_convert + self.format = format + self.bits = bits + self.group_size = group_size + self.packing_bitwidth = packing_bitwidth + self.solver = solver + self.input_dtype = input_dtype + self.output_dtype = output_dtype + self.scales_dtype = scales_dtype + self.zeros_dtype = zeros_dtype + self.quantized_fqns = quantized_fqns + + # Keep the default as `None` in the signature because `torch` is an optional dependency and may not be imported + # when this module is loaded. Resolve it here instead of using `compute_dtype=torch.float16`. + if self.compute_dtype is None: + self.compute_dtype = torch.float16 + elif isinstance(self.compute_dtype, str): + self.compute_dtype = getattr(torch, self.compute_dtype) + + self.post_init() + + def post_init(self): + required_config_fields = ( + "bits", + "group_size", + "packing_bitwidth", + "input_dtype", + "output_dtype", + "scales_dtype", + "zeros_dtype", + ) + missing_config_fields = [field for field in required_config_fields if getattr(self, field) is None] + if missing_config_fields: + raise ValueError( + "Pre-quantized GemLite checkpoints require serialized layout fields in `quantization_config`: " + f"{', '.join(missing_config_fields)}." + ) + + if not all(isinstance(getattr(self, field), int) for field in ("bits", "group_size", "packing_bitwidth")): + raise ValueError("GemLite `bits`, `group_size`, and `packing_bitwidth` must be integers.") + if self.bits <= 0 or self.group_size <= 0 or self.packing_bitwidth <= 0: + raise ValueError("GemLite `bits`, `group_size`, and `packing_bitwidth` must be positive.") + if self.packing_bitwidth not in self._SUPPORTED_PACKING_BITWIDTHS: + raise ValueError( + f"Unsupported GemLite `packing_bitwidth`: {self.packing_bitwidth!r}. " + f"Supported values are: {list(self._SUPPORTED_PACKING_BITWIDTHS)}." + ) + if self.packing_bitwidth % self.bits: + raise ValueError("GemLite `packing_bitwidth` must be divisible by `bits`.") + if self.bits not in self._SUPPORTED_BITS: + raise ValueError( + f"Unsupported GemLite `bits`: {self.bits!r}. Supported values are: {list(self._SUPPORTED_BITS)}." + ) + + unsupported_dtype_fields = [ + field + for field in ("input_dtype", "output_dtype", "scales_dtype", "zeros_dtype") + if getattr(self, field) not in self._SUPPORTED_SERIALIZED_DTYPES + ] + if unsupported_dtype_fields: + raise ValueError( + "Unsupported GemLite serialized dtype in `quantization_config`: " + f"{', '.join(unsupported_dtype_fields)}." + ) + + def to_dict(self) -> dict[str, Any]: + output = copy.deepcopy(self.__dict__) + output["compute_dtype"] = str(output["compute_dtype"]).split(".")[1] + return output + + def to_diff_dict(self) -> dict[str, Any]: + config_dict = self.to_dict() + serializable_config_dict = {"quant_method": config_dict["quant_method"]} + for key, value in config_dict.items(): + if key == "quant_method": + continue + if (key == "compute_dtype" and value != "float16") or (key != "compute_dtype" and value is not None): + serializable_config_dict[key] = value + return serializable_config_dict + + @dataclass class NunchakuLiteQuantizationConfig(QuantizationConfigMixin): """Configuration for loading Nunchaku Lite checkpoints. diff --git a/src/diffusers/utils/__init__.py b/src/diffusers/utils/__init__.py index be63bed71196..beb391d0d0e6 100644 --- a/src/diffusers/utils/__init__.py +++ b/src/diffusers/utils/__init__.py @@ -79,6 +79,8 @@ is_flash_attn_version, is_flashpack_available, is_ftfy_available, + is_gemlite_available, + is_gemlite_version, is_gguf_available, is_gguf_version, is_google_colab, diff --git a/src/diffusers/utils/dummy_gemlite_objects.py b/src/diffusers/utils/dummy_gemlite_objects.py new file mode 100644 index 000000000000..7211bfcfd0d3 --- /dev/null +++ b/src/diffusers/utils/dummy_gemlite_objects.py @@ -0,0 +1,17 @@ +# This file is autogenerated by the command `make fix-copies`, do not edit. +from ..utils import DummyObject, requires_backends + + +class GemLiteConfig(metaclass=DummyObject): + _backends = ["gemlite"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["gemlite"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["gemlite"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["gemlite"]) diff --git a/src/diffusers/utils/import_utils.py b/src/diffusers/utils/import_utils.py index 30e00f74bc49..b09da8183d6b 100644 --- a/src/diffusers/utils/import_utils.py +++ b/src/diffusers/utils/import_utils.py @@ -204,6 +204,7 @@ def _is_package_available(pkg_name: str, get_dist_name: bool = False) -> tuple[b _accelerate_available, _accelerate_version = _is_package_available("accelerate") _xformers_available, _xformers_version = _is_package_available("xformers") _gguf_available, _gguf_version = _is_package_available("gguf") +_gemlite_available, _gemlite_version = _is_package_available("gemlite") _torchao_available, _torchao_version = _is_package_available("torchao") _bitsandbytes_available, _bitsandbytes_version = _is_package_available("bitsandbytes") _optimum_quanto_available, _optimum_quanto_version = _is_package_available("optimum", get_dist_name=True) @@ -354,6 +355,10 @@ def is_gguf_available(): return _gguf_available +def is_gemlite_available(): + return _gemlite_available + + def is_flashpack_available(): return _flashpack_available @@ -548,6 +553,11 @@ def is_av_available(): {0} requires the gguf library but it was not found in your environment. You can install it with pip: `pip install gguf` """ +GEMLITE_IMPORT_ERROR = """ +{0} requires the gemlite library but it was not found in your environment. You can install it with pip: `pip install +gemlite` +""" + TORCHAO_IMPORT_ERROR = """ {0} requires the torchao library but it was not found in your environment. You can install it with pip: `pip install torchao` @@ -603,6 +613,7 @@ def is_av_available(): ("sentencepiece", (is_sentencepiece_available, SENTENCEPIECE_IMPORT_ERROR)), ("imageio", (is_imageio_available, IMAGEIO_IMPORT_ERROR)), ("gguf", (is_gguf_available, GGUF_IMPORT_ERROR)), + ("gemlite", (is_gemlite_available, GEMLITE_IMPORT_ERROR)), ("torchao", (is_torchao_available, TORCHAO_IMPORT_ERROR)), ("quanto", (is_optimum_quanto_available, QUANTO_IMPORT_ERROR)), ("pytorch_retinaface", (is_pytorch_retinaface_available, PYTORCH_RETINAFACE_IMPORT_ERROR)), @@ -803,6 +814,22 @@ def is_bitsandbytes_version(operation: str, version: str): return compare_versions(parse(_bitsandbytes_version), operation, version) +@cache +def is_gemlite_version(operation: str, version: str): + """ + Compares the current gemlite version to a given reference with an operation. + + Args: + operation (`str`): + A string representation of an operator, such as `">"` or `"<="` + version (`str`): + A version string + """ + if not _gemlite_available: + return False + return compare_versions(parse(_gemlite_version), operation, version) + + @cache def is_gguf_version(operation: str, version: str): """ diff --git a/src/diffusers/utils/testing_utils.py b/src/diffusers/utils/testing_utils.py index ce3f1d765d86..01a8c52c8433 100644 --- a/src/diffusers/utils/testing_utils.py +++ b/src/diffusers/utils/testing_utils.py @@ -35,6 +35,7 @@ is_accelerate_available, is_bitsandbytes_available, is_compel_available, + is_gemlite_available, is_gguf_available, is_kernels_available, is_note_seq_available, @@ -535,6 +536,13 @@ def require_quanto(test_case): return unittest.skipUnless(is_optimum_quanto_available(), "test requires quanto")(test_case) +def require_gemlite(test_case): + """ + Decorator marking a test that requires gemlite. These tests are skipped when gemlite isn't installed. + """ + return unittest.skipUnless(is_gemlite_available(), "test requires gemlite")(test_case) + + def require_accelerate(test_case): """ Decorator marking a test that requires accelerate. These tests are skipped when accelerate isn't installed. @@ -625,6 +633,18 @@ def decorator(test_case): return decorator +def require_gemlite_version_greater_or_equal(gemlite_version): + def decorator(test_case): + correct_gemlite_version = is_gemlite_available() and version.parse( + version.parse(importlib.metadata.version("gemlite")).base_version + ) >= version.parse(gemlite_version) + return unittest.skipUnless( + correct_gemlite_version, f"Test requires gemlite with the version greater than {gemlite_version}." + )(test_case) + + return decorator + + def require_torchao_version_greater_or_equal(torchao_version): def decorator(test_case): correct_torchao_version = is_torchao_available() and version.parse( diff --git a/tests/quantization/gemlite/__init__.py b/tests/quantization/gemlite/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/tests/quantization/gemlite/test_gemlite.py b/tests/quantization/gemlite/test_gemlite.py new file mode 100644 index 000000000000..601ab87b06f8 --- /dev/null +++ b/tests/quantization/gemlite/test_gemlite.py @@ -0,0 +1,727 @@ +import contextlib +import gc +import tempfile +import unittest +from unittest import mock + +import numpy as np +import torch +import torch.nn as nn + +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.models.modeling_utils import ModelMixin +from diffusers.quantizers.auto import DiffusersAutoQuantizer +from diffusers.quantizers.gemlite.gemlite_quantizer import ( + _is_in_skip_modules, + _normalize_torch_device, + _replace_with_gemlite_linear, +) +from diffusers.quantizers.quantization_config import GemLiteConfig, QuantizationMethod +from diffusers.utils import get_module_from_name, is_gemlite_available + +from ...testing_utils import ( + backend_empty_cache, + enable_full_determinism, + nightly, + require_accelerate, + require_gemlite, + require_gemlite_version_greater_or_equal, + require_torch_gpu, + torch_device, +) + + +_GEMLITE_SERIALIZED_STATE_NAMES = {"W_q", "scales", "zeros", "metadata", "orig_shape", "meta_scale"} + + +enable_full_determinism() + + +if is_gemlite_available() and torch.cuda.is_available(): + from gemlite.core import DType, GemLiteLinearTriton + + +def _get_gemlite_config(**kwargs): + config = { + "bits": 2, + "group_size": 16, + "packing_bitwidth": 8, + "input_dtype": "fp16", + "output_dtype": "fp16", + "scales_dtype": "fp32", + "zeros_dtype": "fp32", + } + config.update(kwargs) + return GemLiteConfig(**config) + + +def _create_packed_gemlite_state_dict( + in_features=64, + out_features=32, + w_nbits=4, + group_size=64, + packing_bitwidth=32, + scales_dtype=torch.float16, + zeros_dtype=torch.float16, + device=torch_device, + prefix="", +): + source_layer = GemLiteLinearTriton( + W_nbits=w_nbits, + group_size=group_size, + in_features=in_features, + out_features=out_features, + input_dtype=DType.FP16, + output_dtype=DType.FP16, + ) + weight = ( + torch.arange(out_features * in_features, dtype=torch.int32, device=device) + .remainder(2**w_nbits) + .to(torch.uint8) + ) + scales = torch.ones((out_features, in_features // group_size), dtype=scales_dtype, device=device) + zeros = torch.full( + (out_features, in_features // group_size), + (2**w_nbits - 1) // 2, + dtype=zeros_dtype, + device=device, + ) + + source_layer.pack(weight, scales, zeros, bias=None, packing_bitwidth=packing_bitwidth) + + return {f"{prefix}.{name}" if prefix else name: value for name, value in source_layer.state_dict().items()} + + +def _save_packed_gemlite_transformer(transformer, save_directory): + group_size = 32 + quantized_fqns = [ + name + for name, module in transformer.named_modules() + if isinstance(module, nn.Linear) and module.in_features == group_size + ] + quantization_config = GemLiteConfig( + bits=8, + group_size=group_size, + packing_bitwidth=8, + input_dtype="fp16", + output_dtype="fp16", + scales_dtype="fp32", + zeros_dtype="fp32", + quantized_fqns=quantized_fqns, + ) + + from gemlite.helper import A16W8_INT8 + + for name in quantized_fqns: + parent, child_name = get_module_from_name(transformer, name) + linear = getattr(parent, child_name) + setattr( + parent, + child_name, + A16W8_INT8(device=str(linear.weight.device), dtype=linear.weight.dtype).from_linear(linear), + ) + + transformer.register_to_config(quantization_config=quantization_config) + transformer.hf_quantizer = DiffusersAutoQuantizer.from_config(quantization_config, pre_quantized=True) + transformer.save_pretrained(save_directory, safe_serialization=True) + + return quantized_fqns + + +class GemLiteConfigTest(unittest.TestCase): + def test_config_defaults(self): + config = _get_gemlite_config() + + self.assertEqual(config.quant_method, QuantizationMethod.GEMLITE) + self.assertEqual(config.compute_dtype, torch.float16) + self.assertEqual( + config.to_diff_dict(), + { + "quant_method": QuantizationMethod.GEMLITE, + "bits": 2, + "group_size": 16, + "packing_bitwidth": 8, + "input_dtype": "fp16", + "output_dtype": "fp16", + "scales_dtype": "fp32", + "zeros_dtype": "fp32", + }, + ) + + def test_config_requires_serialized_layout(self): + with self.assertRaisesRegex(ValueError, "require serialized layout fields"): + GemLiteConfig() + + def test_config_accepts_64_bit_packing(self): + config = _get_gemlite_config(packing_bitwidth=64) + + self.assertEqual(config.packing_bitwidth, 64) + + def test_config_rejects_invalid_serialized_layout(self): + with self.assertRaisesRegex(ValueError, "must be positive"): + _get_gemlite_config(bits=0) + with self.assertRaisesRegex(ValueError, "Unsupported GemLite `packing_bitwidth`"): + _get_gemlite_config(packing_bitwidth=4) + with self.assertRaisesRegex(ValueError, "must be divisible by `bits`"): + _get_gemlite_config(bits=3) + with self.assertRaisesRegex(ValueError, "Unsupported GemLite `bits`"): + _get_gemlite_config(bits=32, packing_bitwidth=32) + with self.assertRaisesRegex(ValueError, "Unsupported GemLite serialized dtype"): + _get_gemlite_config(input_dtype="fp8") + + def test_config_from_dict(self): + config = DiffusersAutoQuantizer.from_dict( + { + "quant_method": "gemlite", + "compute_dtype": "bfloat16", + "modules_to_not_convert": ["proj_out"], + "bits": 2, + "group_size": 16, + "packing_bitwidth": 8, + "input_dtype": "fp16", + "output_dtype": "fp16", + "scales_dtype": "fp32", + "zeros_dtype": "fp32", + } + ) + + self.assertIsInstance(config, GemLiteConfig) + self.assertEqual(config.compute_dtype, torch.bfloat16) + self.assertEqual(config.modules_to_not_convert, ["proj_out"]) + + def test_config_round_trip(self): + config = GemLiteConfig( + compute_dtype=torch.bfloat16, + modules_to_not_convert=["proj_out"], + format="gemlite-int2-ternary-g128", + bits=2, + group_size=128, + packing_bitwidth=8, + solver="ternary", + input_dtype="fp16", + output_dtype="fp16", + scales_dtype="fp32", + zeros_dtype="fp32", + quantized_fqns=["blocks.0.proj"], + ) + + restored = DiffusersAutoQuantizer.from_dict(config.to_dict()) + + self.assertIsInstance(restored, GemLiteConfig) + self.assertEqual(restored.compute_dtype, torch.bfloat16) + self.assertEqual(restored.modules_to_not_convert, ["proj_out"]) + self.assertEqual(restored.format, "gemlite-int2-ternary-g128") + self.assertEqual(restored.bits, 2) + self.assertEqual(restored.group_size, 128) + self.assertEqual(restored.packing_bitwidth, 8) + self.assertEqual(restored.solver, "ternary") + self.assertEqual(restored.input_dtype, "fp16") + self.assertEqual(restored.output_dtype, "fp16") + self.assertEqual(restored.scales_dtype, "fp32") + self.assertEqual(restored.zeros_dtype, "fp32") + self.assertEqual(restored.quantized_fqns, ["blocks.0.proj"]) + + def test_quantizer_uses_compute_dtype_when_torch_dtype_is_not_provided(self): + quantizer = DiffusersAutoQuantizer.from_config(_get_gemlite_config(compute_dtype=torch.bfloat16)) + + self.assertEqual(quantizer.update_torch_dtype(None), torch.bfloat16) + self.assertEqual(quantizer.update_torch_dtype(torch.bfloat16), torch.bfloat16) + + def test_quantizer_rejects_mismatched_torch_dtype(self): + quantizer = DiffusersAutoQuantizer.from_config(_get_gemlite_config(compute_dtype=torch.bfloat16)) + + with self.assertRaisesRegex(ValueError, "must match `GemLiteConfig.compute_dtype`"): + quantizer.update_torch_dtype(torch.float16) + + def test_quantizer_disables_parallel_loading(self): + quantizer = DiffusersAutoQuantizer.from_config(_get_gemlite_config()) + + self.assertFalse(quantizer.supports_parallel_loading) + + def test_quantizer_rejects_disk_offloading(self): + quantizer = DiffusersAutoQuantizer.from_config(_get_gemlite_config()) + + with self.assertRaisesRegex(ValueError, "does not support disk offloading"): + quantizer.validate_environment(device_map={"transformer_blocks.0": "disk"}) + + def test_quantizer_rejects_unquantized_checkpoints(self): + with self.assertRaisesRegex(ValueError, "only supports loading pre-quantized checkpoints"): + DiffusersAutoQuantizer.from_config(_get_gemlite_config(), pre_quantized=False) + + +class GemLiteQuantizerHelperTest(unittest.TestCase): + def test_is_in_skip_modules_matches_exact_names_and_children(self): + modules_to_not_convert = ["proj_out", "blocks.0"] + + self.assertTrue(_is_in_skip_modules("proj_out", modules_to_not_convert)) + self.assertTrue(_is_in_skip_modules("blocks.0.attn.to_q", modules_to_not_convert)) + self.assertFalse(_is_in_skip_modules("proj_output", modules_to_not_convert)) + self.assertFalse(_is_in_skip_modules("blocks.01.attn.to_q", modules_to_not_convert)) + + def test_normalize_torch_device_handles_int_device_map_values(self): + self.assertEqual(_normalize_torch_device(0), torch.device("cuda:0")) + self.assertEqual(_normalize_torch_device("cpu"), torch.device("cpu")) + self.assertEqual(_normalize_torch_device(torch.device("meta")), torch.device("meta")) + + +@require_gemlite +@require_torch_gpu +class GemLiteQuantizerEnvironmentTest(unittest.TestCase): + @contextlib.contextmanager + def _mock_module_availability(self, gemlite, cuda, old_gemlite): + with contextlib.ExitStack() as stack: + stack.enter_context( + mock.patch("diffusers.quantizers.gemlite.gemlite_quantizer.is_gemlite_available", return_value=gemlite) + ) + stack.enter_context( + mock.patch( + "diffusers.quantizers.gemlite.gemlite_quantizer.is_gemlite_version", + return_value=old_gemlite, + ) + ) + stack.enter_context(mock.patch("torch.cuda.is_available", return_value=cuda)) + yield + + def test_validate_environment_checks_pre_quantized_serialized_state(self): + quantizer = DiffusersAutoQuantizer.from_config(_get_gemlite_config(), pre_quantized=True) + + with self._mock_module_availability(gemlite=True, cuda=True, old_gemlite=False): + quantizer.validate_environment() + + def test_validate_environment_rejects_missing_gemlite(self): + quantizer = DiffusersAutoQuantizer.from_config(_get_gemlite_config(), pre_quantized=True) + + with self._mock_module_availability(gemlite=False, cuda=True, old_gemlite=False): + with self.assertRaisesRegex(ImportError, "requires the gemlite library"): + quantizer.validate_environment() + + def test_validate_environment_rejects_old_gemlite_version(self): + quantizer = DiffusersAutoQuantizer.from_config(_get_gemlite_config(), pre_quantized=True) + + with self._mock_module_availability(gemlite=True, cuda=True, old_gemlite=True): + with self.assertRaisesRegex(ImportError, "requires gemlite>=0.6.0"): + quantizer.validate_environment() + + def test_validate_environment_rejects_broken_gemlite_core_import(self): + quantizer = DiffusersAutoQuantizer.from_config(_get_gemlite_config(), pre_quantized=True) + original_import = __import__ + + def import_with_broken_gemlite_core(name, *args, **kwargs): + if name == "gemlite.core": + raise RuntimeError("broken gemlite core") + return original_import(name, *args, **kwargs) + + with self._mock_module_availability(gemlite=True, cuda=True, old_gemlite=False): + with mock.patch("builtins.__import__", side_effect=import_with_broken_gemlite_core): + with self.assertRaisesRegex(ImportError, "core linear module could not be imported"): + quantizer.validate_environment() + + +@require_gemlite +@require_gemlite_version_greater_or_equal("0.6.0") +@require_torch_gpu +class GemLiteQuantizerTest(unittest.TestCase): + def test_replace_with_gemlite_linear(self): + from diffusers import FluxTransformer2DModel + + skip_patterns = [ + "proj_out", + "x_embedder", + "context_embedder", + "time_text_embed", + "time_guidance_embed", + "norm_out", + "double_stream_modulation_img", + "double_stream_modulation_txt", + "single_stream_modulation", + ] + model = FluxTransformer2DModel( + in_channels=4, + num_layers=1, + num_single_layers=1, + attention_head_dim=16, + num_attention_heads=2, + joint_attention_dim=32, + pooled_projection_dim=32, + axes_dims_rope=[4, 4, 8], + ).to(torch_device) + original_linear_names = {name for name, module in model.named_modules() if isinstance(module, nn.Linear)} + skipped_linear_names = {name for name in original_linear_names if _is_in_skip_modules(name, skip_patterns)} + + replaced = _replace_with_gemlite_linear( + model, + skip_patterns, + GemLiteConfig( + bits=2, + group_size=16, + packing_bitwidth=8, + input_dtype="fp16", + output_dtype="fp16", + scales_dtype="fp32", + zeros_dtype="fp32", + ), + ) + + self.assertEqual(replaced, 20) + for name, module in model.named_modules(): + if name in skipped_linear_names: + self.assertIs(type(module), nn.Linear) + elif name in original_linear_names: + self.assertIsInstance(module, GemLiteLinearTriton) + + def test_replace_preserves_bias_state_key(self): + model = nn.Sequential(nn.Linear(32, 32, bias=True), nn.Linear(32, 32, bias=False)).to(torch_device) + + replaced = _replace_with_gemlite_linear( + model, + [], + GemLiteConfig( + bits=2, + group_size=16, + packing_bitwidth=8, + input_dtype="fp16", + output_dtype="fp16", + scales_dtype="fp32", + zeros_dtype="fp32", + ), + ) + + self.assertEqual(replaced, 2) + state_dict = model.state_dict() + assert _GEMLITE_SERIALIZED_STATE_NAMES.issubset( + {name.removeprefix("0.") for name in state_dict if name.startswith("0.")} + ) + self.assertIn("0.bias", state_dict) + self.assertNotIn("1.bias", state_dict) + + @require_accelerate + def test_prequantized_auto_device_map_uses_packed_state_size(self): + from accelerate.utils import compute_module_sizes + + from diffusers.models.model_loading_utils import _determine_device_map + + class GemLiteDeviceMapTestModel(ModelMixin, ConfigMixin): + _no_split_modules = [] + + @register_to_config + def __init__(self): + super().__init__() + self.proj = nn.Linear(128, 32, bias=False) + + packed_state = _create_packed_gemlite_state_dict( + in_features=128, + out_features=32, + w_nbits=2, + group_size=128, + packing_bitwidth=8, + scales_dtype=torch.float32, + zeros_dtype=torch.float32, + device=torch_device, + ) + packed_size = sum(value.numel() * value.element_size() for value in packed_state.values()) + + quantizer = DiffusersAutoQuantizer.from_config( + GemLiteConfig( + format="gemlite-int2-ternary-g128", + bits=2, + group_size=128, + packing_bitwidth=8, + solver="ternary", + input_dtype="fp16", + output_dtype="fp16", + scales_dtype="fp32", + zeros_dtype="fp32", + quantized_fqns=["proj"], + ), + pre_quantized=True, + ) + gemlite_model = GemLiteDeviceMapTestModel().to("meta") + quantizer.preprocess_model(gemlite_model, device_map="auto", keep_in_fp32_modules=[]) + placeholder_state_dict = gemlite_model.state_dict() + for name, value in packed_state.items(): + placeholder = placeholder_state_dict[f"proj.{name}"] + self.assertEqual(placeholder.shape, value.shape) + self.assertEqual(placeholder.dtype, value.dtype) + + module_sizes = compute_module_sizes( + gemlite_model, + dtype=torch.float16, + special_dtypes=quantizer.get_special_dtypes_update(gemlite_model, torch.float16), + ) + self.assertEqual(module_sizes["proj"], packed_size) + with self.assertWarnsRegex(UserWarning, "Current model requires .* bytes of buffer for offloaded layers"): + gemlite_device_map = _determine_device_map( + gemlite_model, + "auto", + {0: packed_size // 2, "cpu": packed_size * 2}, + torch.float16, + hf_quantizer=quantizer, + ) + + self.assertNotIn(0, gemlite_device_map.values()) + + def test_process_model_before_weight_loading_replaces_pre_quantized_linears(self): + from diffusers import FluxTransformer2DModel + + skip_patterns = [ + "proj_out", + "x_embedder", + "context_embedder", + "time_text_embed", + "time_guidance_embed", + "norm_out", + "double_stream_modulation_img", + "double_stream_modulation_txt", + "single_stream_modulation", + ] + quantization_config = GemLiteConfig( + modules_to_not_convert=skip_patterns[:4], + bits=2, + group_size=16, + packing_bitwidth=8, + input_dtype="fp16", + output_dtype="fp16", + scales_dtype="fp32", + zeros_dtype="fp32", + ) + quantizer = DiffusersAutoQuantizer.from_config(quantization_config, pre_quantized=True) + model = FluxTransformer2DModel( + in_channels=4, + num_layers=1, + num_single_layers=1, + attention_head_dim=16, + num_attention_heads=2, + joint_attention_dim=32, + pooled_projection_dim=32, + axes_dims_rope=[4, 4, 8], + ).to(torch_device) + original_linears = {name: module for name, module in model.named_modules() if isinstance(module, nn.Linear)} + skipped_linear_names = {name for name in original_linears if _is_in_skip_modules(name, skip_patterns)} + + quantizer._process_model_before_weight_loading( + model, device_map=None, keep_in_fp32_modules=[*skip_patterns[4:], None] + ) + + processed_modules = dict(model.named_modules()) + for name, linear in original_linears.items(): + if name in skipped_linear_names: + self.assertIs(processed_modules[name], linear) + else: + self.assertIsInstance(processed_modules[name], GemLiteLinearTriton) + self.assertEqual(quantizer.modules_to_not_convert, skip_patterns) + self.assertIs(model.config.quantization_config, quantization_config) + + def test_check_if_quantized_param_pre_quantized(self): + quantizer = DiffusersAutoQuantizer.from_config(_get_gemlite_config(), pre_quantized=True) + + plain_model = nn.Sequential(nn.Linear(32, 32, bias=False)).to(torch_device) + self.assertFalse( + quantizer.check_if_quantized_param( + plain_model, torch.ones(4, 4, dtype=torch.uint8, device=torch_device), "0.W_q", {} + ) + ) + + gemlite_model = nn.Sequential(GemLiteLinearTriton()).to(torch_device) + self.assertFalse( + quantizer.check_if_quantized_param(gemlite_model, torch.ones(32, 32, device=torch_device), "0.weight", {}) + ) + self.assertTrue( + quantizer.check_if_quantized_param( + gemlite_model, torch.ones(4, 4, dtype=torch.uint8, device=torch_device), "0.W_q", {} + ) + ) + + def test_create_quantized_param_loads_pre_quantized_state(self): + in_features = 64 + out_features = 32 + gemlite_state_dict = _create_packed_gemlite_state_dict(in_features, out_features, device=torch_device) + model = nn.Sequential(GemLiteLinearTriton()).to(torch_device) + model[0]._gemlite_loaded_param_names = set() + quantizer = DiffusersAutoQuantizer.from_config(_get_gemlite_config(), pre_quantized=True) + + for name, value in gemlite_state_dict.items(): + with self.subTest(name=name): + quantizer.create_quantized_param(model, value, f"0.{name}", torch_device) + + loaded_value = getattr(model[0], name) + self.assertIsInstance(loaded_value, nn.Parameter) + self.assertIs(model[0]._parameters[name], loaded_value) + self.assertEqual(loaded_value.device, value.device) + self.assertEqual(loaded_value.dtype, value.dtype) + self.assertFalse(loaded_value.requires_grad) + self.assertTrue(torch.equal(loaded_value, value)) + + self.assertEqual(model[0]._gemlite_loaded_param_names, set(gemlite_state_dict)) + + def test_create_quantized_param_preserves_serialized_dtypes(self): + quantizer = DiffusersAutoQuantizer.from_config(_get_gemlite_config(), pre_quantized=True) + model = nn.Sequential(GemLiteLinearTriton()).to(torch_device) + model[0]._gemlite_loaded_param_names = set() + serialized_state = { + "0.W_q": torch.arange(4, dtype=torch.float16, device=torch_device).reshape(2, 2).to(torch.float8_e5m2), + "0.scales": torch.tensor([[0.12345679], [0.9876543]], dtype=torch.float32, device=torch_device), + } + + for param_name, original_value in serialized_state.items(): + with self.subTest(param_name=param_name): + quantizer.create_quantized_param( + model, + original_value.to(torch.float16), + param_name, + torch_device, + state_dict=serialized_state, + ) + + loaded_value = getattr(model[0], param_name.removeprefix("0.")) + self.assertEqual(loaded_value.dtype, original_value.dtype) + self.assertTrue(torch.equal(loaded_value, original_value)) + + def test_get_state_dict_and_metadata_sets_data_contiguous_for_serialization(self): + quantizer = DiffusersAutoQuantizer.from_config(_get_gemlite_config(), pre_quantized=True) + w_q = torch.arange(12, dtype=torch.uint8, device=torch_device).reshape(3, 4).t() + metadata = torch.tensor([0, 8, 4, 255, 1, 1, 1, 0, 0, 0, 2, 0], dtype=torch.int32, device=torch_device) + state_dict = { + "0.W_q": w_q, + "0.metadata": metadata, + "0.orig_shape": torch.tensor([3, 4], dtype=torch.int32, device=torch_device), + } + + serialized_state_dict, _ = quantizer.get_state_dict_and_metadata(state_dict, safe_serialization=True) + + self.assertFalse(w_q.is_contiguous()) + self.assertEqual(metadata[-1].item(), 0) + self.assertTrue(serialized_state_dict["0.W_q"].is_contiguous()) + self.assertTrue(torch.equal(serialized_state_dict["0.W_q"], w_q)) + self.assertEqual(serialized_state_dict["0.metadata"][-1].item(), 1) + + def test_create_quantized_param_rejects_non_gemlite_serialized_name(self): + model = nn.Sequential(GemLiteLinearTriton()).to(torch_device) + quantizer = DiffusersAutoQuantizer.from_config(_get_gemlite_config(), pre_quantized=True) + + with self.assertRaisesRegex(ValueError, "is not a GemLite serialized tensor"): + quantizer.create_quantized_param(model, torch.ones(32, 64, device=torch_device), "0.weight", torch_device) + + +@nightly +@require_gemlite +@require_gemlite_version_greater_or_equal("0.6.0") +@require_torch_gpu +@require_accelerate +class GemLiteKrea2TransformerIntegrationTests(unittest.TestCase): + model_id = "hf-internal-testing/tiny-krea2-modular-pipe" + torch_dtype = torch.float16 + maximum_quantized_memory_fraction = 0.8 + + def tearDown(self): + gc.collect() + backend_empty_cache(torch_device) + + def test_inference_matches_unquantized_transformer(self): + from diffusers import Krea2Transformer2DModel + + reference_transformer = Krea2Transformer2DModel.from_pretrained( + self.model_id, + subfolder="transformer", + dtype=self.torch_dtype, + ).to(torch_device) + unquantized_transformer_memory = reference_transformer.get_memory_footprint() + inputs = { + "hidden_states": torch.randn((1, 4, 16), device=torch_device, dtype=self.torch_dtype), + "encoder_hidden_states": torch.randn((1, 4, 12, 16), device=torch_device, dtype=self.torch_dtype), + "timestep": torch.tensor([0.5], device=torch_device, dtype=self.torch_dtype), + "position_ids": torch.zeros((8, 3), device=torch_device), + "encoder_attention_mask": torch.tensor([[True, True, True, False]], device=torch_device), + } + with torch.inference_mode(): + reference_output = reference_transformer(**inputs).sample + + with tempfile.TemporaryDirectory() as model_dir: + quantized_fqns = _save_packed_gemlite_transformer(reference_transformer, model_dir) + self.assertTrue(quantized_fqns) + + del reference_transformer + gc.collect() + backend_empty_cache(torch_device) + + transformer = Krea2Transformer2DModel.from_pretrained( + model_dir, + dtype=self.torch_dtype, + device_map={"": torch_device}, + ) + + self.assertTrue(any(isinstance(module, GemLiteLinearTriton) for module in transformer.modules())) + self.assertLessEqual( + transformer.get_memory_footprint(), + unquantized_transformer_memory * self.maximum_quantized_memory_fraction, + ) + with torch.inference_mode(): + gemlite_output = transformer(**inputs).sample + + torch.testing.assert_close(gemlite_output, reference_output, rtol=5e-2, atol=5e-2) + + +@nightly +@require_gemlite +@require_gemlite_version_greater_or_equal("0.6.0") +@require_torch_gpu +@require_accelerate +class GemLiteFluxPipelineIntegrationTests(unittest.TestCase): + model_id = "hf-internal-testing/tiny-flux-pipe" + torch_dtype = torch.float16 + maximum_quantized_memory_fraction = 0.8 + + def tearDown(self): + gc.collect() + backend_empty_cache(torch_device) + + def test_inference_matches_unquantized_pipeline(self): + from diffusers import FluxPipeline, FluxTransformer2DModel + + pipe_inputs = { + "prompt": "a photo of a cat", + "height": 32, + "width": 32, + "num_inference_steps": 2, + "output_type": "np", + } + reference_pipe = FluxPipeline.from_pretrained(self.model_id, dtype=self.torch_dtype).to(torch_device) + unquantized_transformer_memory = reference_pipe.transformer.get_memory_footprint() + reference_pipe.set_progress_bar_config(disable=True) + with torch.inference_mode(): + reference_output = reference_pipe( + **pipe_inputs, + generator=torch.Generator(device=torch_device).manual_seed(0), + ).images + + with tempfile.TemporaryDirectory() as model_dir: + quantized_fqns = _save_packed_gemlite_transformer(reference_pipe.transformer, model_dir) + self.assertTrue(quantized_fqns) + + del reference_pipe + gc.collect() + backend_empty_cache(torch_device) + + transformer = FluxTransformer2DModel.from_pretrained( + model_dir, + dtype=self.torch_dtype, + device_map={"": torch_device}, + ) + + self.assertTrue(any(isinstance(module, GemLiteLinearTriton) for module in transformer.modules())) + self.assertLessEqual( + transformer.get_memory_footprint(), + unquantized_transformer_memory * self.maximum_quantized_memory_fraction, + ) + gemlite_pipe = FluxPipeline.from_pretrained( + self.model_id, + transformer=transformer, + dtype=self.torch_dtype, + ).to(torch_device) + gemlite_pipe.set_progress_bar_config(disable=True) + with torch.inference_mode(): + gemlite_output = gemlite_pipe( + **pipe_inputs, + generator=torch.Generator(device=torch_device).manual_seed(0), + ).images + + np.testing.assert_allclose(gemlite_output, reference_output, rtol=5e-2, atol=5e-2) diff --git a/tests/testing_utils.py b/tests/testing_utils.py index 0fb03e92028f..5b79e1907832 100644 --- a/tests/testing_utils.py +++ b/tests/testing_utils.py @@ -36,6 +36,7 @@ is_bitsandbytes_available, is_compel_available, is_flashpack_available, + is_gemlite_available, is_gguf_available, is_kernels_available, is_note_seq_available, @@ -738,6 +739,13 @@ def require_quanto(test_case): return pytest.mark.skipif(not is_optimum_quanto_available(), reason="test requires quanto")(test_case) +def require_gemlite(test_case): + """ + Decorator marking a test that requires gemlite. These tests are skipped when gemlite isn't installed. + """ + return pytest.mark.skipif(not is_gemlite_available(), reason="test requires gemlite")(test_case) + + def require_accelerate(test_case): """ Decorator marking a test that requires accelerate. These tests are skipped when accelerate isn't installed. @@ -837,6 +845,19 @@ def decorator(test_case): return decorator +def require_gemlite_version_greater_or_equal(gemlite_version): + def decorator(test_case): + correct_gemlite_version = is_gemlite_available() and version.parse( + version.parse(importlib.metadata.version("gemlite")).base_version + ) >= version.parse(gemlite_version) + return pytest.mark.skipif( + not correct_gemlite_version, + reason=f"Test requires gemlite with the version greater than {gemlite_version}.", + )(test_case) + + return decorator + + def require_torchao_version_greater_or_equal(torchao_version): def decorator(test_case): correct_torchao_version = is_torchao_available() and version.parse( From 77a21e21e94ca677410385bf0fa8dcef8b8f95fc Mon Sep 17 00:00:00 2001 From: Gabriel Date: Mon, 27 Jul 2026 22:57:01 -0500 Subject: [PATCH 02/10] Refactor GemLite pre-quantized loading --- .../quantizers/gemlite/gemlite_quantizer.py | 141 ++++++++---------- tests/quantization/gemlite/test_gemlite.py | 36 +++-- 2 files changed, 90 insertions(+), 87 deletions(-) diff --git a/src/diffusers/quantizers/gemlite/gemlite_quantizer.py b/src/diffusers/quantizers/gemlite/gemlite_quantizer.py index b1fe16ec5e1f..db5c27ee9318 100644 --- a/src/diffusers/quantizers/gemlite/gemlite_quantizer.py +++ b/src/diffusers/quantizers/gemlite/gemlite_quantizer.py @@ -33,97 +33,79 @@ def _normalize_torch_device(device: Any) -> "torch.device": return torch.device(device) -def _replace_with_gemlite_linear( - model: "ModelMixin", modules_to_not_convert: list[str], quantization_config: "GemLiteConfig" -) -> int: - """ - Replace eligible `nn.Linear` modules in `model` with GemLite modules whose serialized tensors have meta shapes. +class _GemLiteDiffusersProcessor: + """Implement GemLite's `patch_model` protocol for Diffusers' pre-quantized loading path.""" + + def __init__(self, quantization_config: "GemLiteConfig", modules_to_not_convert: list[str]): + from gemlite.core import DType + from gemlite.dtypes import DTYPE_TO_TORCH, PACKING_BITWIDTH_TO_TORCH_DTYPE + + gemlite_dtypes = { + "fp16": DType.FP16, + "float16": DType.FP16, + "bf16": DType.BF16, + "bfloat16": DType.BF16, + "fp32": DType.FP32, + "float32": DType.FP32, + } + self.bits = quantization_config.bits + self.group_size = quantization_config.group_size + self.input_dtype = gemlite_dtypes[quantization_config.input_dtype] + self.output_dtype = gemlite_dtypes[quantization_config.output_dtype] + self.scales_gemlite_dtype = gemlite_dtypes[quantization_config.scales_dtype] + self.scales_dtype = DTYPE_TO_TORCH[self.scales_gemlite_dtype.value] + self.zeros_dtype = DTYPE_TO_TORCH[gemlite_dtypes[quantization_config.zeros_dtype].value] + self.packed_dtype = PACKING_BITWIDTH_TO_TORCH_DTYPE[quantization_config.packing_bitwidth] + self.elements_per_sample = quantization_config.packing_bitwidth // quantization_config.bits + self.modules_to_not_convert = modules_to_not_convert + self.quantized_fqns = ( + set(quantization_config.quantized_fqns) if quantization_config.quantized_fqns is not None else None + ) - Returns the number of replaced modules. Modules in `modules_to_not_convert` are left unchanged. - """ - from gemlite.core import DType, GemLiteLinearTriton - from gemlite.dtypes import DTYPE_TO_TORCH, PACKING_BITWIDTH_TO_TORCH_DTYPE - - gemlite_dtypes = { - "fp16": DType.FP16, - "float16": DType.FP16, - "bf16": DType.BF16, - "bfloat16": DType.BF16, - "fp32": DType.FP32, - "float32": DType.FP32, - } - input_dtype = gemlite_dtypes[quantization_config.input_dtype] - output_dtype = gemlite_dtypes[quantization_config.output_dtype] - scales_gemlite_dtype = gemlite_dtypes[quantization_config.scales_dtype] - scales_dtype = DTYPE_TO_TORCH[scales_gemlite_dtype.value] - zeros_dtype = DTYPE_TO_TORCH[gemlite_dtypes[quantization_config.zeros_dtype].value] - packed_dtype = PACKING_BITWIDTH_TO_TORCH_DTYPE[quantization_config.packing_bitwidth] - - elements_per_sample = quantization_config.packing_bitwidth // quantization_config.bits - quantized_fqns = ( - set(quantization_config.quantized_fqns) if quantization_config.quantized_fqns is not None else None - ) - - def initialize_serialized_tensors(gemlite_linear: "nn.Module", linear: "nn.Linear") -> None: - gemlite_linear.elements_per_sample = elements_per_sample - gemlite_linear.meta_dtype = scales_gemlite_dtype + def from_linear(self, linear: "nn.Linear") -> "nn.Module": + from gemlite.core import GemLiteLinearTriton + + should_skip_linear = ( + self.quantized_fqns is not None and linear.name not in self.quantized_fqns + ) or _is_in_skip_modules(linear.name, self.modules_to_not_convert) + if should_skip_linear: + return linear + + gemlite_linear = GemLiteLinearTriton( + W_nbits=self.bits, + group_size=self.group_size, + in_features=linear.in_features, + out_features=linear.out_features, + input_dtype=self.input_dtype, + output_dtype=self.output_dtype, + ).to(linear.weight.device) + gemlite_linear.elements_per_sample = self.elements_per_sample + gemlite_linear.meta_dtype = self.scales_gemlite_dtype gemlite_linear.channel_scale_mode = 0 gemlite_linear.W_group_mode = 0 gemlite_linear.data_contiguous = True gemlite_linear.W_q = torch.empty( - linear.in_features // elements_per_sample, + linear.in_features // self.elements_per_sample, linear.out_features, - dtype=packed_dtype, + dtype=self.packed_dtype, device=linear.weight.device, ) gemlite_linear.scales = torch.empty( - linear.in_features // quantization_config.group_size, + linear.in_features // self.group_size, linear.out_features, - dtype=scales_dtype, + dtype=self.scales_dtype, device=linear.weight.device, ) gemlite_linear.zeros = torch.empty( - linear.in_features // quantization_config.group_size, + linear.in_features // self.group_size, linear.out_features, - dtype=zeros_dtype, + dtype=self.zeros_dtype, device=linear.weight.device, ) - gemlite_linear.metadata = torch.empty( - len(gemlite_linear.get_meta_args()), dtype=torch.int32, device=linear.weight.device - ) - gemlite_linear.orig_shape = torch.empty(2, dtype=torch.int32, device=linear.weight.device) - gemlite_linear.meta_scale = torch.empty((), dtype=torch.float32, device=linear.weight.device) - - def replace(module: "nn.Module", prefix: str = "") -> int: - replaced = 0 - for name, child in module.named_children(): - child_name = f"{prefix}.{name}" if prefix else name - should_replace = ( - isinstance(child, nn.Linear) - and (quantized_fqns is None or child_name in quantized_fqns) - and not _is_in_skip_modules(child_name, modules_to_not_convert) - ) - if should_replace: - gemlite_linear = GemLiteLinearTriton( - W_nbits=quantization_config.bits, - group_size=quantization_config.group_size, - in_features=child.in_features, - out_features=child.out_features, - input_dtype=input_dtype, - output_dtype=output_dtype, - ).to(child.weight.device) - # Match the checkpoint's serialized tensor shapes and dtypes so Accelerate sizes the device map correctly. - initialize_serialized_tensors(gemlite_linear, child) - if child.bias is not None: - gemlite_linear.bias = nn.Parameter(torch.empty_like(child.bias), requires_grad=False) - gemlite_linear._gemlite_loaded_param_names = set() - setattr(module, name, gemlite_linear) - replaced += 1 - else: - replaced += replace(child, child_name) - return replaced - - return replace(model) + if linear.bias is not None: + gemlite_linear.bias = nn.Parameter(torch.empty_like(linear.bias), requires_grad=False) + gemlite_linear._gemlite_loaded_param_names = set() + return gemlite_linear class GemLiteQuantizer(DiffusersQuantizer): @@ -255,7 +237,14 @@ def _process_model_before_weight_loading( self.modules_to_not_convert.extend(keep_in_fp32_modules) self.modules_to_not_convert = [module for module in self.modules_to_not_convert if module is not None] - _replace_with_gemlite_linear(model, self.modules_to_not_convert, self.quantization_config) + from gemlite.helper import patch_model + + processor = _GemLiteDiffusersProcessor(self.quantization_config, self.modules_to_not_convert) + patch_model( + model, + device=next(model.parameters()).device, + processor=processor, + ) model.config.quantization_config = self.quantization_config def _process_model_after_weight_loading(self, model: "ModelMixin", **kwargs): diff --git a/tests/quantization/gemlite/test_gemlite.py b/tests/quantization/gemlite/test_gemlite.py index 601ab87b06f8..ed371f61a220 100644 --- a/tests/quantization/gemlite/test_gemlite.py +++ b/tests/quantization/gemlite/test_gemlite.py @@ -12,9 +12,9 @@ from diffusers.models.modeling_utils import ModelMixin from diffusers.quantizers.auto import DiffusersAutoQuantizer from diffusers.quantizers.gemlite.gemlite_quantizer import ( + _GemLiteDiffusersProcessor, _is_in_skip_modules, _normalize_torch_device, - _replace_with_gemlite_linear, ) from diffusers.quantizers.quantization_config import GemLiteConfig, QuantizationMethod from diffusers.utils import get_module_from_name, is_gemlite_available @@ -321,7 +321,9 @@ def import_with_broken_gemlite_core(name, *args, **kwargs): @require_gemlite_version_greater_or_equal("0.6.0") @require_torch_gpu class GemLiteQuantizerTest(unittest.TestCase): - def test_replace_with_gemlite_linear(self): + def test_processor_replaces_gemlite_linears(self): + from gemlite.helper import patch_model + from diffusers import FluxTransformer2DModel skip_patterns = [ @@ -348,9 +350,7 @@ def test_replace_with_gemlite_linear(self): original_linear_names = {name for name, module in model.named_modules() if isinstance(module, nn.Linear)} skipped_linear_names = {name for name in original_linear_names if _is_in_skip_modules(name, skip_patterns)} - replaced = _replace_with_gemlite_linear( - model, - skip_patterns, + processor = _GemLiteDiffusersProcessor( GemLiteConfig( bits=2, group_size=16, @@ -360,21 +360,28 @@ def test_replace_with_gemlite_linear(self): scales_dtype="fp32", zeros_dtype="fp32", ), + skip_patterns, + ) + patch_model( + model, + device=torch_device, + processor=processor, + skip_modules=[], ) - self.assertEqual(replaced, 20) + self.assertEqual(sum(isinstance(module, GemLiteLinearTriton) for module in model.modules()), 20) for name, module in model.named_modules(): if name in skipped_linear_names: self.assertIs(type(module), nn.Linear) elif name in original_linear_names: self.assertIsInstance(module, GemLiteLinearTriton) - def test_replace_preserves_bias_state_key(self): + def test_processor_preserves_bias_state_key(self): + from gemlite.helper import patch_model + model = nn.Sequential(nn.Linear(32, 32, bias=True), nn.Linear(32, 32, bias=False)).to(torch_device) - replaced = _replace_with_gemlite_linear( - model, - [], + processor = _GemLiteDiffusersProcessor( GemLiteConfig( bits=2, group_size=16, @@ -384,9 +391,16 @@ def test_replace_preserves_bias_state_key(self): scales_dtype="fp32", zeros_dtype="fp32", ), + [], + ) + patch_model( + model, + device=torch_device, + processor=processor, + skip_modules=[], ) - self.assertEqual(replaced, 2) + self.assertEqual(sum(isinstance(module, GemLiteLinearTriton) for module in model.modules()), 2) state_dict = model.state_dict() assert _GEMLITE_SERIALIZED_STATE_NAMES.issubset( {name.removeprefix("0.") for name in state_dict if name.startswith("0.")} From 56d5d02e81621917c37d2179ab659c215d635f1f Mon Sep 17 00:00:00 2001 From: Gabriel Date: Mon, 27 Jul 2026 23:34:18 -0500 Subject: [PATCH 03/10] Remove redundant GemLite core import check --- .../quantizers/gemlite/gemlite_quantizer.py | 4 ---- tests/quantization/gemlite/test_gemlite.py | 14 -------------- 2 files changed, 18 deletions(-) diff --git a/src/diffusers/quantizers/gemlite/gemlite_quantizer.py b/src/diffusers/quantizers/gemlite/gemlite_quantizer.py index db5c27ee9318..46006a9c8ada 100644 --- a/src/diffusers/quantizers/gemlite/gemlite_quantizer.py +++ b/src/diffusers/quantizers/gemlite/gemlite_quantizer.py @@ -184,10 +184,6 @@ def validate_environment(self, *args, **kwargs): f"Using GemLite quantization requires gemlite>={_GEMLITE_MIN_VERSION}. " "Please upgrade with `pip install -U gemlite`." ) - try: - __import__("gemlite.core") - except Exception as error: - raise ImportError("GemLite is installed but its core linear module could not be imported.") from error def check_if_quantized_param( self, diff --git a/tests/quantization/gemlite/test_gemlite.py b/tests/quantization/gemlite/test_gemlite.py index ed371f61a220..be843a469bf8 100644 --- a/tests/quantization/gemlite/test_gemlite.py +++ b/tests/quantization/gemlite/test_gemlite.py @@ -302,20 +302,6 @@ def test_validate_environment_rejects_old_gemlite_version(self): with self.assertRaisesRegex(ImportError, "requires gemlite>=0.6.0"): quantizer.validate_environment() - def test_validate_environment_rejects_broken_gemlite_core_import(self): - quantizer = DiffusersAutoQuantizer.from_config(_get_gemlite_config(), pre_quantized=True) - original_import = __import__ - - def import_with_broken_gemlite_core(name, *args, **kwargs): - if name == "gemlite.core": - raise RuntimeError("broken gemlite core") - return original_import(name, *args, **kwargs) - - with self._mock_module_availability(gemlite=True, cuda=True, old_gemlite=False): - with mock.patch("builtins.__import__", side_effect=import_with_broken_gemlite_core): - with self.assertRaisesRegex(ImportError, "core linear module could not be imported"): - quantizer.validate_environment() - @require_gemlite @require_gemlite_version_greater_or_equal("0.6.0") From b6dedd48c22444e838eab6a23454bfa2987665ee Mon Sep 17 00:00:00 2001 From: Gabriel Date: Mon, 27 Jul 2026 23:38:29 -0500 Subject: [PATCH 04/10] Remove GemLite CUDA device coercion --- src/diffusers/quantizers/gemlite/gemlite_quantizer.py | 7 ------- tests/quantization/gemlite/test_gemlite.py | 11 +---------- 2 files changed, 1 insertion(+), 17 deletions(-) diff --git a/src/diffusers/quantizers/gemlite/gemlite_quantizer.py b/src/diffusers/quantizers/gemlite/gemlite_quantizer.py index 46006a9c8ada..7ef2777c7b58 100644 --- a/src/diffusers/quantizers/gemlite/gemlite_quantizer.py +++ b/src/diffusers/quantizers/gemlite/gemlite_quantizer.py @@ -27,12 +27,6 @@ def _is_in_skip_modules(name: str, modules_to_not_convert: list[str]) -> bool: return any((key + "." in name) or (key == name) for key in modules_to_not_convert) -def _normalize_torch_device(device: Any) -> "torch.device": - if isinstance(device, int): - return torch.device(f"cuda:{device}") - return torch.device(device) - - class _GemLiteDiffusersProcessor: """Implement GemLite's `patch_model` protocol for Diffusers' pre-quantized loading path.""" @@ -209,7 +203,6 @@ def create_quantized_param( **kwargs, ): module, tensor_name = get_module_from_name(model, param_name) - target_device = _normalize_torch_device(target_device) if tensor_name not in GEMLITE_STATE_NAMES: raise ValueError(f"`{param_name}` is not a GemLite serialized tensor.") diff --git a/tests/quantization/gemlite/test_gemlite.py b/tests/quantization/gemlite/test_gemlite.py index be843a469bf8..a95f11d98368 100644 --- a/tests/quantization/gemlite/test_gemlite.py +++ b/tests/quantization/gemlite/test_gemlite.py @@ -11,11 +11,7 @@ from diffusers.configuration_utils import ConfigMixin, register_to_config from diffusers.models.modeling_utils import ModelMixin from diffusers.quantizers.auto import DiffusersAutoQuantizer -from diffusers.quantizers.gemlite.gemlite_quantizer import ( - _GemLiteDiffusersProcessor, - _is_in_skip_modules, - _normalize_torch_device, -) +from diffusers.quantizers.gemlite.gemlite_quantizer import _GemLiteDiffusersProcessor, _is_in_skip_modules from diffusers.quantizers.quantization_config import GemLiteConfig, QuantizationMethod from diffusers.utils import get_module_from_name, is_gemlite_available @@ -258,11 +254,6 @@ def test_is_in_skip_modules_matches_exact_names_and_children(self): self.assertFalse(_is_in_skip_modules("proj_output", modules_to_not_convert)) self.assertFalse(_is_in_skip_modules("blocks.01.attn.to_q", modules_to_not_convert)) - def test_normalize_torch_device_handles_int_device_map_values(self): - self.assertEqual(_normalize_torch_device(0), torch.device("cuda:0")) - self.assertEqual(_normalize_torch_device("cpu"), torch.device("cpu")) - self.assertEqual(_normalize_torch_device(torch.device("meta")), torch.device("meta")) - @require_gemlite @require_torch_gpu From 9b118c527ff5452d3dd588a2328f2c0572c679dd Mon Sep 17 00:00:00 2001 From: Gabriel Date: Mon, 27 Jul 2026 23:51:54 -0500 Subject: [PATCH 05/10] Finalize GemLite layers during loading --- src/diffusers/models/modeling_utils.py | 9 ++-- .../quantizers/gemlite/gemlite_quantizer.py | 31 ++++++----- tests/quantization/gemlite/test_gemlite.py | 51 +++++++++++++++++++ 3 files changed, 73 insertions(+), 18 deletions(-) diff --git a/src/diffusers/models/modeling_utils.py b/src/diffusers/models/modeling_utils.py index 6e292c4ae747..b8d4c048c3cd 100644 --- a/src/diffusers/models/modeling_utils.py +++ b/src/diffusers/models/modeling_utils.py @@ -1411,11 +1411,6 @@ 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 = { @@ -1425,6 +1420,10 @@ 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) diff --git a/src/diffusers/quantizers/gemlite/gemlite_quantizer.py b/src/diffusers/quantizers/gemlite/gemlite_quantizer.py index 7ef2777c7b58..d86879981970 100644 --- a/src/diffusers/quantizers/gemlite/gemlite_quantizer.py +++ b/src/diffusers/quantizers/gemlite/gemlite_quantizer.py @@ -143,6 +143,14 @@ def update_torch_dtype(self, torch_dtype: "torch.dtype | None") -> "torch.dtype" ) return torch_dtype + def maybe_update_loaded_keys(self, loaded_keys: list[str], checkpoint_files: list[str]) -> list[str]: + self._gemlite_module_state_names = {} + for key in loaded_keys: + module_name, _, tensor_name = key.rpartition(".") + if tensor_name in GEMLITE_STATE_NAMES: + self._gemlite_module_state_names.setdefault(module_name, set()).add(tensor_name) + return loaded_keys + def get_special_dtypes_update(self, model, torch_dtype: "torch.dtype") -> dict[str, "torch.dtype"]: special_dtypes = super().get_special_dtypes_update(model, torch_dtype) @@ -212,6 +220,11 @@ def create_quantized_param( setattr(module, tensor_name, torch.nn.Parameter(value, requires_grad=False)) module._gemlite_loaded_param_names.add(tensor_name) + module_name, _, _ = param_name.rpartition(".") + expected_state_names = self._gemlite_module_state_names[module_name] + if module._gemlite_loaded_param_names == expected_state_names: + module.load_state_dict({name: getattr(module, name) for name in expected_state_names}) + def _process_model_before_weight_loading( self, model: "ModelMixin", @@ -223,6 +236,10 @@ def _process_model_before_weight_loading( Replace `nn.Linear` modules with GemLite modules before the checkpoint tensors are loaded. The serialized layout in the config creates correctly shaped GemLite tensors. Modules in `keep_in_fp32_modules` are excluded. """ + state_dict = kwargs.get("state_dict") + if state_dict is not None: + self.maybe_update_loaded_keys(list(state_dict), []) + self.modules_to_not_convert.extend(keep_in_fp32_modules) self.modules_to_not_convert = [module for module in self.modules_to_not_convert if module is not None] @@ -239,19 +256,7 @@ def _process_model_before_weight_loading( def _process_model_after_weight_loading(self, model: "ModelMixin", **kwargs): from gemlite.core import GemLiteLinearTriton - has_gemlite_linear = False - for module in model.modules(): - if not isinstance(module, GemLiteLinearTriton): - continue - has_gemlite_linear = True - state_dict = { - name: getattr(module, name) - for name in GEMLITE_STATE_NAMES - if name in module._gemlite_loaded_param_names - } - module.load_state_dict(state_dict) - - if not has_gemlite_linear: + if not any(isinstance(module, GemLiteLinearTriton) for module in model.modules()): logger.warning("No linear modules are using GemLite.") return model diff --git a/tests/quantization/gemlite/test_gemlite.py b/tests/quantization/gemlite/test_gemlite.py index a95f11d98368..fa27c600f4dd 100644 --- a/tests/quantization/gemlite/test_gemlite.py +++ b/tests/quantization/gemlite/test_gemlite.py @@ -1,5 +1,6 @@ import contextlib import gc +import os import tempfile import unittest from unittest import mock @@ -124,6 +125,15 @@ def _save_packed_gemlite_transformer(transformer, save_directory): return quantized_fqns +class GemLiteTestModel(ModelMixin, ConfigMixin): + _no_split_modules = [] + + @register_to_config + def __init__(self): + super().__init__() + self.proj = nn.Linear(256, 64, bias=False) + + class GemLiteConfigTest(unittest.TestCase): def test_config_defaults(self): config = _get_gemlite_config() @@ -529,6 +539,7 @@ def test_create_quantized_param_loads_pre_quantized_state(self): model = nn.Sequential(GemLiteLinearTriton()).to(torch_device) model[0]._gemlite_loaded_param_names = set() quantizer = DiffusersAutoQuantizer.from_config(_get_gemlite_config(), pre_quantized=True) + quantizer.maybe_update_loaded_keys([f"0.{name}" for name in gemlite_state_dict], []) for name, value in gemlite_state_dict.items(): with self.subTest(name=name): @@ -543,6 +554,7 @@ def test_create_quantized_param_loads_pre_quantized_state(self): self.assertTrue(torch.equal(loaded_value, value)) self.assertEqual(model[0]._gemlite_loaded_param_names, set(gemlite_state_dict)) + self.assertEqual(model[0].W_group_mode, gemlite_state_dict["metadata"][10].item()) def test_create_quantized_param_preserves_serialized_dtypes(self): quantizer = DiffusersAutoQuantizer.from_config(_get_gemlite_config(), pre_quantized=True) @@ -552,6 +564,7 @@ def test_create_quantized_param_preserves_serialized_dtypes(self): "0.W_q": torch.arange(4, dtype=torch.float16, device=torch_device).reshape(2, 2).to(torch.float8_e5m2), "0.scales": torch.tensor([[0.12345679], [0.9876543]], dtype=torch.float32, device=torch_device), } + quantizer.maybe_update_loaded_keys([f"0.{name}" for name in _GEMLITE_SERIALIZED_STATE_NAMES], []) for param_name, original_value in serialized_state.items(): with self.subTest(param_name=param_name): @@ -567,6 +580,44 @@ def test_create_quantized_param_preserves_serialized_dtypes(self): self.assertEqual(loaded_value.dtype, original_value.dtype) self.assertTrue(torch.equal(loaded_value, original_value)) + @require_accelerate + def test_finalizes_gemlite_modules_before_dispatch(self): + quantization_config = _get_gemlite_config( + bits=4, + group_size=64, + packing_bitwidth=32, + quantized_fqns=["proj"], + ) + packed_state = _create_packed_gemlite_state_dict( + in_features=256, + out_features=64, + w_nbits=4, + group_size=64, + packing_bitwidth=32, + device=torch_device, + ) + self.assertNotEqual(packed_state["metadata"][10].item(), 0) + + model = GemLiteTestModel().to(torch_device) + model.proj = GemLiteLinearTriton() + model.proj.load_state_dict(packed_state.copy()) + model.register_to_config(quantization_config=quantization_config) + + with tempfile.TemporaryDirectory() as model_dir: + model.save_pretrained(model_dir, safe_serialization=False, max_shard_size="1KB") + self.assertTrue(any(name.endswith(".index.json") for name in os.listdir(model_dir))) + + def assert_finalized_before_dispatch(loaded_model, **kwargs): + self.assertEqual(loaded_model.proj.W_group_mode, packed_state["metadata"][10].item()) + self.assertTrue(torch.equal(loaded_model.state_dict()["proj.metadata"], packed_state["metadata"])) + + with mock.patch( + "diffusers.models.modeling_utils.dispatch_model", side_effect=assert_finalized_before_dispatch + ) as dispatch_model: + GemLiteTestModel.from_pretrained(model_dir, dtype=torch.float16, device_map={"": torch_device}) + + dispatch_model.assert_called_once() + def test_get_state_dict_and_metadata_sets_data_contiguous_for_serialization(self): quantizer = DiffusersAutoQuantizer.from_config(_get_gemlite_config(), pre_quantized=True) w_q = torch.arange(12, dtype=torch.uint8, device=torch_device).reshape(3, 4).t() From db04c49283543e76d008348f14ccc05af14aefb7 Mon Sep 17 00:00:00 2001 From: Gabriel Date: Mon, 27 Jul 2026 23:59:54 -0500 Subject: [PATCH 06/10] Fix GemLite documentation link --- docs/source/en/quantization/gemlite.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/source/en/quantization/gemlite.md b/docs/source/en/quantization/gemlite.md index 78ed6cfe5334..3413094be903 100644 --- a/docs/source/en/quantization/gemlite.md +++ b/docs/source/en/quantization/gemlite.md @@ -13,7 +13,7 @@ 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 +[GemLite](https://github.com/dropbox/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. From fd7b31b82fbcff71a9708de34706d80f7039cbc2 Mon Sep 17 00:00:00 2001 From: Gabriel Date: Tue, 28 Jul 2026 10:48:27 -0500 Subject: [PATCH 07/10] Validate GemLite compute dtypes --- src/diffusers/quantizers/quantization_config.py | 7 ++++++- tests/quantization/gemlite/test_gemlite.py | 10 ++++++++++ 2 files changed, 16 insertions(+), 1 deletion(-) diff --git a/src/diffusers/quantizers/quantization_config.py b/src/diffusers/quantizers/quantization_config.py index 521caecaa52e..6618ee875a3d 100644 --- a/src/diffusers/quantizers/quantization_config.py +++ b/src/diffusers/quantizers/quantization_config.py @@ -514,7 +514,12 @@ def __init__( if self.compute_dtype is None: self.compute_dtype = torch.float16 elif isinstance(self.compute_dtype, str): - self.compute_dtype = getattr(torch, self.compute_dtype) + compute_dtype = getattr(torch, self.compute_dtype, None) + if not isinstance(compute_dtype, torch.dtype): + raise ValueError(f"Unsupported GemLite compute dtype: {self.compute_dtype!r}.") + self.compute_dtype = compute_dtype + elif not isinstance(self.compute_dtype, torch.dtype): + raise ValueError("GemLite compute_dtype must be a torch.dtype.") self.post_init() diff --git a/tests/quantization/gemlite/test_gemlite.py b/tests/quantization/gemlite/test_gemlite.py index fa27c600f4dd..ae573e4ef86d 100644 --- a/tests/quantization/gemlite/test_gemlite.py +++ b/tests/quantization/gemlite/test_gemlite.py @@ -195,6 +195,16 @@ def test_config_from_dict(self): self.assertEqual(config.compute_dtype, torch.bfloat16) self.assertEqual(config.modules_to_not_convert, ["proj_out"]) + def test_config_rejects_invalid_compute_dtype(self): + config_dict = _get_gemlite_config().to_dict() + + for compute_dtype in ("foo", "nn"): + with self.subTest(compute_dtype=compute_dtype): + config_dict["compute_dtype"] = compute_dtype + + with self.assertRaisesRegex(ValueError, "Unsupported GemLite compute dtype"): + DiffusersAutoQuantizer.from_dict(config_dict) + def test_config_round_trip(self): config = GemLiteConfig( compute_dtype=torch.bfloat16, From abe8d1dc57b2064e12becf099b3cc9d93a855590 Mon Sep 17 00:00:00 2001 From: Gabriel Date: Tue, 28 Jul 2026 11:00:45 -0500 Subject: [PATCH 08/10] Clarify GemLite version requirement --- src/diffusers/utils/testing_utils.py | 2 +- tests/testing_utils.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/src/diffusers/utils/testing_utils.py b/src/diffusers/utils/testing_utils.py index 01a8c52c8433..c6eac86cefe0 100644 --- a/src/diffusers/utils/testing_utils.py +++ b/src/diffusers/utils/testing_utils.py @@ -639,7 +639,7 @@ def decorator(test_case): version.parse(importlib.metadata.version("gemlite")).base_version ) >= version.parse(gemlite_version) return unittest.skipUnless( - correct_gemlite_version, f"Test requires gemlite with the version greater than {gemlite_version}." + correct_gemlite_version, f"Test requires gemlite with the version greater than or equal to {gemlite_version}." )(test_case) return decorator diff --git a/tests/testing_utils.py b/tests/testing_utils.py index 5b79e1907832..691cec571f49 100644 --- a/tests/testing_utils.py +++ b/tests/testing_utils.py @@ -852,7 +852,7 @@ def decorator(test_case): ) >= version.parse(gemlite_version) return pytest.mark.skipif( not correct_gemlite_version, - reason=f"Test requires gemlite with the version greater than {gemlite_version}.", + reason=f"Test requires gemlite with the version greater than or equal to {gemlite_version}.", )(test_case) return decorator From f3d69de8ba3de213d6491066d83dec8ebf889a3d Mon Sep 17 00:00:00 2001 From: Gabriel Date: Tue, 28 Jul 2026 11:20:12 -0500 Subject: [PATCH 09/10] Test GemLite serialization round trip --- tests/quantization/gemlite/test_gemlite.py | 107 +++++++++++++++------ 1 file changed, 76 insertions(+), 31 deletions(-) diff --git a/tests/quantization/gemlite/test_gemlite.py b/tests/quantization/gemlite/test_gemlite.py index ae573e4ef86d..412213c5ead7 100644 --- a/tests/quantization/gemlite/test_gemlite.py +++ b/tests/quantization/gemlite/test_gemlite.py @@ -654,19 +654,25 @@ def test_create_quantized_param_rejects_non_gemlite_serialized_name(self): quantizer.create_quantized_param(model, torch.ones(32, 64, device=torch_device), "0.weight", torch_device) +class GemLiteIntegrationTestCase(unittest.TestCase): + torch_dtype = torch.float16 + maximum_quantized_memory_fraction = 0.8 + + def _clear_memory(self): + gc.collect() + backend_empty_cache(torch_device) + + def tearDown(self): + self._clear_memory() + + @nightly @require_gemlite @require_gemlite_version_greater_or_equal("0.6.0") @require_torch_gpu @require_accelerate -class GemLiteKrea2TransformerIntegrationTests(unittest.TestCase): +class GemLiteKrea2TransformerIntegrationTests(GemLiteIntegrationTestCase): model_id = "hf-internal-testing/tiny-krea2-modular-pipe" - torch_dtype = torch.float16 - maximum_quantized_memory_fraction = 0.8 - - def tearDown(self): - gc.collect() - backend_empty_cache(torch_device) def test_inference_matches_unquantized_transformer(self): from diffusers import Krea2Transformer2DModel @@ -692,24 +698,42 @@ def test_inference_matches_unquantized_transformer(self): self.assertTrue(quantized_fqns) del reference_transformer - gc.collect() - backend_empty_cache(torch_device) + self._clear_memory() - transformer = Krea2Transformer2DModel.from_pretrained( + first_load_transformer = Krea2Transformer2DModel.from_pretrained( model_dir, dtype=self.torch_dtype, device_map={"": torch_device}, ) - self.assertTrue(any(isinstance(module, GemLiteLinearTriton) for module in transformer.modules())) self.assertLessEqual( - transformer.get_memory_footprint(), + first_load_transformer.get_memory_footprint(), unquantized_transformer_memory * self.maximum_quantized_memory_fraction, ) with torch.inference_mode(): - gemlite_output = transformer(**inputs).sample + first_load_output = first_load_transformer(**inputs).sample + + with tempfile.TemporaryDirectory() as model_dir: + first_load_transformer.save_pretrained(model_dir, safe_serialization=True) + del first_load_transformer + self._clear_memory() - torch.testing.assert_close(gemlite_output, reference_output, rtol=5e-2, atol=5e-2) + second_load_transformer = Krea2Transformer2DModel.from_pretrained( + model_dir, + dtype=self.torch_dtype, + device_map={"": torch_device}, + ) + + self.assertLessEqual( + second_load_transformer.get_memory_footprint(), + unquantized_transformer_memory * self.maximum_quantized_memory_fraction, + ) + with torch.inference_mode(): + second_load_output = second_load_transformer(**inputs).sample + + torch.testing.assert_close(first_load_output, reference_output, rtol=5e-2, atol=5e-2) + torch.testing.assert_close(second_load_output, reference_output, rtol=5e-2, atol=5e-2) + torch.testing.assert_close(second_load_output, first_load_output, rtol=5e-2, atol=5e-2) @nightly @@ -717,14 +741,8 @@ def test_inference_matches_unquantized_transformer(self): @require_gemlite_version_greater_or_equal("0.6.0") @require_torch_gpu @require_accelerate -class GemLiteFluxPipelineIntegrationTests(unittest.TestCase): +class GemLiteFluxPipelineIntegrationTests(GemLiteIntegrationTestCase): model_id = "hf-internal-testing/tiny-flux-pipe" - torch_dtype = torch.float16 - maximum_quantized_memory_fraction = 0.8 - - def tearDown(self): - gc.collect() - backend_empty_cache(torch_device) def test_inference_matches_unquantized_pipeline(self): from diffusers import FluxPipeline, FluxTransformer2DModel @@ -750,30 +768,57 @@ def test_inference_matches_unquantized_pipeline(self): self.assertTrue(quantized_fqns) del reference_pipe - gc.collect() - backend_empty_cache(torch_device) + self._clear_memory() + + first_load_transformer = FluxTransformer2DModel.from_pretrained( + model_dir, + dtype=self.torch_dtype, + device_map={"": torch_device}, + ) + + self.assertLessEqual( + first_load_transformer.get_memory_footprint(), + unquantized_transformer_memory * self.maximum_quantized_memory_fraction, + ) + first_load_pipe = FluxPipeline.from_pretrained( + self.model_id, + transformer=first_load_transformer, + dtype=self.torch_dtype, + ).to(torch_device) + first_load_pipe.set_progress_bar_config(disable=True) + with torch.inference_mode(): + first_load_output = first_load_pipe( + **pipe_inputs, + generator=torch.Generator(device=torch_device).manual_seed(0), + ).images + + with tempfile.TemporaryDirectory() as model_dir: + first_load_transformer.save_pretrained(model_dir, safe_serialization=True) + del first_load_pipe, first_load_transformer + self._clear_memory() - transformer = FluxTransformer2DModel.from_pretrained( + second_load_transformer = FluxTransformer2DModel.from_pretrained( model_dir, dtype=self.torch_dtype, device_map={"": torch_device}, ) - self.assertTrue(any(isinstance(module, GemLiteLinearTriton) for module in transformer.modules())) self.assertLessEqual( - transformer.get_memory_footprint(), + second_load_transformer.get_memory_footprint(), unquantized_transformer_memory * self.maximum_quantized_memory_fraction, ) - gemlite_pipe = FluxPipeline.from_pretrained( + second_load_pipe = FluxPipeline.from_pretrained( self.model_id, - transformer=transformer, + transformer=second_load_transformer, dtype=self.torch_dtype, ).to(torch_device) - gemlite_pipe.set_progress_bar_config(disable=True) + second_load_pipe.set_progress_bar_config(disable=True) with torch.inference_mode(): - gemlite_output = gemlite_pipe( + second_load_output = second_load_pipe( **pipe_inputs, generator=torch.Generator(device=torch_device).manual_seed(0), ).images - np.testing.assert_allclose(gemlite_output, reference_output, rtol=5e-2, atol=5e-2) + np.testing.assert_allclose(first_load_output, reference_output, rtol=5e-2, atol=5e-2) + np.testing.assert_allclose(second_load_output, reference_output, rtol=5e-2, atol=5e-2) + np.testing.assert_allclose(second_load_output, first_load_output, rtol=5e-2, atol=5e-2) From 5021969ad240657890ab93ff9eebec042f493fe8 Mon Sep 17 00:00:00 2001 From: Gabriel Date: Tue, 28 Jul 2026 11:43:19 -0500 Subject: [PATCH 10/10] Fix GemLite quantizer tests --- .../quantizers/gemlite/gemlite_quantizer.py | 9 +-- tests/quantization/gemlite/test_gemlite.py | 68 +------------------ 2 files changed, 6 insertions(+), 71 deletions(-) diff --git a/src/diffusers/quantizers/gemlite/gemlite_quantizer.py b/src/diffusers/quantizers/gemlite/gemlite_quantizer.py index d86879981970..90c4d16d8629 100644 --- a/src/diffusers/quantizers/gemlite/gemlite_quantizer.py +++ b/src/diffusers/quantizers/gemlite/gemlite_quantizer.py @@ -113,8 +113,9 @@ class GemLiteQuantizer(DiffusersQuantizer): This quantizer only loads pre-quantized checkpoints. It replaces `torch.nn.Linear` modules with `GemLiteLinearTriton` modules before weight loading, then restores the serialized GemLite state (`W_q`, `scales`, `zeros`, `metadata`, `orig_shape`, `meta_scale`, ...) through the low-memory loader. The quantization config - provides the serialized layout, so the replacement tensors have the checkpoint's exact shapes and dtypes before - Accelerate estimates module sizes for `device_map="auto"`. + provides the packed weight layout, so replacement `W_q`, `scales`, and `zeros` tensors have the checkpoint's + exact shapes and dtypes before Accelerate estimates module sizes for `device_map="auto"`. Serialization metadata + is restored during weight loading and is not used for device-map sizing. Modules listed in `modules_to_not_convert` are skipped and left in their original dtype. """ @@ -233,8 +234,8 @@ def _process_model_before_weight_loading( **kwargs, ): """ - Replace `nn.Linear` modules with GemLite modules before the checkpoint tensors are loaded. The serialized - layout in the config creates correctly shaped GemLite tensors. Modules in `keep_in_fp32_modules` are excluded. + Replace `nn.Linear` modules with GemLite modules before the checkpoint tensors are loaded. The packed layout + in the config creates correctly shaped compute tensors. Modules in `keep_in_fp32_modules` are excluded. """ state_dict = kwargs.get("state_dict") if state_dict is not None: diff --git a/tests/quantization/gemlite/test_gemlite.py b/tests/quantization/gemlite/test_gemlite.py index 412213c5ead7..8b7840bc8329 100644 --- a/tests/quantization/gemlite/test_gemlite.py +++ b/tests/quantization/gemlite/test_gemlite.py @@ -405,72 +405,6 @@ def test_processor_preserves_bias_state_key(self): self.assertIn("0.bias", state_dict) self.assertNotIn("1.bias", state_dict) - @require_accelerate - def test_prequantized_auto_device_map_uses_packed_state_size(self): - from accelerate.utils import compute_module_sizes - - from diffusers.models.model_loading_utils import _determine_device_map - - class GemLiteDeviceMapTestModel(ModelMixin, ConfigMixin): - _no_split_modules = [] - - @register_to_config - def __init__(self): - super().__init__() - self.proj = nn.Linear(128, 32, bias=False) - - packed_state = _create_packed_gemlite_state_dict( - in_features=128, - out_features=32, - w_nbits=2, - group_size=128, - packing_bitwidth=8, - scales_dtype=torch.float32, - zeros_dtype=torch.float32, - device=torch_device, - ) - packed_size = sum(value.numel() * value.element_size() for value in packed_state.values()) - - quantizer = DiffusersAutoQuantizer.from_config( - GemLiteConfig( - format="gemlite-int2-ternary-g128", - bits=2, - group_size=128, - packing_bitwidth=8, - solver="ternary", - input_dtype="fp16", - output_dtype="fp16", - scales_dtype="fp32", - zeros_dtype="fp32", - quantized_fqns=["proj"], - ), - pre_quantized=True, - ) - gemlite_model = GemLiteDeviceMapTestModel().to("meta") - quantizer.preprocess_model(gemlite_model, device_map="auto", keep_in_fp32_modules=[]) - placeholder_state_dict = gemlite_model.state_dict() - for name, value in packed_state.items(): - placeholder = placeholder_state_dict[f"proj.{name}"] - self.assertEqual(placeholder.shape, value.shape) - self.assertEqual(placeholder.dtype, value.dtype) - - module_sizes = compute_module_sizes( - gemlite_model, - dtype=torch.float16, - special_dtypes=quantizer.get_special_dtypes_update(gemlite_model, torch.float16), - ) - self.assertEqual(module_sizes["proj"], packed_size) - with self.assertWarnsRegex(UserWarning, "Current model requires .* bytes of buffer for offloaded layers"): - gemlite_device_map = _determine_device_map( - gemlite_model, - "auto", - {0: packed_size // 2, "cpu": packed_size * 2}, - torch.float16, - hf_quantizer=quantizer, - ) - - self.assertNotIn(0, gemlite_device_map.values()) - def test_process_model_before_weight_loading_replaces_pre_quantized_linears(self): from diffusers import FluxTransformer2DModel @@ -614,7 +548,7 @@ def test_finalizes_gemlite_modules_before_dispatch(self): model.register_to_config(quantization_config=quantization_config) with tempfile.TemporaryDirectory() as model_dir: - model.save_pretrained(model_dir, safe_serialization=False, max_shard_size="1KB") + model.save_pretrained(model_dir, safe_serialization=True, max_shard_size="1KB") self.assertTrue(any(name.endswith(".index.json") for name in os.listdir(model_dir))) def assert_finalized_before_dispatch(loaded_model, **kwargs):