Nunchaku Lite quantization in place

docs/source/en/quantization/nunchaku.md

@@ -122,6 +122,34 @@ List each module you want to quantize under `svdq_w4a4` or `awq_w4a16`. A module
 }
 ```
 
+## Data-free quantization on load
+
+Pass `pre_quantized=False` to quantize an *unquantized* checkpoint at load time with data-free SVDQuant — no calibration data is needed. Each targeted linear gets weight-span smoothing, a rank-`r` SVD low-rank branch, and int4 or nvfp4 group quantization of the residual, packed directly into the kernel layout. Only `svdq_w4a4` targets are supported in this mode, and each target's `in_features`/`out_features` must be multiples of 128 (with `rank` a multiple of 16, or 0 to disable the low-rank branch).
+
+When `targets` is omitted, eligible targets are inferred automatically from the model's structure: quantization is restricted to the repeated transformer-block stacks (so embedders, final projections, and modulation heads outside the stacks stay unquantized), adaLN-style linears are skipped via the default `("norm", "modulation")` name patterns, and every remaining `nn.Linear` satisfying the packing constraints is selected. The model's `_keep_in_fp32_modules` is always honored. No configuration is needed for typical DiTs:
+
+```python
+import torch
+from diffusers import Flux2Transformer2DModel, NunchakuLiteQuantizationConfig
+
+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",
+)
+```
+
+Pass `exclude_targets` (substring match; it replaces the default name patterns) to exclude further modules, or an explicit `targets` list for full control. Quantization happens per weight as the checkpoint streams in, so peak memory stays near the quantized model size. The result is identical to loading a checkpoint produced offline by a data-free SVDQuant exporter.
+
+Two quality knobs are worth tuning per model. `rank` sets the size of the bf16 low-rank branch that absorbs the largest singular components before 4-bit coding — larger ranks trade a modest size increase for fidelity, and the gain can be substantial (on a 19B video DiT, int4 at `rank=128` recovered ~2.7 dB over `rank=32` for ~16% more transformer memory). `smooth_exponent` (default `0.5`) sets the weight-span smoothing strength; it interacts with `rank` — larger low-rank branches tend to prefer weaker smoothing (e.g. `0.25` at `rank=128`) — so when raising `rank`, sweep the exponent rather than assuming the default.
+
+Quantize-on-load runs one SVD per target at load time (minutes for large models on a fast GPU). For repeated loads of the same configuration, prefer a pre-quantized checkpoint produced by an offline exporter — loading packed weights takes seconds, and the packed format is identical.
+
 ## Fused kernels
 
 The original [Nunchaku](https://github.com/nunchaku-ai/nunchaku) engine gets much of its speed from model-specific fused execution paths. It combines the Q, K, and V projections with RMSNorm and RoPE, and uses a fused GELU kernel for the MLP. Nunchaku Lite instead uses the standard Diffusers model with generic quantized linear layers, so it does not include these fusions.

src/diffusers/models/model_loading_utils.py

@@ -234,6 +234,18 @@ def load_model_dict_into_meta(
 
     for param_name, param in state_dict.items():
         if param_name not in empty_state_dict:
+            # A quantize-on-load quantizer may claim checkpoint keys that no longer
+            # exist on the model (e.g. `weight` of a module replaced by a quantized
+            # linear) and materialize the module's packed parameters from them.
+            if (
+                is_quantized
+                and not hf_quantizer.pre_quantized
+                and hf_quantizer.check_if_quantized_param(model, param, param_name, state_dict)
+            ):
+                param_device = _determine_param_device(param_name, device_map)
+                hf_quantizer.create_quantized_param(
+                    model, param, param_name, param_device, state_dict, unexpected_keys, dtype=dtype
+                )
             continue
 
         set_module_kwargs = {}

src/diffusers/quantizers/nunchaku/nunchaku_quantizer.py

@@ -6,6 +6,8 @@
 
 
 if TYPE_CHECKING:
+    import torch
+
     from ...models.modeling_utils import ModelMixin
 
 
@@ -19,7 +21,9 @@ class NunchakuLiteQuantizer(DiffusersQuantizer):
     def __init__(self, quantization_config, **kwargs):
         super().__init__(quantization_config, **kwargs)
         self.compute_dtype = quantization_config.compute_dtype
-        self.pre_quantized = quantization_config.pre_quantized
+        # Quantize on load when either the loader inferred an unquantized
+        # checkpoint or the config explicitly requested `pre_quantized=False`.
+        self.pre_quantized = self.pre_quantized and quantization_config.pre_quantized
 
     def validate_environment(self, *args, **kwargs):
         if not is_kernels_available():
@@ -66,13 +70,87 @@ def _process_model_before_weight_loading(
     ):
         from .utils import check_strict_state_dict_match, replace_with_nunchaku_linear
 
+        svdq_config = self.quantization_config.svdq_w4a4
+        if not self.pre_quantized and svdq_config is not None and svdq_config.get("targets") is None:
+            from .svdquant import infer_data_free_targets
+
+            svdq_config["targets"] = infer_data_free_targets(
+                model,
+                group_size=svdq_config["group_size"],
+                exclude_targets=self.quantization_config.exclude_targets or (),
+            )
+            logger.info(f"Inferred {len(svdq_config['targets'])} data-free quantization targets.")
+
         quantization_config = self.quantization_config.to_dict()
         num_replaced = replace_with_nunchaku_linear(model, quantization_config, self.compute_dtype)
 
-        if state_dict is not None:
+        if self.pre_quantized and state_dict is not None:
             check_strict_state_dict_match(model, state_dict)
         logger.info(f"Applied Nunchaku quantization config with {num_replaced} targets.")
 
+    def update_missing_keys(self, model, missing_keys: list[str], prefix: str) -> list[str]:
+        if self.pre_quantized:
+            return missing_keys
+        # In data-free mode the checkpoint holds `weight`/`bias` while the model
+        # expects the packed parameters; those are produced at load time.
+        from .svdquant import DATA_FREE_PARAMETER_NAMES
+
+        return [key for key in missing_keys if key.rpartition(".")[2] not in DATA_FREE_PARAMETER_NAMES]
+
+    def check_if_quantized_param(
+        self,
+        model: "ModelMixin",
+        param_value: "torch.Tensor",
+        param_name: str,
+        state_dict: dict[str, Any],
+        **kwargs,
+    ) -> bool:
+        if self.pre_quantized:
+            return False
+        from .utils import SVDQW4A4Linear
+
+        module_name, _, tensor_name = param_name.rpartition(".")
+        if tensor_name not in ("weight", "bias") or not module_name:
+            return False
+        try:
+            module = model.get_submodule(module_name)
+        except AttributeError:
+            return False
+        return isinstance(module, SVDQW4A4Linear)
+
+    def create_quantized_param(
+        self,
+        model: "ModelMixin",
+        param_value: "torch.Tensor",
+        param_name: str,
+        target_device: "torch.device",
+        state_dict: dict[str, Any] | None = None,
+        unexpected_keys: list[str] | None = None,
+        **kwargs,
+    ):
+        import torch
+
+        from .svdquant import pack_data_free_bias, quantize_linear_data_free
+
+        module_name, _, tensor_name = param_name.rpartition(".")
+        module = model.get_submodule(module_name)
+        if unexpected_keys is not None and param_name in unexpected_keys:
+            unexpected_keys.remove(param_name)
+        if tensor_name == "bias":
+            packed_bias = pack_data_free_bias(param_value.to(target_device), torch_dtype=self.compute_dtype)
+            module._parameters["bias"] = torch.nn.Parameter(packed_bias, requires_grad=False)
+            return
+        quantized = quantize_linear_data_free(
+            param_value.to(target_device),
+            precision=module.precision,
+            group_size=module.group_size,
+            rank=module.rank,
+            torch_dtype=self.compute_dtype,
+            smooth_exponent=self.quantization_config.svdq_w4a4.get("smooth_exponent", 0.5),
+        )
+        for name, tensor in quantized.items():
+            module._parameters[name] = torch.nn.Parameter(tensor.to(target_device), requires_grad=False)
+
     def _process_model_after_weight_loading(self, model: "ModelMixin", **kwargs):
         return model