Add Nunchaku Lite single-file quantization#14100
Conversation
github-actions
Bot
added
size/L
sayakpaul
left a comment
sayakpaul
left a comment
There was a problem hiding this comment.
Thanks for getting started! Just did a first pass and left high-level reviews.
| @@ -0,0 +1,161 @@ | |||
| import json | |||
There was a problem hiding this comment.
For tests, WDYT of adding a mixin to https://github.com/huggingface/diffusers/blob/main/tests/models/testing_utils/quantization.py and then extending a popular model like Flux to use that mixin?
There was a problem hiding this comment.
Yes, let's do it that way
|
I just did some benchmark on RTX PRO 6000, here is the visual result between bf16 and nvfp4 checkpoint for ERNIE-Image-Turbo
By replacing Nunchaku Linear, we have reduced the latency of these linear operations by 2x with large shape
One note here, AWQ only benefit with adaLN layer, so other modules like time embedding or final linear can stay as bf16
|
|
I have just implemented the native loading feature, which now can load by
import torch
from diffusers import ErnieImagePipeline
pipe = ErnieImagePipeline.from_pretrained(
"rootonchair/ERNIE-Image-Turbo-nunchaku-lite-int4",
torch_dtype=torch.bfloat16,
).to("cuda")
image = pipe(
prompt="A modern red armchair in a quiet studio, soft window light, realistic product photography",
height=1024,
width=1024,
num_inference_steps=8,
guidance_scale=1.0,
use_pe=False,
).images[0]
image.save("ernie-image-turbo-nunchaku-lite-int4.png")Quantization config now change to: If we agree to use this schema, I will remove the old metadata/from_single_file approach |
sayakpaul
left a comment
sayakpaul
left a comment
There was a problem hiding this comment.
Looking good. I think we can remove all metadata related code?
sayakpaul
left a comment
sayakpaul
left a comment
There was a problem hiding this comment.
No rush but let us know once you would like another round of review.
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
github-actions
Bot
added
documentation
|
@sayakpaul I think this is ready for the next review |
![sergereview[bot]](https://avatars.githubusercontent.com/in/3673521?s=60&v=4){kind=small}
There was a problem hiding this comment.
🤗 Serge says:
This adds a new nunchaku_lite quantization backend (config, quantizer, runtime linear layers backed by the HF kernels package, docs, and tests). The overall structure follows existing quantizer integrations, but there are a few blocking issues.
Security
_get_ops()insrc/diffusers/quantizers/nunchaku/utils.pycallsget_kernel(..., trust_remote_code=True)against a personal user repo (rootonchair/nunchaku-lite-kernels). This silently enables remote code execution for anyone who loads a Nunchaku Lite checkpoint, with no user opt-in. No other kernel usage in the repo does this (ggufandattention_dispatchcallget_kernelwithouttrust_remote_code). This needs to be removed or made an explicit user decision, and the kernel repo should ideally live under an org namespace with pinned revisions.
Description vs. diff mismatch
- The PR description claims "The loader reads safetensors metadata during
from_single_fileso Nunchaku Lite checkpoints can use their embedded runtime manifest" — but there are no changes tosingle_file_model.pyor any loader in this diff. Correspondingly, themetadataparameter ofNunchakuLiteQuantizer._process_model_before_weight_loadingis never populated by any call site and is dead code.
Correctness
NunchakuLiteTesterMixin._test_quantized_layersis tautological: it setsexpected_quantized_layers = num_quantized_layersand then asserts they're equal, so the count check can never fail. It should compare against the number of targets in the quantization config (like the base mixin compares against linear-layer count).- The docs and the
NunchakuLiteQuantizationConfigdocstring repeatedly referencemodel.json, but Diffusers model configs (and the test in this PR) useconfig.json. As written, users following the doc will put the quantization config in a file Diffusers never reads.
Style
nunchaku_quantizer.pyhas trailing whitespace (line 64) —make stylewas apparently not run.- New source files are missing the Apache license header used across
src/diffusers. NunchakuLiteQuantizationConfigis appended to_import_structurevia a standalone statement instead of being listed in the dict literal like every other unconditional export.import itertoolsis buried insidecheck_strict_state_dict_match; move it to module top level.SVDQW4A4Linearre-validatesprecision/group_sizethatNunchakuLiteQuantizationConfig.post_initalready enforces — per repo guidelines, drop the duplicated defensive checks (the forward path hard-codes group sizes 16/64 anyway, so thegroup_sizeargument is effectively unused at runtime).
serge v0.1.0 · model: claude-fable-5 · 10 LLM turns · 22 tool calls · 390.5s · 437374 in / 30106 out tokens
| if _ops is None: | ||
| from kernels import get_kernel | ||
|
|
||
| _ops = get_kernel(_HF_KERNEL_REPO, version=_HF_KERNEL_VERSION, trust_remote_code=True).ops |
There was a problem hiding this comment.
Security: trust_remote_code=True silently enables execution of arbitrary remote code from a personal user repo whenever a Nunchaku Lite checkpoint is loaded — the user never opts in. Neither of the existing get_kernel call sites in this repo (quantizers/gguf/utils.py, models/attention_dispatch.py) passes trust_remote_code. Please drop it (publish the kernel as a standard prebuilt kernels repo that doesn't require remote code), or at minimum surface this as an explicit user-facing opt-in. Hosting under a personal namespace (rootonchair/...) rather than an org also makes this a supply-chain risk for everyone using the backend.
| self, | ||
| model: "ModelMixin", | ||
| state_dict: dict[str, Any] | None = None, | ||
| metadata: dict[str, str] | None = None, |
There was a problem hiding this comment.
metadata is never passed by any caller — neither modeling_utils.py nor single_file_model.py forwards safetensors metadata to preprocess_model. The PR description claims the loader "reads safetensors metadata during from_single_file", but no loader changes exist in this diff. Per the repo guidelines (no unused parameters "for API consistency"), remove this parameter.
| self._verify_if_layer_quantized(name, module, config_kwargs) | ||
| num_quantized_layers += 1 | ||
|
|
||
| expected_quantized_layers = num_quantized_layers |
There was a problem hiding this comment.
This assertion is tautological: expected_quantized_layers is set to num_quantized_layers, so the num_quantized_layers == expected_quantized_layers check below can never fail, and num_fp32_modules is always 0. The only effective check left is num_quantized_layers > 0. Compute the expected count from the config instead, e.g. the total number of entries in svdq_w4a4["targets"] + awq_w4a16["targets"], so the test actually verifies all targets were replaced.
| The exported state dict must match the target Diffusers model architecture exactly. Checkpoints quantized with | ||
| fused QKV projections won't load into a model config that expects separate Q, K, and V projection modules. | ||
|
|
||
| Example compact `model.json` config: |
There was a problem hiding this comment.
Same as the doc page: Diffusers reads the model config from config.json, not model.json. Please fix the filename here and in the "quantization_config stored in model.json" sentence above so users don't package their checkpoints incorrectly.
| ], | ||
| } | ||
|
|
||
| _import_structure["quantizers.quantization_config"].append("NunchakuLiteQuantizationConfig") |
There was a problem hiding this comment.
Since this export is unconditional, add "NunchakuLiteQuantizationConfig" directly to the "quantizers.quantization_config" list in the _import_structure dict literal above instead of appending via a standalone statement — that matches how every other unconditional export is declared. Note also that NunchakuLiteQuantizationConfig.__init__ references torch unconditionally, while all other quantization configs exported here are gated behind is_torch_available(); consider whether this one needs the same guard.
| if device is None: | ||
| device = torch.device("cpu") | ||
|
|
||
| if precision not in {"int4", "nvfp4"}: |
There was a problem hiding this comment.
These precision/group_size checks duplicate validation that NunchakuLiteQuantizationConfig.post_init already enforces (and post_init is stricter: it pins group_size to 16 for fp4 / 64 for int4, while this accepts any positive value that forward then ignores — the activation-scale layout hard-codes 16/64). Per the repo's no-defensive-code guideline, drop the re-validation here and rely on the config.
Co-authored-by: sergereview[bot] <283583894+sergereview[bot]@users.noreply.github.com>
Co-authored-by: sergereview[bot] <283583894+sergereview[bot]@users.noreply.github.com>
Co-authored-by: sergereview[bot] <283583894+sergereview[bot]@users.noreply.github.com>
| if _ops is None: | ||
| from kernels import get_kernel | ||
|
|
||
| _ops = get_kernel(_HF_KERNEL_REPO, version=_HF_KERNEL_VERSION, trust_remote_code=True).ops |
There was a problem hiding this comment.
Should we make it a trusted organization @sayakpaul ? not sure if this is good to leave trust_remote_code=True. Also, I think it is fine to not pin the version unless we really need it no ?
There was a problem hiding this comment.
It's ALWAYS recommended to pin a version when using kernels. For now, trust_remote_code=True is fine. We can flip it once there's usage and more community trust.
| _HF_KERNEL_REPO = "rootonchair/nunchaku-lite-kernels" | ||
| _HF_KERNEL_VERSION = 1 | ||
|
|
||
|
|
||
| _ops = None | ||
|
|
||
|
|
||
| def _get_ops(): | ||
| global _ops | ||
| if _ops is None: | ||
| from kernels import get_kernel |
There was a problem hiding this comment.
let's make this more specific since this function is just there to load the nunchaku_kernels. maybe just all it nunchaku_kernels ?
There was a problem hiding this comment.
Additionally, can we avoid using global?
There was a problem hiding this comment.
This for loading the kernels lazily, I don't want the user to load the nunchaku kernels if they don't use it. Or am I just overthink it? 😅
There was a problem hiding this comment.
I think this has been resolved
| def check_strict_state_dict_match(model: nn.Module, state_dict: dict[str, Any]) -> None: | ||
| expected_keys = {n for n, _ in itertools.chain(model.named_parameters(), model.named_buffers())} | ||
| loaded_keys = set(state_dict.keys()) | ||
| missing_keys = sorted(expected_keys - loaded_keys) | ||
| unexpected_keys = sorted(loaded_keys - expected_keys) | ||
| if missing_keys or unexpected_keys: | ||
| message = "Nunchaku checkpoint keys must exactly match the patched model state dict." | ||
| if missing_keys: | ||
| message += f" Missing keys: {missing_keys[:10]}" | ||
| if len(missing_keys) > 10: | ||
| message += f" and {len(missing_keys) - 10} more" | ||
| message += "." | ||
| if unexpected_keys: | ||
| message += f" Unexpected keys: {unexpected_keys[:10]}" | ||
| if len(unexpected_keys) > 10: | ||
| message += f" and {len(unexpected_keys) - 10} more" | ||
| message += "." | ||
| raise ValueError(message) |
There was a problem hiding this comment.
we don't need this no ? we already print some logs when there are missing or unexpected keys.
There was a problem hiding this comment.
I want a hard fail on this as partial loaded model will make inference incorrect and produce noisy results
| precision="nvfp4" if precision == "fp4" else precision, | ||
| group_size=group_size, | ||
| torch_dtype=compute_dtype, | ||
| device=_module_device(module), |
There was a problem hiding this comment.
the module is practically always on meta device when loading with low_cpu_mem_usage. I think it should be fine to leave device to None no ? Also, if you want to be sure that the modules are loaded on meta device, we usually use
ctx = init_empty_weights if is_accelerate_available() else nullcontext
with ctx():
There was a problem hiding this comment.
yes, I think using ctx would be more correct
| def update_torch_dtype(self, torch_dtype): | ||
| if torch_dtype is None: | ||
| torch_dtype = self.compute_dtype | ||
| return torch_dtype |
There was a problem hiding this comment.
what is the torch_dtype passed in from_pretrained don't match the compute_dtype specified in the quantization_config ? there will be an issue no ? Should be overwrite compute_dtype in that case ?
| if op == "svdq_w4a4": | ||
| replacement = SVDQW4A4Linear( |
There was a problem hiding this comment.
this only works on specific machines, should we catch this beforehand or just let the user get an error when trying to get/run the kernels ?
There was a problem hiding this comment.
I think we should let the kernels library handling it when loading. @sayakpaul WDYT
There was a problem hiding this comment.
I have added the validation to validate_environment for this
| out_features, | ||
| rank=rank, | ||
| bias=bias is not None, | ||
| precision="nvfp4" if precision == "fp4" else precision, |
There was a problem hiding this comment.
why not just use nvfp4 instead of fp4 when naming the precision ?
There was a problem hiding this comment.
I think it will be a good idea to make the precision name shorter. But I will change it back
| | `svdq_w4a4` | `fp4` | 16 | Uses NVFP4 runtime kernels with SVDQ low-rank correction. | | ||
| | `svdq_w4a4` | `int4` | 64 | Uses INT4 W4A4 kernels with SVDQ low-rank correction. | | ||
| | `awq_w4a16` | `int4` | 64 | Uses INT4 weight-only AWQ-style kernels. | | ||
|
|
There was a problem hiding this comment.
maybe we should add some recommendation of hardware to be used with ?
There was a problem hiding this comment.
yes, I will be a good addition
There was a problem hiding this comment.
I think it was not addressed? I think should mention a matrix of the quant schemes and their supported hardware / CUDA capabilities. Also, the note on hopper since we cannot support it.
|
nice, I tested it and made a krea model, it works pretty good! 
import torch
from diffusers import DiffusionPipeline
pipe = DiffusionPipeline.from_pretrained("OzzyGT/Krea_2_Turbo_nunchaku_lite_nvfp4", torch_dtype=torch.bfloat16)
pipe.to("cuda")
prompt = (
"A cozy corner bookstore-cafe on a rainy evening, cinematic wide shot. "
'A large hand-lettered chalkboard sign in the window reads "FRESH COFFEE & OLD BOOKS" '
"and below it in smaller chalk letters \"open 'til late\". "
"Warm golden light spills onto wet cobblestones that mirror pink and blue neon reflections. "
"Inside, tall mahogany shelves are packed with hundreds of colorful book spines with tiny legible titles, "
"a barista in a striped apron pours delicate latte art, steam curling upward, "
"a tabby cat sleeps on a windowsill beside a stack of paperbacks. "
"Intricate detail, sharp focus, shallow depth of field, photorealistic, rich color grading."
)
image = pipe(
prompt,
num_inference_steps=8,
guidance_scale=0.0,
height=1024,
width=1024,
generator=torch.Generator("cuda").manual_seed(7),
).images[0]
image.save("sample.png") |
sayakpaul
left a comment
sayakpaul
left a comment
There was a problem hiding this comment.
Nice, thanks a lot for the further updates. I will post about the benchmarking I conducted through HF Jobs and Claude Code below.
| Compile the quantized transformer after loading the pipeline for faster inference. | ||
|
|
||
| ```python | ||
| pipe.transformer = torch.compile(pipe.transformer, mode="default", fullgraph=False) |
There was a problem hiding this comment.
Why can't we do it fullgraph=True? Can you open an issue on the nunchaku-lite repo with the trace?
|
|
||
| # Nunchaku Lite | ||
|
|
||
| Nunchaku Lite is a quantization backend for loading prequantized checkpoints in Diffusers. Use |
There was a problem hiding this comment.
Should provide some citations to the original Nunchaku repo since Nunchaku Lite is based off of it?
There was a problem hiding this comment.
Not only the repo, I will reference their paper too
There was a problem hiding this comment.
this is resolved
| use_pe=False, | ||
| ).images[0] | ||
| image.save("ernie-image-turbo-nunchaku-lite.png") | ||
| ``` |
There was a problem hiding this comment.
Is it compatible with offloading, too?
| | `svdq_w4a4` | `fp4` | 16 | Uses NVFP4 runtime kernels with SVDQ low-rank correction. | | ||
| | `svdq_w4a4` | `int4` | 64 | Uses INT4 W4A4 kernels with SVDQ low-rank correction. | | ||
| | `awq_w4a16` | `int4` | 64 | Uses INT4 weight-only AWQ-style kernels. | | ||
|
|
There was a problem hiding this comment.
I think it was not addressed? I think should mention a matrix of the quant schemes and their supported hardware / CUDA capabilities. Also, the note on hopper since we cannot support it.
| self.compute_dtype = quantization_config.compute_dtype | ||
| self.pre_quantized = quantization_config.pre_quantized | ||
|
|
||
| def validate_environment(self, *args, **kwargs): |
There was a problem hiding this comment.
Hopper validation added
| def update_torch_dtype(self, torch_dtype): | ||
| if torch_dtype is None: | ||
| torch_dtype = self.compute_dtype | ||
| return torch_dtype |
|
|
||
| self.post_init() | ||
|
|
||
| def post_init(self): |
| @@ -0,0 +1,161 @@ | |||
| import json | |||
|
Even though the below is generated with Claude Code (I have manually verified them). Nunchaku Lite (PR #14100) — testing notesI ran the snippet from Reproduced the doc benchmark ✅Ran the exact doc snippet (
These match the doc's reported 2.271 s → 1.675 s almost exactly, so the NVFP4 + Rough edge 1: capability check is more permissive than the shipped kernels
Consequence: on an H100 (sm_90) the INT4 path passes which is a confusing failure mode (looks like a runtime bug, not an unsupported-hardware message). So Hopper is effectively unsupported for INT4 today even though the check green-lights it. Suggestions (pick whatever fits):
Rough edge 2: kernels only ship for torch 2.11 / 2.12
For Blackwell specifically, the working combo is torch 2.11 + cu128 ( Suggestion: document the supported torch versions in Environment
|
|
Now running the tests. |
|
Testing report (again manually verified). Nunchaku Lite (PR #14100) — test run on HF JobsRan the PR's test suite Setup
Result9 passed in 4.66s ✅ (pytest exit code 0), including One caveat (not a PR bug)The 3 tests that download the kernel initially failed due to an external dependency bug:
Suggestion: report to |
Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>
rootonchair
force-pushed
the
feature/nunchaku-lite-single-file
branch
from
03399a0 to
5522954
Compare
rootonchair
force-pushed
the
feature/nunchaku-lite-single-file
branch
from
5522954 to
86cd7c8
Compare
…chair/diffusers into feature/nunchaku-lite-single-file
5 participants
What does this PR do?
Adds Nunchaku Lite single-file checkpoint loading for Diffusers models.
This introduces
NunchakuLiteQuantizationConfigand a new Nunchaku Lite quantizer that can patch supported nn.Linear modules into runtime SVDQ/AWQ linear layers before strict checkpoint loading. The loader reads safetensors metadata duringfrom_single_fileso Nunchaku Lite checkpoints can use their embedded runtime manifest to decide which modules to replace.Deprecated API
New API for
from_single_fileuseFixes # (issue)
Before submitting
.ai/review-rules.md?documentation guidelines, and
here are tips on formatting docstrings.
Who can review?
Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.