Add GemLite quantization backend support for pre-quantized checkpoints - #14286
Add GemLite quantization backend support for pre-quantized checkpoints#14286gabe-engineers wants to merge 10 commits 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
There was a problem hiding this comment.
You're right about that, I checked those too. I couldn't find a good fit for this PR, though, because from_packed_weights receives the full packed tensors for a layer (W_q, scales, zeros, etc.), and those aren't available during Diffusers' preprocessing phase. So I'll go ahead with the custom processor while reusing patch_model. I think that was a good suggestion on your end.
|
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
| 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.
There was a problem hiding this comment.
I don't currently have a suitable org namespace to move this checkpoint into, so for now it would need to stay under my personal namespace. If the maintainers would prefer it under an HF-owned/testing namespace, I'm happy to move or copy it there if there's an appropriate destination.
| 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.
There was a problem hiding this comment.
I added validation for it, good point though, thanks for point it out.
Hello! I updated the link, feel free to check again. I also ran some benchmarks, compared First inference (cold) means the triton autotuning ran, which takes a while. NVIDIA L40, no torch.compile:
NVIDIA L40, with torch.compile:
NVIDIA RTX Pro 6000, no torch.compile:
NVIDIA RTX Pro 6000, with torch.compile:
|
|
@gabe-engineers when you measure torch.compile time, you need to do 2-3 iterations for warm-up/Triton autotuning. |
@mobicham thank for the note.The benchmark reports the cold first inference separately, then runs two warm-up iterations for torch.compile/Triton autotuning before collecting five reported hot runs. The warm-up iterations are excluded from the reported hot median and mean timings. The script used is here: https://gist.github.com/gabe-engineers/ed105533c3ca645d28133e6ba5e749c8 |
|
Thanks! Will try to review tomorrow. |
|
For reference, I like to use the |
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