up

Author: sayakpaul

src/diffusers/quantizers/nunchaku/nunchaku_quantizer.py

@@ -20,11 +20,6 @@
 if is_torch_available():
     import torch
 
-if is_accelerate_available():
-    pass
-
-if is_nunchaku_available():
-    from .utils import replace_with_nunchaku_linear
 
 logger = logging.get_logger(__name__)
 
@@ -79,13 +74,14 @@ def check_if_quantized_param(
         state_dict: Dict[str, Any],
         **kwargs,
     ):
-        from nunchaku.models.linear import SVDQW4A4Linear
-
-        module, tensor_name = get_module_from_name(model, param_name)
-        if self.pre_quantized and isinstance(module, SVDQW4A4Linear):
-            return True
-
-        return False
+        # TODO: revisit
+        # Check if the param_name is not in self.modules_to_not_convert
+        if any((key + "." in param_name) or (key == param_name) for key in self.modules_to_not_convert):
+            return False
+        else:
+            # We only quantize the weight of nn.Linear
+            module, _ = get_module_from_name(model, param_name)
+            return isinstance(module, torch.nn.Linear)
 
     def create_quantized_param(
         self,
@@ -112,13 +108,32 @@ def create_quantized_param(
                 module._buffers[tensor_name] = torch.nn.Parameter(param_value.to(target_device))
 
         elif isinstance(module, torch.nn.Linear):
-            if tensor_name in module._parameters:
-                module._parameters[tensor_name] = torch.nn.Parameter(param_value).to(device=target_device)
-            if tensor_name in module._buffers:
-                module._buffers[tensor_name] = torch.nn.Parameter(param_value).to(target_device)
-
-            new_module = SVDQW4A4Linear.from_linear(module)
-            setattr(model, param_name, new_module)
+            # TODO: this returns an `SVDQW4A4Linear` layer initialized from the corresponding `linear` module.
+            # But we need to have a utility that can take a pretrained param value and quantize it. Not sure
+            # how to do that yet.
+            # Essentially, we need something like `bnb.nn.Params4bit.from_prequantized`. Or is there a better
+            # way to do it?
+            is_param = tensor_name in module._parameters
+            is_buffer = tensor_name in module._buffers
+            new_module = SVDQW4A4Linear.from_linear(
+                module, precision=self.quantization_config.precision, rank=self.quantization_config.rank
+            )
+            module_name = ".".join(param_name.split(".")[:-1])
+            if "." in module_name:
+                parent_name, leaf = module_name.rsplit(".", 1)
+                parent = model.get_submodule(parent_name)
+            else:
+                parent, leaf = model, module_name
+
+            # rebind
+            # this will result into
+            # AttributeError: 'SVDQW4A4Linear' object has no attribute 'weight'. Did you mean: 'qweight'.
+            if is_param:
+                new_module._parameters[tensor_name] = torch.nn.Parameter(param_value).to(device=target_device)
+            elif is_buffer:
+                new_module._buffers[tensor_name] = torch.nn.Parameter(param_value).to(device=target_device)
+
+            setattr(parent, leaf, new_module)
 
     def adjust_max_memory(self, max_memory: Dict[str, Union[int, str]]) -> Dict[str, Union[int, str]]:
         max_memory = {key: val * 0.90 for key, val in max_memory.items()}
@@ -157,24 +172,25 @@ def _process_model_before_weight_loading(
         keep_in_fp32_modules: List[str] = [],
         **kwargs,
     ):
-        # TODO: deal with `device_map`
         self.modules_to_not_convert = self.quantization_config.modules_to_not_convert
 
         if not isinstance(self.modules_to_not_convert, list):
             self.modules_to_not_convert = [self.modules_to_not_convert]
 
         self.modules_to_not_convert.extend(keep_in_fp32_modules)
+
+        # TODO: revisit
+        # Extend `self.modules_to_not_convert` to keys that are supposed to be offloaded to `cpu` or `disk`
+        # if isinstance(device_map, dict) and len(device_map.keys()) > 1:
+        #     keys_on_cpu = [key for key, value in device_map.items() if value in ["disk", "cpu"]]
+        #     self.modules_to_not_convert.extend(keys_on_cpu)
+
         # Purge `None`.
         # Unlike `transformers`, we don't know if we should always keep certain modules in FP32
         # in case of diffusion transformer models. For language models and others alike, `lm_head`
         # and tied modules are usually kept in FP32.
         self.modules_to_not_convert = [module for module in self.modules_to_not_convert if module is not None]
 
-        model = replace_with_nunchaku_linear(
-            model,
-            modules_to_not_convert=self.modules_to_not_convert,
-            quantization_config=self.quantization_config,
-        )
         model.config.quantization_config = self.quantization_config
 
     def _process_model_after_weight_loading(self, model, **kwargs):