Nunchaku Lite quantization in place

[rootonchair]

What does this PR do?

Adds data-free quantize-on-load to the Nunchaku Lite backend: pass pre_quantized=False to quantize an unquantized checkpoint into the Nunchaku W4A4 packed format while its weights stream in — no calibration data, no offline conversion step.

transformer = Flux2Transformer2DModel.from_pretrained(
    "black-forest-labs/FLUX.2-klein-9B",
    subfolder="transformer",
    quantization_config=NunchakuLiteQuantizationConfig(
        svdq_w4a4={"precision": "nvfp4", "group_size": 16, "rank": 32},
        pre_quantized=False,
    ),
    torch_dtype=torch.bfloat16,
    device_map="cuda",
)

For each targeted linear: weight-span smoothing, a rank-r SVD low-rank branch, and int4/nvfp4 group quantization of the residual, packed directly into the kernel layout (quantizers/nunchaku/svdquant.py, pure PyTorch). When targets is omitted they are inferred structurally (repeated block stacks, adaLN/_keep_in_fp32_modules exclusions, packing-constraint filter) — zero configuration on typical DiTs. Includes a fix in load_model_dict_into_meta letting a quantize-on-load quantizer claim checkpoint keys that no longer exist on the model (replaced modules have qweight/… instead of weight), and an optional smooth_exponent knob.

Validation

• Packing verified bit-identical to DeepCompressor-style exporters (weight codes, group/micro scales, vectors, low-rank layout); end-to-end reconstruction error matches an offline exporter to four decimals. Loading published pre-quantized checkpoints with this backend cross-validates both directions.
• Auto-inference selected the correct target sets with no configuration on two architectures: 1,344 layers on LTX-2.3 (19B audio-video DiT), 144 on FLUX.2-klein-9B — matching hand-curated published checkpoints exactly.
• Quality: seed-matched against bf16 on both models, data-free output is statistically equivalent to calibrated (smoothing-search + GPTQ) checkpoints of the same config — e.g. LTX-2.3 int4: 13.3 dB (data-free) vs 12.6 dB (calibrated) vs the same bf16 reference; nvfp4: 17.3 vs 15.3 dB. The trade against pre-quantized checkpoints is purely load time (~15 min of on-load SVDs for 19B vs ~40 s), not fidelity.
• Tuning: rank is the effective quality lever (int4 rank=128: +2.7 dB over rank=32 for +16% transformer memory) and interacts with smooth_exponent (larger ranks prefer weaker smoothing) — documented in the quantization docs.

Hardware notes

nvfp4 requires Blackwell (sm_100+); int4 runs on Turing/Ampere/Ada and is the pre-Blackwell path. Known kernel-level issue, not introduced here: the int4 path is ~3-4x slower than nvfp4 on sm_120 (fine on A100), reproducible with existing pre-quantized checkpoints.

Self-review notes (AI-assisted development)

• The meta-tensor loading fix was found by the first real-checkpoint run; NunchakuLiteTesterMixin.test_nunchaku_lite_quantize_on_load exists but no test class subclasses it yet, so quantize-on-load still lacks an e2e regression test against a real model — known gap, happy to wire a small-model test if reviewers want it in this PR.
• A smooth_overrides mechanism (injecting precomputed per-layer smoothing vectors) was built for ablations and deliberately left out to keep the API surface minimal.
• CPU-side unit tests validate the packing math against reference unpackers (15 tests).

[rootonchair]

Experimental validation of data-free quantize-on-load

