From b6d2aa7fd87e7527aa147e5a75630b1b9233dad0 Mon Sep 17 00:00:00 2001 From: Disty0 Date: Wed, 7 May 2025 21:50:31 +0300 Subject: [PATCH] NNCF more optimizations --- modules/model_quant_nncf.py | 76 +++++++++++++++---------------------- 1 file changed, 31 insertions(+), 45 deletions(-) diff --git a/modules/model_quant_nncf.py b/modules/model_quant_nncf.py index d85e53ea5..d40d57487 100644 --- a/modules/model_quant_nncf.py +++ b/modules/model_quant_nncf.py @@ -1,7 +1,8 @@ -from typing import Any, Dict, List, Optional, Tuple, Union +from typing import Any, Dict, List, Optional, Union from dataclasses import dataclass from enum import Enum +import os import torch from diffusers.quantizers.base import DiffusersQuantizer from diffusers.quantizers.quantization_config import QuantizationConfigMixin @@ -13,6 +14,8 @@ from accelerate.utils import CustomDtype from modules import devices, shared +debug = os.environ.get('SD_QUANT_DEBUG', None) is not None + torch_dtype_dict = { "int8": torch.int8, "uint8": torch.uint8, @@ -112,27 +115,25 @@ def nncf_compress_layer(layer, num_bits, is_asym_mode, torch_dtype=None, quant_c scale=scale.data, zero_point=zero_point.data, compressed_weight_shape=compressed_weight.shape, - result_shape=layer.weight.shape, - result_dtype=torch_dtype + result_dtype=torch_dtype, ) else: decompressor = INT4SymmetricWeightsDecompressor( scale=scale.data, compressed_weight_shape=compressed_weight.shape, - result_shape=layer.weight.shape, - result_dtype=torch_dtype + result_dtype=torch_dtype, ) else: if is_asym_mode: decompressor = INT8AsymmetricWeightsDecompressor( scale=scale.data, zero_point=zero_point.data, - result_dtype=torch_dtype + result_dtype=torch_dtype, ) else: decompressor = INT8SymmetricWeightsDecompressor( scale=scale.data, - result_dtype=torch_dtype + result_dtype=torch_dtype, ) compressed_weight = decompressor.pack_weight(compressed_weight) @@ -377,15 +378,14 @@ def unpack_uint4(packed_tensor: torch.Tensor) -> torch.Tensor: return torch.stack((torch.bitwise_and(packed_tensor, 15), torch.bitwise_right_shift(packed_tensor, 4)), dim=-1) -def unpack_int4(packed_tensor: torch.Tensor) -> torch.Tensor: +def unpack_int4(packed_tensor: torch.Tensor, dtype: Optional[torch.dtype] = torch.int8) -> torch.Tensor: t = unpack_uint4(packed_tensor) - return t.to(dtype=torch.int8) - 8 + return t.to(dtype=dtype) - 8 def pack_uint4(tensor: torch.Tensor) -> torch.Tensor: if tensor.dtype != torch.uint8: - msg = f"Invalid tensor dtype {tensor.type}. torch.uint8 type is supported." - raise RuntimeError(msg) + raise RuntimeError(f"Invalid tensor dtype {tensor.type}. torch.uint8 type is supported.") packed_tensor = tensor.contiguous() packed_tensor = packed_tensor.reshape(-1, 2) packed_tensor = torch.bitwise_and(packed_tensor[..., ::2], 15) | packed_tensor[..., 1::2] << 4 @@ -394,14 +394,13 @@ def pack_uint4(tensor: torch.Tensor) -> torch.Tensor: def pack_int4(tensor: torch.Tensor) -> torch.Tensor: if tensor.dtype != torch.int8: - msg = f"Invalid tensor dtype {tensor.type}. torch.int8 type is supported." - raise RuntimeError(msg) + raise RuntimeError(f"Invalid tensor dtype {tensor.type}. torch.int8 type is supported.") tensor = tensor + 8 return pack_uint4(tensor.to(dtype=torch.uint8)) class INT8AsymmetricWeightsDecompressor(torch.nn.Module): - def __init__(self, scale: torch.Tensor, zero_point: torch.Tensor, result_dtype: Optional[torch.dtype] = None): + def __init__(self, scale: torch.Tensor, zero_point: torch.Tensor, result_dtype: torch.dtype): super().__init__() self.scale = scale self.zero_point = zero_point @@ -416,12 +415,9 @@ class INT8AsymmetricWeightsDecompressor(torch.nn.Module): return "asymmetric" def pack_weight(self, weight: torch.Tensor) -> torch.Tensor: - if torch.is_floating_point(weight): - msg = f"Invalid weight dtype {weight.type}. Integer types are supported." - raise ValueError(msg) - if torch.any((weight < 0) | (weight > 255)): - msg = "Weight values are not in [0, 255]." - raise ValueError(msg) + if debug: + if torch.any((weight < 0) | (weight > 255)): + raise ValueError("Weight values are not in [0, 255].") return weight.to(dtype=torch.uint8) def forward(self, x, *args, return_decompressed_only=False): @@ -434,7 +430,7 @@ class INT8AsymmetricWeightsDecompressor(torch.nn.Module): class INT8SymmetricWeightsDecompressor(torch.nn.Module): - def __init__(self, scale: torch.Tensor, result_dtype: Optional[torch.dtype] = None): + def __init__(self, scale: torch.Tensor, result_dtype: torch.dtype): super().__init__() self.scale = scale self.result_dtype = result_dtype @@ -448,9 +444,9 @@ class INT8SymmetricWeightsDecompressor(torch.nn.Module): return "symmetric" def pack_weight(self, weight: torch.Tensor) -> torch.Tensor: - if torch.any((weight < -128) | (weight > 127)): - msg = "Weight values are not in [-128, 127]." - raise ValueError(msg) + if debug: + if torch.any((weight < -128) | (weight > 127)): + raise ValueError("Weight values are not in [-128, 127].") return weight.to(dtype=torch.int8) def forward(self, x, *args, return_decompressed_only=False): @@ -468,15 +464,13 @@ class INT4AsymmetricWeightsDecompressor(torch.nn.Module): self, scale: torch.Tensor, zero_point: torch.Tensor, - compressed_weight_shape: Tuple[int, ...], - result_shape: Optional[Tuple[int, ...]] = None, - result_dtype: Optional[torch.dtype] = None, + compressed_weight_shape: torch.Size, + result_dtype: torch.dtype, ): super().__init__() self.scale = scale self.zero_point = zero_point self.compressed_weight_shape = compressed_weight_shape - self.result_shape = result_shape self.result_dtype = result_dtype @property @@ -488,9 +482,9 @@ class INT4AsymmetricWeightsDecompressor(torch.nn.Module): return "asymmetric" def pack_weight(self, weight: torch.Tensor) -> torch.Tensor: - if torch.any((weight < 0) | (weight > 15)): - msg = "Weight values are not in [0, 15]." - raise ValueError(msg) + if debug: + if torch.any((weight < 0) | (weight > 15)): + raise ValueError("Weight values are not in [0, 15].") return pack_uint4(weight.to(dtype=torch.uint8)) def forward(self, x, *args, return_decompressed_only=False): @@ -498,7 +492,6 @@ class INT4AsymmetricWeightsDecompressor(torch.nn.Module): result = result.reshape(self.compressed_weight_shape) 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: @@ -511,15 +504,12 @@ class INT4SymmetricWeightsDecompressor(torch.nn.Module): def __init__( self, scale: torch.Tensor, - compressed_weight_shape: Tuple[int, ...], - result_shape: Optional[Tuple[int, ...]] = None, - result_dtype: Optional[torch.dtype] = None, + compressed_weight_shape: torch.Size, + result_dtype: torch.dtype, ): super().__init__() self.scale = scale - self.compressed_weight_shape = compressed_weight_shape - self.result_shape = result_shape self.result_dtype = result_dtype @property @@ -531,20 +521,16 @@ class INT4SymmetricWeightsDecompressor(torch.nn.Module): return "symmetric" def pack_weight(self, weight: torch.Tensor) -> torch.Tensor: - if torch.is_floating_point(weight): - msg = f"Invalid weight dtype {weight.type}. Integer types are supported." - raise ValueError(msg) - if torch.any((weight < -8) | (weight > 7)): - msg = "Tensor values are not in [-8, 7]." - raise ValueError(msg) + if debug: + if torch.any((weight < -8) | (weight > 7)): + raise ValueError("Tensor values are not in [-8, 7].") return pack_int4(weight.to(dtype=torch.int8)) def forward(self, x, *arg, return_decompressed_only=False): - result = unpack_int4(x.weight) + result = unpack_int4(x.weight, dtype=self.scale.dtype) result = result.reshape(self.compressed_weight_shape) 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: