SDNQ add training related changes

modules/sdnq/dequantizer.py

@@ -164,30 +164,54 @@ def re_quantize_matmul_packed_int_symmetric(weight: torch.ByteTensor, scale: tor
 
 @devices.inference_context()
 def dequantize_layer_weight(self: torch.nn.Module, inplace: bool = False):
-    weight = self.sdnq_dequantizer(self.weight, self.scale, self.zero_point, self.svd_up, self.svd_down, skip_quantized_matmul=self.sdnq_dequantizer.use_quantized_matmul)
+    weight = torch.nn.Parameter(self.sdnq_dequantizer(self.weight, self.scale, self.zero_point, self.svd_up, self.svd_down, skip_quantized_matmul=self.sdnq_dequantizer.use_quantized_matmul), requires_grad=True)
+    forward = getattr(torch.nn, self.sdnq_dequantizer.layer_class_name).forward
     if inplace:
-        self.weight.data = weight
-        self.forward = getattr(torch.nn, self.sdnq_dequantizer.layer_class_name).forward
+        self.weight = weight
+        self.forward = forward
+        self.forward = self.forward.__get__(self, self.__class__)
         del self.sdnq_dequantizer, self.scale, self.zero_point, self.svd_up, self.svd_down
-    return weight
+        return self
+    else:
+        return weight, forward
 
 
 @devices.inference_context()
-def dequantize_sdnq_model(model: torch.nn.Module):
+def dequantize_sdnq_module(model: torch.nn.Module):
     if hasattr(model, "sdnq_dequantizer"):
-        model.weight.data = dequantize_layer_weight(model, inplace=True)
+        model = dequantize_layer_weight(model, inplace=True)
     has_children = list(model.children())
     if not has_children:
         return model
     for module_name, module in model.named_children():
         if hasattr(module, "sdnq_dequantizer"):
-            module.weight.data = dequantize_layer_weight(module, inplace=True)
-            setattr(model, module_name, module)
+            setattr(model, module_name, dequantize_layer_weight(module, inplace=True))
         else:
             setattr(model, module_name, dequantize_sdnq_model(module))
     return model
 
 
+@devices.inference_context()
+def dequantize_sdnq_model(model: torch.nn.Module):
+    model = dequantize_sdnq_module(model)
+    if hasattr(model, "quantization_method"):
+        del model.quantization_method
+    if hasattr(model, "quantization_config"):
+        del model.quantization_config
+    if hasattr(model, "config"):
+        try:
+            if hasattr(model.config, "quantization_config"):
+                del model.config.quantization_config
+        except Exception:
+            pass
+        try:
+            if hasattr(model.config, "pop"):
+                model.config.pop("quantization_config", None)
+        except Exception:
+            pass
+    return model
+
+
 class SDNQDequantizer():
     def __init__(
         self,
@@ -207,6 +231,7 @@ class SDNQDequantizer():
     ):
         self.is_packed = dtype_dict[weights_dtype]["is_packed"]
         self.is_unsigned = dtype_dict[weights_dtype]["is_unsigned"]
+        self.is_integer = dtype_dict[weights_dtype]["is_integer"]
         self.result_dtype = result_dtype
         self.result_shape = result_shape
         self.original_shape = original_shape

modules/sdnq/forward.py

@@ -2,10 +2,10 @@
 
 from typing import Callable
 
-from .common import conv_types, conv_transpose_types
+from .common import conv_types, conv_transpose_types, use_tensorwise_fp8_matmul
 
 
-def get_forward_func(layer_class_name: str, use_quantized_matmul: bool, is_integer: bool, use_tensorwise_fp8_matmul: bool) -> Callable: # pylint: disable=inconsistent-return-statements
+def get_forward_func(layer_class_name: str, use_quantized_matmul: bool, is_integer: bool) -> Callable: # pylint: disable=inconsistent-return-statements
     if layer_class_name in conv_types:
         if use_quantized_matmul:
             if is_integer:

modules/sdnq/quantizer.py

@@ -24,6 +24,7 @@ from .forward import get_forward_func
 
 class QuantizationMethod(str, Enum):
     SDNQ = "sdnq"
+    SDNQ_TRAINING = "sdnq_training"
 
 
 @devices.inference_context()
@@ -396,13 +397,13 @@ def sdnq_quantize_layer(layer, weights_dtype="int8", torch_dtype=None, group_siz
         layer.svd_down = torch.nn.Parameter(layer.svd_down.to(return_device, non_blocking=non_blocking), requires_grad=False)
 
     layer = layer.to(return_device, non_blocking=non_blocking)
-    layer.forward = get_forward_func(layer_class_name, use_quantized_matmul, dtype_dict[weights_dtype]["is_integer"], use_tensorwise_fp8_matmul)
+    layer.forward = get_forward_func(layer_class_name, use_quantized_matmul, layer.sdnq_dequantizer.is_integer)
     layer.forward = layer.forward.__get__(layer, layer.__class__)
     return layer
 
 
 @devices.inference_context()
-def apply_sdnq_to_module(model, weights_dtype="int8", torch_dtype=None, group_size=0, svd_rank=32, svd_steps=8, use_svd=False, quant_conv=False, use_quantized_matmul=False, use_quantized_matmul_conv=False, use_stochastic_rounding=False, dequantize_fp32=False, non_blocking=False, quantization_device=None, return_device=None, modules_to_not_convert: List[str] = None, modules_dtype_dict: Dict[str, List[str]] = None, full_param_name="", op=None): # pylint: disable=unused-argument
+def apply_sdnq_to_module(model, weights_dtype="int8", torch_dtype=None, group_size=0, svd_rank=32, svd_steps=8, use_svd=False, quant_conv=False, use_quantized_matmul=False, use_quantized_matmul_conv=False, use_stochastic_rounding=False, dequantize_fp32=False, non_blocking=False, quantization_device=None, return_device=None, modules_to_not_convert: List[str] = None, modules_dtype_dict: Dict[str, List[str]] = None, full_param_name=""): # pylint: disable=unused-argument
     has_children = list(model.children())
     if not has_children:
         return model
@@ -460,15 +461,14 @@ def apply_sdnq_to_module(model, weights_dtype="int8", torch_dtype=None, group_si
                 modules_to_not_convert=modules_to_not_convert,
                 modules_dtype_dict=modules_dtype_dict,
                 full_param_name=param_name,
-                op=op,
             ))
     return model
 
 
 @devices.inference_context()
 def sdnq_post_load_quant(
-    model,
-    weights_dtype="int8",
+    model: torch.nn.Module,
+    weights_dtype: str = "int8",
     torch_dtype: torch.dtype = None,
     group_size: int = 0,
     svd_rank: int = 32,
@@ -481,11 +481,10 @@ def sdnq_post_load_quant(
     dequantize_fp32: bool = False,
     non_blocking: bool = False,
     add_skip_keys:bool = True,
-    quantization_device: torch.device = None,
-    return_device: torch.device = None,
+    quantization_device: Optional[torch.device] = None,
+    return_device: Optional[torch.device] = None,
     modules_to_not_convert: List[str] = None,
     modules_dtype_dict: Dict[str, List[str]] = None,
-    op=None,
 ):
     if modules_to_not_convert is None:
         modules_to_not_convert = []
@@ -516,7 +515,6 @@ def sdnq_post_load_quant(
         return_device=return_device,
         modules_to_not_convert=modules_to_not_convert,
         modules_dtype_dict=modules_dtype_dict,
-        op=op,
     )
     model.quantization_config = SDNQConfig(
         weights_dtype=weights_dtype,
@@ -527,8 +525,10 @@ def sdnq_post_load_quant(
         quant_conv=quant_conv,
         use_quantized_matmul=use_quantized_matmul,
         use_quantized_matmul_conv=use_quantized_matmul_conv,
+        use_stochastic_rounding=use_stochastic_rounding,
         dequantize_fp32=dequantize_fp32,
         non_blocking=non_blocking,
+        add_skip_keys=add_skip_keys,
         quantization_device=quantization_device,
         return_device=return_device,
         modules_to_not_convert=modules_to_not_convert,
@@ -703,7 +703,11 @@ class SDNQQuantizer(DiffusersQuantizer, HfQuantizer):
             quantization_config_dict.pop("return_device", None)
             quantization_config_dict.pop("non_blocking", None)
             quantization_config_dict.pop("add_skip_keys", None)
+            quantization_config_dict.pop("use_static_quantization", None)
             quantization_config_dict.pop("use_stochastic_rounding", None)
+            quantization_config_dict.pop("quantized_matmul_dtype", None)
+            quantization_config_dict.pop("use_grad_ckpt", None)
+            quantization_config_dict.pop("is_training", None)
             with init_empty_weights():
                 model = sdnq_post_load_quant(model, add_skip_keys=False, **quantization_config_dict)
 
@@ -725,6 +729,17 @@ class SDNQQuantizer(DiffusersQuantizer, HfQuantizer):
         model.quantization_config = self.quantization_config
 
     def _process_model_after_weight_loading(self, model, **kwargs): # pylint: disable=unused-argument
+        if self.quantization_config.is_training:
+            from .training import convert_sdnq_model_to_training
+            model = convert_sdnq_model_to_training(
+                model,
+                dtype=self.torch_dtype,
+                quantized_matmul_dtype=self.quantization_config.quantized_matmul_dtype,
+                use_grad_ckpt=self.quantization_config.use_grad_ckpt,
+                use_quantized_matmul=self.quantization_config.use_quantized_matmul,
+                use_stochastic_rounding=self.quantization_config.use_stochastic_rounding,
+                dequantize_fp32=self.quantization_config.dequantize_fp32,
+            )
         if shared.opts.diffusers_offload_mode != "none":
             try:
                 model = model.to(device=devices.cpu)
@@ -743,30 +758,14 @@ class SDNQQuantizer(DiffusersQuantizer, HfQuantizer):
         return self.get_accelerator_warm_up_factor()
 
     def _dequantize(self, model):
-        model = dequantize_sdnq_model(model)
-        if hasattr(model, "quantization_method"):
-            del model.quantization_method
-        if hasattr(model, "quantization_config"):
-            del model.quantization_config
-        if hasattr(model, "config"):
-            try:
-                if hasattr(model.config, "quantization_config"):
-                    del model.config.quantization_config
-            except Exception:
-                pass
-            try:
-                if hasattr(model.config, "pop"):
-                    model.config.pop("quantization_config", None)
-            except Exception:
-                pass
-        return model
+        return dequantize_sdnq_model(model)
 
     def is_serializable(self, *args, **kwargs) -> bool:  # pylint: disable=unused-argument, invalid-overridden-method
-        return True
+        return not self.quantization_config.is_training
 
     @property
     def is_trainable(self):
-        return False
+        return self.quantization_config.is_training
 
     @property
     def is_compileable(self):
@@ -818,13 +817,16 @@ class SDNQConfig(QuantizationConfigMixin):
     def __init__( # pylint: disable=super-init-not-called
         self,
         weights_dtype: str = "int8",
+        quantized_matmul_dtype: str = "int8",
         group_size: int = 0,
         svd_rank: int = 32,
         svd_steps: int = 8,
         use_svd: bool = False,
+        use_grad_ckpt: bool = True,
         quant_conv: bool = False,
         use_quantized_matmul: bool = False,
         use_quantized_matmul_conv: bool = False,
+        use_static_quantization: bool = True,
         use_stochastic_rounding: bool = False,
         dequantize_fp32: bool = False,
         non_blocking: bool = False,
@@ -833,17 +835,25 @@ class SDNQConfig(QuantizationConfigMixin):
         return_device: Optional[torch.device] = None,
         modules_to_not_convert: Optional[List[str]] = None,
         modules_dtype_dict: Optional[Dict[str, List[str]]] = None,
+        is_training: bool = False,
         **kwargs, # pylint: disable=unused-argument
     ):
         self.weights_dtype = weights_dtype
-        self.quant_method = QuantizationMethod.SDNQ
+        self.quantized_matmul_dtype = quantized_matmul_dtype
+        self.is_training = is_training
+        if self.is_training:
+            self.quant_method = QuantizationMethod.SDNQ_TRAINING
+        else:
+            self.quant_method = QuantizationMethod.SDNQ
         self.group_size = group_size
         self.svd_rank = svd_rank
         self.svd_steps = svd_steps
         self.use_svd = use_svd
+        self.use_grad_ckpt = use_grad_ckpt
         self.quant_conv = quant_conv
         self.use_quantized_matmul = use_quantized_matmul
         self.use_quantized_matmul_conv = use_quantized_matmul_conv
+        self.use_static_quantization = use_static_quantization
         self.use_stochastic_rounding = use_stochastic_rounding
         self.dequantize_fp32 = dequantize_fp32
         self.non_blocking = non_blocking