Nunchaku Lite quantization in place - #14608
Draft
rootonchair wants to merge 7 commits into
Draft
Conversation
Support `pre_quantized=False` in NunchakuLiteQuantizationConfig: targeted linears of an unquantized checkpoint are quantized at load time with data-free SVDQuant (weight-span smoothing, rank-r SVD low-rank branch, int4/nvfp4 group quantization) and packed directly into the kernel layout SVDQW4A4Linear consumes — no calibration data needed. The math lives in quantizers/nunchaku/data_free.py, which is pure torch and stays importable without the `kernels` package; quantization happens per weight in create_quantized_param so peak memory stays near the quantized model size. Packed outputs are tensor-for-tensor identical to DeepCompressor's Nunchaku W4A4 converter (verified against it on random weights; qweight byte-identical). `awq_w4a16` targets are not supported in this mode and raise. CPU tests validate shapes, round-trip reconstruction error, bias packing, and the quantizer flow; a gated GPU mixin test runs quantize-on-load end to end where kernels are available. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_013NAtDkGmAw2fzbcvjfC79w
When `pre_quantized=False` and `svdq_w4a4.targets` is omitted, the quantizer now infers targets from the model at load time: every nn.Linear whose dimensions satisfy the Nunchaku packing constraints is selected, minus modules matched by the new `modules_to_not_convert` config option or listed in the model's `_keep_in_fp32_modules`. Explicit target lists keep working. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_013NAtDkGmAw2fzbcvjfC79w
Auto-inference no longer needs a hand-written modules_to_not_convert list:
targets are restricted to the model's repeated transformer-block stacks
(identical-class nn.ModuleLists), which structurally excludes embedders,
final projections, and modulation heads, and adaLN-style linears inside
blocks are skipped via default ("norm", "modulation") name patterns. An
explicit modules_to_not_convert replaces the default patterns. For
FLUX.2-klein-9B the zero-config inferred target set matches the curated
list exactly (144 targets).
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_013NAtDkGmAw2fzbcvjfC79w
Clearer pairing with the svdq_w4a4 `targets` field, and avoids implying the bnb/torchao semantics of keeping modules in high precision at load: the option only filters data-free target inference. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_013NAtDkGmAw2fzbcvjfC79w
Match the bitsandbytes loader contract: filter the load-time-produced packed parameter names out of missing_keys via update_missing_keys, and remove the consumed `weight`/`bias` checkpoint keys from unexpected_keys inside create_quantized_param. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_013NAtDkGmAw2fzbcvjfC79w
The module implements the SVDQuant math (smoothing, low-rank split, quantization, kernel packing) as opposed to utils.py's kernel runtime; name it after the algorithm. Data-free stays in the function names, where it describes the mode. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_013NAtDkGmAw2fzbcvjfC79w
Name the integration test after the loader mechanism (pre_quantized=False) rather than the algorithm mode, and rename the companion class attribute to quantize_on_load_config_dict to match. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_013NAtDkGmAw2fzbcvjfC79w
github-actions
Bot
added
documentation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What does this PR do?
Fixes # (issue)
Before submitting
self-reviewskill on the diff?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.