Evidence appendix for the claims in the PR description. All quality numbers use a seed-matched protocol: identical prompts, seeds, and GPU per comparison, scored as PSNR/SSIM against a bf16 reference — so they measure pure quantization-induced divergence. All measurements were taken with guidance genuinely disabled; note that on current main, guidance_scale=1.0 does not disable guidance for LTX-2 (see #14650) — any quality measurement of these models is contaminated without those fixes or explicit flags.

1. Data-free vs calibrated checkpoints (LTX-2.3-Distilled v1.1, 19B, t2v, 5 seeds)

Calibrated = published checkpoints produced offline with per-layer grid-searched smoothing (scored on calibration activations) + GPTQ residual rounding, on identical targets/rank/group size. Same-GPU bf16 references.

run	GPU	PSNR vs bf16	SSIM
data-free nvfp4 (rank 32)	RTX PRO 6000	17.33 ± 4.07	0.60
calibrated nvfp4	RTX PRO 6000	15.34 ± 5.59	0.52
data-free int4 (rank 32)	A100	13.29 ± 2.97	0.42
calibrated int4	A100	12.62 ± 3.26	0.39

Data-free matches or nominally exceeds calibrated on both precisions (consistent across three protocol variants). An ablation transplanting the calibrated checkpoints' per-layer smoothing vectors into the data-free pipeline changed nothing (13.18 vs 13.71 dB, within noise) — the calibrated pipeline's extra machinery does not help at the trajectory level on this model. The practical difference is load time only: ~40 s (pre-quantized) vs ~15–22 min of on-load SVDs (19B model).

2. Tuning: smooth_exponent and rank (int4, same protocol)

Smoothing-exponent sweep (at the default rank 32) — flat within noise; the default 0.5 is fine and the knob alone is not a quality lever:

smooth_exponent	0.1	0.25	0.5 (default)
PSNR vs bf16	13.39 ± 3.19	12.53 ± 1.94	13.29 ± 2.97
SSIM	0.44	0.41	0.42

Rank sweep (at the default exponent 0.5) — rank moves quality substantially:

rank	32 (default)	64	128
PSNR vs bf16	13.29 ± 2.97	15.11 ± 4.47	13.65 ± 2.47
SSIM	0.42	0.52	0.45
transformer size	10.6 GB	11.2 GB	12.3 GB

Full grid — the two knobs interact: the rank-128 optimum sits at a weaker exponent (0.25), plausibly because strong smoothing and a large SVD branch compete for the same outlier energy. Each row is the exponent sweep at a fixed rank; each column the rank sweep at a fixed exponent:

PSNR / SSIM	e=0.0	e=0.1	e=0.25	e=0.5
rank 32	—	13.39 / 0.44	12.53 / 0.41	13.29 / 0.42
rank 64	13.25 / 0.43	13.68 / 0.46	14.27 / 0.47	15.11 / 0.52
rank 128	13.93 / 0.48	14.57 / 0.53	15.99 / 0.55	13.65 / 0.45

Best cell: rank 128 / exponent 0.25 — +2.7 dB over the defaults for ~16% more transformer memory and ~unchanged inference time (wins on 5/5 samples). Practical guidance documented in this PR's quantization docs: when raising rank, sweep the exponent rather than assuming 0.5.

Both knobs are set in the svdq_w4a4 section:

transformer = LTX2VideoTransformer3DModel.from_pretrained(
    "rootonchair/LTX-2.3-Distilled-v1.1-Diffusers",
    subfolder="transformer",
    quantization_config=NunchakuLiteQuantizationConfig(
        svdq_w4a4={
            "precision": "int4",
            "group_size": 64,
            "rank": 128,            # low-rank branch size (multiple of 16, or 0 to disable)
            "smooth_exponent": 0.25,  # optional; defaults to 0.5
        },
        pre_quantized=False,
    ),
    torch_dtype=torch.bfloat16,
    device_map="cuda",
)

smooth_exponent controls the data-free weight-span smoothing. A per-input-channel scale s_j = 1 / absmax(W[:, j]) ** smooth_exponent is computed, the weight columns are multiplied by it (W' = W · diag(s), equalizing per-channel magnitudes) before the SVD and residual quantization, and at runtime the kernel divides the corresponding activation channels by the same s — so W' · (x / s) = W · x exactly, and the output needs no correction. Smoothing never changes the linear map; it only relocates where the 4-bit quantization error lands across channels. 0.0 disables smoothing, 1.0 fully flattens per-channel weight spans. It is the weight-side analogue of SmoothQuant's exponent — with no calibration data there is no activation term, which is why it is a single scalar rather than a per-layer searched (α, β) pair.

3. Image-to-video

Same protocol, conditioning image + prompt + seed matched, LTX2ImageToVideoPipeline:

run	PSNR vs bf16	SSIM	frame-0 adherence to conditioning
bf16	—	—	33.70 dB
data-free int4, r128/e0.25	21.91 ± 2.45	0.86	33.45 dB
data-free int4, r64/e0.5	21.61 ± 1.94	0.85	33.34 dB

Conditioning collapses the divergence: quantized i2v is near-reference (~22 dB / 0.86 SSIM vs ~16 dB / 0.55 unconditioned), and the conditioning pathway itself is unaffected (frame-0 adherence equal to bf16; the ~34 dB ceiling is the model's own CRF-33 image recompression).

4. Second architecture: FLUX.2-klein-9B (image, 4-step distilled)

Auto-inference selected all 144 eligible block linears with zero configuration (fused SwiGLU/QKV-MLP projections included; dims-based exclusions correct). nvfp4: 18.1 ± 3.3 dB / SSIM 0.70 vs bf16; int4: 16.4 ± 2.3 / 0.61; the smoothing exponent is quality-neutral on this architecture across 0.0–0.75.

Hardware caveat (pre-existing, kernel-level — not introduced here): the int4 kernel path is ~3–4× slower than nvfp4 on sm_120 (reproduces with published pre-quantized checkpoints; e.g. 427 s vs 113 s per five 121-frame clips). On A100 (sm_80) int4 performs normally.

---

Quantize-on-load cost for reference (one torch.linalg.svd per target):

model	targets	GPU	quantize-on-load	pre-quantized load
FLUX.2-klein-9B	144	RTX PRO 6000	~3.5 min	—
LTX-2.3-Distilled 19B	1,344	RTX PRO 6000	~14.5 min	~40 s
LTX-2.3-Distilled 19B	1,344	A100	~22 min	~50 s

Times are effectively independent of precision, rank, and smooth_exponent — the full SVD dominates and is shared by all configurations (across the entire rank/exponent grid on the A100, load times span 1327–1341 s, ~1%). Tuning those knobs costs nothing at quantize time.

Per-sample raw data (frames, metrics, seeds) for every table is retained and available on request.