Add GemLite quantization backend support for pre-quantized checkpoints#14286
Add GemLite quantization backend support for pre-quantized checkpoints#14286gabe-engineers wants to merge 1 commit into
Conversation
github-actions
Bot
added
documentation
| return torch.device(device) | ||
|
|
||
|
|
||
| def _replace_with_gemlite_linear( |
There was a problem hiding this comment.
Don't do this manually, the helper functions do all of that and manage some edge cases. Also using the helper functions makes it easier to upgrade without breaking things.
You can simply use: https://github.com/dropbox/gemlite/blob/master/gemlite/helper.py#L35
There was a problem hiding this comment.
That's a good callout, thanks! If we create a custom processor to be used with patch_model, then it would work. I'm thinking something like GemLiteDiffusersLoadingAdapter that would put the model in the state the diffuser loader expects before the actual weight loading.
The custom processor would still be needed because of a few reasons:
-
Looks like currently the patch_model function uses
processor.from_linear()for the regular non hqq path, which in the existing processors in the helper file, seem to be designed to quantize the weights rather than loading an already packed layer. That seems to be reserved forfrom_packed_weights; however, it needs the real packed tensors and therefore cannot run during Diffusers’ pre-load phase. -
During the diffusers loading, if the device_map="auto" option is passed, then we need to let accelerate figure out the best device mapping, and to do this, we need to allocate buffers in the correct shape and dtype (I think this could also be done within GemLiteLinearTriton itself, over here in case the feature shapes and dtype arguments are passed in the constructor). If we don't do this, is not a fatal issue necessarily, but it might result in sub-optimal auto device allocations and OOM errors in some scenarios.
There was a problem hiding this comment.
In the end I'm thinking about something like this:
class GemLiteDiffusersLoadingAdapter:
def __init__(self, quantization_config: "GemLiteConfig"):
# Store the quantization_config attributes using the GemLite dtypes
def from_linear(self, linear: "nn.Linear") -> "nn.Module":
# Instantiate GemLiteLinearTriton with the params from the config, then allocate the correctly sized buffers for accelerate, but don't load anything into any devices.
loading_adapter = GemLiteDiffusersLoadingAdapter(self.quantization_config)
patch_model(
model,
device=next(model.parameters()).device,
processor=loading_adapter,
skip_modules=self.module_to_not_convert,
)There was a problem hiding this comment.
seem to be designed to quantize the weights rather than loading an already packed layer yeah, but that's specific to patch_model, the other helper functions (individual patching functions) do support loading pre-quantized weights, they are used in vLLM to load various pre-quantized models. You can ask codex to take a look at: https://github.com/dropbox/gemlite/tree/master/gemlite/vllm
|
Hi there, I cannot open the resultant image, sadly. 
Could you update the link? Additionally, could you provide details on the following:
|
![sergereview[bot]](https://avatars.githubusercontent.com/in/3673521?s=60&v=4){kind=small}
There was a problem hiding this comment.
🤗 Serge says:
Overall this is a well-structured backend addition that mostly follows the existing quantizer conventions (auto.py registration, config in quantization_config.py, dummy objects, docs, nightly CI entry). A few things need attention before merge, chiefly the change to the shared loading path in modeling_utils.py and some deviations from how other backends implement helpers.
Correctness
- Reordering
postprocess_modelbeforedispatch_modelinmodels/modeling_utils.pyis a global behavior change for every quantizer, not just GemLite. bnb/torchao/quanto/modelopt/autoround all run their post-load hooks after dispatch today;autoround'spost_initin particular moves buffers to devices and freezes params, and ModelOpt's hook mutates module state. Please justify this ordering change with a test that covers the other backends, or find a GemLite-local way to finalize state (the finalization here is onlymodule.load_state_dict(...)on GemLite modules, which could equally happen without moving the shared hook). _normalize_torch_devicehardcodescuda:{device}for integer device-map values. Accelerate emits bare ints for any accelerator, so this breaks on XPU/other backends; no other quantizer in the tree does this coercion — they passparam_devicestraight to.to(...).validate_environmentdoes a bare__import__("gemlite.core")inside atry/except Exception. This is exactly the kind of defensive fallback the repo guide asks to avoid:is_gemlite_available()plus the version check already covers the realistic failure modes, and every subsequent code path importsgemlite.coredirectly and would surface a clear traceback anyway.
Style / conventions
_replace_with_gemlite_linearnests two closures (initialize_serialized_tensors,replace) inside a module-level function. Comparereplace_with_bnb_linear/_replace_with_quanto_layers, which are flat recursive functions. Flattening this would make the flow readable top-to-bottom.GemLiteConfig.__init__acceptscompute_dtypeas astrand silentlygetattr(torch, ...)it.GGUFQuantizationConfigandBitsAndBytesConfigonly taketorch.dtype; deserialization fromconfig.jsongoes throughto_dict/from_dict, so consider whether the string branch is needed or whether it should raise for unknown names instead of an opaqueAttributeError.GEMLITE_STATE_NAMESincludes"bias", socreate_quantized_paramwill route bias tensors through the quantized path for any GemLite module; worth a comment confirming that is intentional given bias is a plain fp param.
Docs
docs/source/en/quantization/gemlite.mdpoints tohttps://github.com/mobiusml/gemlitewhile the PR description sayshttps://github.com/dropbox/gemlite. Pick one canonical upstream URL.- The doc example uses a personal-namespace model id (
gabe-engineers/...). Other quantization docs reference official/organization checkpoints; consider whether this should live under a stable org.
Tests
- Test coverage of the config/quantizer units is good. However the two integration tests construct the packed checkpoint at runtime via
_save_packed_gemlite_transformer; the description's claim of a "persisted pre-quantized checkpoint → from_pretrained → forward/save/reload test" is only partially met — there is no save→reload round trip after loading (i.e.get_state_dict_and_metadatais only exercised in a unit test with a synthetic state dict, never end-to-end). require_gemlite_version_greater_or_equal's skip reason says "greater than" but the comparison is>=. Minor, but it mirrors the existing gguf/torchao helpers' wording bug — fine to leave, just noting.
serge v0.1.0 · model: claude-opus-5 · 31 LLM turns · 47 tool calls · 689.4s · 2106001 in / 50203 out tokens
|
|
||
| # Quantizers may need to finalize module state before dispatch reads `model.state_dict()`. | ||
| if hf_quantizer is not None: | ||
| hf_quantizer.postprocess_model(model) |
There was a problem hiding this comment.
Moving postprocess_model ahead of dispatch_model changes ordering for every quantizer, not just GemLite.
Today AutoRoundQuantizer._process_model_after_weight_loading calls post_init(model, self.used_backends), which the docstring in that file describes as "repacking weights into the kernel's expected memory layout, moving buffers to the correct device" — running that before dispatch_model may fight with the hooks accelerate installs. NVIDIAModelOptQuantizer similarly mutates per-module ModelOpt state.
Could you either (a) keep the shared hook where it was and do the GemLite finalization in a way that doesn't require reordering, or (b) add coverage demonstrating bnb/torchao/quanto/modelopt/autoround still load correctly with device_map under the new ordering? As written this is an unreviewed behavior change to the shared load path.
|
|
||
| def _normalize_torch_device(device: Any) -> "torch.device": | ||
| if isinstance(device, int): | ||
| return torch.device(f"cuda:{device}") |
There was a problem hiding this comment.
Hardcoding cuda: makes this incorrect on non-CUDA accelerators (XPU, etc.), where accelerate also hands back bare integer device-map values.
Note that no other quantizer normalizes target_device at all — bnb, gguf, torchao and modelopt all pass the value straight into .to(...) / set_module_tensor_to_device, which already handles ints. Unless you have a concrete failure this is guarding against, dropping _normalize_torch_device and using target_device directly would match the rest of the tree and avoid the CUDA assumption.
| f"Using GemLite quantization requires gemlite>={_GEMLITE_MIN_VERSION}. " | ||
| "Please upgrade with `pip install -U gemlite`." | ||
| ) | ||
| try: |
There was a problem hiding this comment.
This __import__ + broad except Exception block is defensive code of the kind the contributing guide asks us not to add: is_gemlite_available() above already confirms the distribution is installed, and every other method in this class (get_special_dtypes_update, check_if_quantized_param, _process_model_after_weight_loading) imports gemlite.core directly and would raise a far more informative traceback than "core linear module could not be imported".
| try: |
| 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: |
There was a problem hiding this comment.
Two nested closures inside a module-level helper makes this hard to follow. replace_with_bnb_linear and _replace_with_quanto_layers are both flat recursive functions with the tensor setup inlined at the replacement site.
Since initialize_serialized_tensors has exactly one caller, inlining it into the should_replace branch (and making replace a plain recursive module-level function) would match the repo's "prefer inlining single-caller helpers" guidance and let a reader follow the whole flow in one pass.
|
|
||
| # GemLite | ||
|
|
||
| [GemLite](https://github.com/mobiusml/gemlite) is a quantization backend for loading prequantized checkpoints in |
There was a problem hiding this comment.
This links to github.com/mobiusml/gemlite, but the PR description describes GemLite as github.com/dropbox/gemlite. One of the two is wrong (or one is a redirect) — please settle on the canonical upstream URL so the docs don't point users at a stale mirror.
| import torch | ||
| from diffusers import DiffusionPipeline | ||
|
|
||
| model_id = "gabe-engineers/bonsai-image-ternary-4B-gemlite-2bit-unpacked-encoder" |
There was a problem hiding this comment.
The only example checkpoint is under a personal namespace. Other quantization docs point at org-owned or hf-internal-testing repos, which are less likely to disappear or be renamed. Worth moving this under a stable org before merge.
| 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) |
There was a problem hiding this comment.
getattr(torch, self.compute_dtype) will raise a bare AttributeError: module 'torch' has no attribute 'foo' for a typo'd or malicious value in a checkpoint's config.json, and would happily resolve a non-dtype attribute (e.g. "nn").
Neither GGUFQuantizationConfig nor BitsAndBytesConfig accept string dtypes at all — to_dict/from_dict round-trips go through the same path here, so if you keep the string branch please validate the result is actually a torch.dtype and raise a concise error otherwise.
3 participants
What does this PR do?
This PR adds GemLite as a quantization backend in Diffusers. It's scoped to only loading already pre-quantized checkpoints.
GemLite provides Triton kernels for low-bit matrix operations and supports packed low-bit weight formats, including binary and ternary weights.
The integration enables Diffusers models with pre-quantized GemLite weights to be loaded and executed directly through the standard Diffusers quantization interface.
Motivation
Several low-bit image-generation models, including the binary and ternary variants of Bonsai Image, distribute their weights using GemLite’s packed representation.
Diffusers currently cannot load these checkpoints directly. Users must instead rely on custom model-loading code or maintain separate forks and patches.
With this backend, models using GemLite low-bit packing can be loaded and run through Diffusers without model-specific loading logic. For the Bonsai Image series of models this PR will not be sufficient but it would be a part 1 of 2. With it you will be able to load the transformer checkpoints which are packed in GemLite format, however the encoder is quantized and packed with HQQ, and cannot be loaded yet (a good follow up but wanted to keep the PR scoped).
Fixes partially #13925 missing working HQQ support afterwards
Last self-review report:
The blocking issue, turned out to be a stale metadata result from codex's search tool. The needed gemlite 0.6.0 version is live in PyPI
I took both suggestions since they were good ones, added 2 end to end tests one with the Krea2 transformer and another with mini-flux's pipeline. Then I added an example to the documentation, I tested the example with a L40 GPU and verified the image. One note on the example is we are using a bonsai model with an unpacked encoder, this is because the encoder is quantized and packed using HQQ, and the HQQ support seems to be broken in transformers (huggingface/transformers#45147) and not available yet natively in diffusers (again a good follow up perhaps, but I wanted to keep this PR scoped).
Testing
Ran the gemlite tests with
RUN_NIGHTLY=1 python -m pytest tests/quantization/gemlite/test_gemlite.py. Output here.Also manually ran inference with:
Resulted in this image
Before submitting
self-reviewskill on the diff?documentation guidelines, and
here are tips on formatting docstrings.
Who can review?
Primary: @sayakpaul
Secondary: @DN6
Optional: @mobicham