Make NNCF INT4 quant run 75% faster and don't force fp32 decompress

installer.py

@@ -1163,7 +1163,7 @@ def install_optional():
     install('albumentations==1.4.3', ignore=True)
     install('pydantic==1.10.21', ignore=True)
     reload('pydantic', '1.10.21')
-    install('nncf==2.16.0', ignore=True) # requires older pandas
+    install('nncf==2.16.0', ignore=True)
     install('gguf', ignore=True)
     install('av', ignore=True)
     try:

modules/model_quant_nncf.py

@@ -42,7 +42,7 @@ class QuantizationMethod(str, Enum):
     NNCF = "nncf"
 
 
-# de-abstracted and modified slightly from the actual quant functions of nncf 2.16.0:
+# de-abstracted and modified from the actual quant functions of nncf 2.16.0:
 def nncf_compress_layer(layer, num_bits, is_asym_mode, torch_dtype=None, quant_conv=False, param_name=None):
     if layer.__class__.__name__ in allowed_types:
         if torch_dtype is None:
@@ -77,7 +77,7 @@ def nncf_compress_layer(layer, num_bits, is_asym_mode, torch_dtype=None, quant_c
 
             scale = torch.where(torch.abs(scale) < eps, eps, scale)
             zero_point = level_low - torch.round(min_values / scale)
-            zero_point = torch.clip(zero_point.to(dtype=torch.int32), level_low, level_high)
+            zero_point = torch.clip(zero_point.to(dtype=torch.int32), level_low, level_high).to(dtype=torch.float32)
         else:
             factor = 2 ** (num_bits - 1)
 
@@ -96,8 +96,12 @@ def nncf_compress_layer(layer, num_bits, is_asym_mode, torch_dtype=None, quant_c
         level_high = 2**num_bits - 1 if is_asym_mode else 2 ** (num_bits - 1) - 1
 
         compressed_weight = layer.weight.data / scale
+        if not shared.opts.nncf_decompress_fp32:
+           scale = scale.to(torch_dtype)
+
         if zero_point is not None:
-            compressed_weight += zero_point.to(dtype=layer.weight.dtype)
+            compressed_weight += zero_point
+            zero_point = zero_point.to(scale.dtype)
 
         compressed_weight = torch.round(compressed_weight)
         compressed_weight = torch.clip(compressed_weight, level_low, level_high).to(dtype)
@@ -359,7 +363,6 @@ class NNCF_T5DenseGatedActDense(torch.nn.Module): # forward can't find what self
 
 def decompress_asymmetric(input: torch.Tensor, scale: torch.Tensor, zero_point: torch.Tensor) -> torch.Tensor:
     input = input.to(dtype=scale.dtype)
-    zero_point = zero_point.to(dtype=scale.dtype)
     decompressed_input = (input - zero_point) * scale
     return decompressed_input
 
@@ -401,7 +404,7 @@ class INT8AsymmetricWeightsDecompressor(torch.nn.Module):
     def __init__(self, scale: torch.Tensor, zero_point: torch.Tensor, result_dtype: Optional[torch.dtype] = None):
         super().__init__()
         self.scale = scale
-        self.zero_point = self.pack_weight(zero_point)
+        self.zero_point = zero_point
         self.result_dtype = result_dtype
 
     @property
@@ -453,6 +456,7 @@ class INT8SymmetricWeightsDecompressor(torch.nn.Module):
     def forward(self, x, *args, return_decompressed_only=False):
         result = decompress_symmetric(x.weight, self.scale)
         result = result.to(dtype=self.result_dtype)
+
         if return_decompressed_only:
             return result
         else:
@@ -470,10 +474,7 @@ class INT4AsymmetricWeightsDecompressor(torch.nn.Module):
     ):
         super().__init__()
         self.scale = scale
-
-        self.zero_point_shape = zero_point.shape
-        self.zero_point = self.pack_weight(zero_point)
-
+        self.zero_point = zero_point
         self.compressed_weight_shape = compressed_weight_shape
         self.result_shape = result_shape
         self.result_dtype = result_dtype
@@ -496,12 +497,10 @@ class INT4AsymmetricWeightsDecompressor(torch.nn.Module):
         result = unpack_uint4(x.weight)
         result = result.reshape(self.compressed_weight_shape)
 
-        zero_point = unpack_uint4(self.zero_point)
-        zero_point = zero_point.reshape(self.zero_point_shape)
-
-        result = decompress_asymmetric(result, self.scale, zero_point)
+        result = decompress_asymmetric(result, self.scale, self.zero_point)
         result = result.reshape(self.result_shape) if self.result_shape is not None else result
         result = result.to(dtype=self.result_dtype)
+
         if return_decompressed_only:
             return result
         else:
@@ -547,6 +546,7 @@ class INT4SymmetricWeightsDecompressor(torch.nn.Module):
         result = decompress_symmetric(result, self.scale)
         result = result.reshape(self.result_shape) if self.result_shape is not None else result
         result = result.to(dtype=self.result_dtype)
+
         if return_decompressed_only:
             return result
         else:

modules/shared.py

@@ -532,7 +532,7 @@ options_templates.update(options_section(('quantization', "Quantization Settings
     "optimum_quanto_weights": OptionInfo([], "Quantization enabled", gr.CheckboxGroup, {"choices": ["Model", "Transformer", "VAE", "TE", "Video", "LLM", "ControlNet"], "visible": native}),
     "optimum_quanto_weights_type": OptionInfo("qint8", "Quantization weights type", gr.Dropdown, {"choices": ['qint8', 'qfloat8_e4m3fn', 'qfloat8_e5m2', 'qint4', 'qint2'], "visible": native}),
     "optimum_quanto_activations_type": OptionInfo("none", "Quantization activations type ", gr.Dropdown, {"choices": ['none', 'qint8', 'qfloat8_e4m3fn', 'qfloat8_e5m2'], "visible": native}),
-    "optimum_quanto_shuffle_weights": OptionInfo(False, "Shuffle weights", gr.Checkbox, {"visible": native}),
+    "optimum_quanto_shuffle_weights": OptionInfo(False, "Shuffle weights in post mode", gr.Checkbox, {"visible": native}),
 
     "torchao_sep": OptionInfo("<h2>TorchAO</h2>", "", gr.HTML),
     "torchao_quantization": OptionInfo([], "Quantization enabled", gr.CheckboxGroup, {"choices": ["Model", "Transformer", "VAE", "TE", "Video", "LLM"], "visible": native}),
@@ -546,9 +546,10 @@ options_templates.update(options_section(('quantization', "Quantization Settings
     "nncf_compress_weights_raito": OptionInfo(0, "Compress ratio", gr.Slider, {"minimum": 0, "maximum": 1, "step": 0.01, "visible": cmd_opts.use_openvino}),
     "nncf_compress_weights_group_size": OptionInfo(0, "Group size", gr.Slider, {"minimum": -1, "maximum": 512, "step": 1, "visible": cmd_opts.use_openvino}),
     "nncf_quantize": OptionInfo([], "OpenVINO enabled", gr.CheckboxGroup, {"choices": ["Model", "VAE", "TE"], "visible": cmd_opts.use_openvino}),
-    "nncf_quantize_mode": OptionInfo("INT8", "OpenVINO mode", gr.Dropdown, {"choices": ['INT8', 'FP8_E4M3', 'FP8_E5M2'], "visible": cmd_opts.use_openvino}),
+    "nncf_quantize_mode": OptionInfo("INT8", "OpenVINO activations mode", gr.Dropdown, {"choices": ['INT8', 'FP8_E4M3', 'FP8_E5M2'], "visible": cmd_opts.use_openvino}),
     "nncf_quantize_conv_layers": OptionInfo(False, "Quantize the convolutional layers", gr.Checkbox, {"visible": native}),
-    "nncf_quantize_shuffle_weights": OptionInfo(False, "Shuffle weights", gr.Checkbox, {"visible": native}),
+    "nncf_decompress_fp32": OptionInfo(False, "Decompress using full precision", gr.Checkbox, {"visible": native}),
+    "nncf_quantize_shuffle_weights": OptionInfo(False, "Shuffle weights in post mode", gr.Checkbox, {"visible": native}),
 
     "layerwise_quantization_sep": OptionInfo("<h2>Layerwise Casting</h2>", "", gr.HTML),
     "layerwise_quantization": OptionInfo([], "Layerwise casting enabled", gr.CheckboxGroup, {"choices": ["Model", "Transformer", "TE"], "visible": native}),