Public/automatic
1
0
Fork 0
mirror of https://github.com/vladmandic/automatic synced 2026-09-19 09:14:35 +02:00
Code Issues Packages Projects Releases Wiki Activity

NNCF more optimizations

Browse Source
This commit is contained in:
Disty0
2025-05-07 21:50:31 +03:00
parent a57c7087b8
commit b6d2aa7fd8
1 changed files with 31 additions and 45 deletions
Download Patch File Download Diff File Expand all files Collapse all files
modules/model_quant_nncf.py
+31 -45
View File
@@ -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: