mirror of
https://github.com/vladmandic/automatic
synced 2026-09-20 01:31:13 +02:00
SDNQ optimize input quantization and use the word quantize instead of compress
This commit is contained in:
Disty0
+13
-13
@@ -45,8 +45,8 @@ def network_backup_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.n
|
||||
self.network_weights_backup = True
|
||||
else:
|
||||
self.network_weights_backup = weight.clone().to(devices.cpu)
|
||||
if hasattr(self, "sdnq_decompressor"):
|
||||
self.sdnq_decompressor_backup = self.sdnq_decompressor.to(devices.cpu)
|
||||
if hasattr(self, "sdnq_dequantizer"):
|
||||
self.sdnq_dequantizer_backup = self.sdnq_dequantizer.to(devices.cpu)
|
||||
|
||||
if bias_backup is None:
|
||||
if getattr(self, 'bias', None) is not None:
|
||||
@@ -79,8 +79,8 @@ def network_calc_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn.
|
||||
continue
|
||||
try:
|
||||
t0 = time.time()
|
||||
if hasattr(self, "sdnq_decompressor"):
|
||||
weight = self.sdnq_decompressor.to(devices.device)(self.weight.to(devices.device), skip_quantized_matmul=self.sdnq_decompressor.use_quantized_matmul)
|
||||
if hasattr(self, "sdnq_dequantizer"):
|
||||
weight = self.sdnq_dequantizer.to(devices.device)(self.weight.to(devices.device), skip_quantized_matmul=self.sdnq_dequantizer.use_quantized_matmul)
|
||||
else:
|
||||
weight = self.weight.to(devices.device) # must perform calc on gpu due to performance
|
||||
updown, ex_bias = module.calc_updown(weight)
|
||||
@@ -136,20 +136,20 @@ def network_add_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn.G
|
||||
# weight._quantize(devices.device) / weight.to(device=device)
|
||||
except Exception as e:
|
||||
shared.log.error(f'Network load: type=LoRA quant=bnb cls={self.__class__.__name__} type={self.quant_type} blocksize={self.blocksize} state={vars(self.quant_state)} weight={self.weight} bias={lora_weights} {e}')
|
||||
elif not bias and hasattr(self, "sdnq_decompressor"):
|
||||
elif not bias and hasattr(self, "sdnq_dequantizer"):
|
||||
try:
|
||||
from modules.sdnq import sdnq_quantize_layer
|
||||
if hasattr(self, "sdnq_decompressor_backup"):
|
||||
sdnq_decompressor = self.sdnq_decompressor_backup.to(devices.device)
|
||||
if hasattr(self, "sdnq_dequantizer_backup"):
|
||||
sdnq_dequantizer = self.sdnq_dequantizer_backup.to(devices.device)
|
||||
else:
|
||||
sdnq_decompressor = self.sdnq_decompressor.to(devices.device)
|
||||
dequant_weight = sdnq_decompressor(model_weights.to(devices.device), skip_quantized_matmul=sdnq_decompressor.use_quantized_matmul)
|
||||
sdnq_dequantizer = self.sdnq_dequantizer.to(devices.device)
|
||||
dequant_weight = sdnq_dequantizer(model_weights.to(devices.device), skip_quantized_matmul=sdnq_dequantizer.use_quantized_matmul)
|
||||
new_weight = dequant_weight.to(devices.device, dtype=torch.float32) + lora_weights.to(devices.device, dtype=torch.float32)
|
||||
self.weight = torch.nn.Parameter(new_weight, requires_grad=False)
|
||||
self.sdnq_decompressor = None
|
||||
self.sdnq_dequantizer = None
|
||||
self = sdnq_quantize_layer(
|
||||
self,
|
||||
sdnq_decompressor.weights_dtype,
|
||||
sdnq_dequantizer.weights_dtype,
|
||||
torch_dtype=devices.dtype,
|
||||
group_size=shared.opts.sdnq_quantize_weights_group_size,
|
||||
quant_conv=shared.opts.sdnq_quantize_conv_layers,
|
||||
@@ -223,8 +223,8 @@ def network_apply_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn
|
||||
network_add_weights(self, model_weights=weights_backup, lora_weights=updown, deactivate=deactivate, device=device, bias=False)
|
||||
else:
|
||||
self.weight = torch.nn.Parameter(weights_backup.to(device), requires_grad=False)
|
||||
if hasattr(self, "sdnq_decompressor_backup"):
|
||||
self.sdnq_decompressor = self.sdnq_decompressor_backup.to(device)
|
||||
if hasattr(self, "sdnq_dequantizer_backup"):
|
||||
self.sdnq_dequantizer = self.sdnq_dequantizer_backup.to(device)
|
||||
|
||||
if bias_backup is not None:
|
||||
self.bias = None
|
||||
|
||||
+19
-22
@@ -10,7 +10,7 @@ from diffusers.utils import get_module_from_name
|
||||
from modules import devices, shared
|
||||
|
||||
from .common import dtype_dict, use_tensorwise_fp8_matmul, quantized_matmul_dtypes, allowed_types, conv_types, conv_transpose_types
|
||||
from .decompressor import decompressor_dict
|
||||
from .dequantizer import dequantizer_dict
|
||||
from .forward import get_forward_func
|
||||
|
||||
|
||||
@@ -123,14 +123,8 @@ def sdnq_quantize_layer(layer, weights_dtype="int8", torch_dtype=None, group_siz
|
||||
else:
|
||||
layer.weight.data = layer.weight.to(dtype=torch.float32)
|
||||
|
||||
if dtype_dict[weights_dtype]["is_unsigned"]:
|
||||
scale, zero_point = get_scale_asymmetric(layer.weight, reduction_axes, weights_dtype)
|
||||
else:
|
||||
scale = get_scale_symmetric(layer.weight, reduction_axes, weights_dtype)
|
||||
zero_point = None
|
||||
layer.weight.data = quantize_weight(layer.weight, scale, zero_point, weights_dtype)
|
||||
|
||||
if not shared.opts.sdnq_decompress_fp32 and not (use_quantized_matmul and not dtype_dict[weights_dtype]["is_integer"] and not use_tensorwise_fp8_matmul):
|
||||
layer.weight.data, scale, zero_point = quantize_weight(layer.weight, reduction_axes, weights_dtype)
|
||||
if not shared.opts.sdnq_dequantize_fp32 and not (use_quantized_matmul and not dtype_dict[weights_dtype]["is_integer"] and not use_tensorwise_fp8_matmul):
|
||||
scale = scale.to(torch_dtype)
|
||||
if zero_point is not None:
|
||||
zero_point = zero_point.to(torch_dtype)
|
||||
@@ -146,17 +140,17 @@ def sdnq_quantize_layer(layer, weights_dtype="int8", torch_dtype=None, group_siz
|
||||
if not use_tensorwise_fp8_matmul:
|
||||
scale = scale.to(torch.float32)
|
||||
|
||||
layer.sdnq_decompressor = decompressor_dict[weights_dtype](
|
||||
layer.sdnq_dequantizer = dequantizer_dict[weights_dtype](
|
||||
scale=scale,
|
||||
zero_point=zero_point,
|
||||
compressed_weight_shape=layer.weight.shape,
|
||||
quantized_weight_shape=layer.weight.shape,
|
||||
result_dtype=torch_dtype,
|
||||
result_shape=result_shape,
|
||||
weights_dtype=weights_dtype,
|
||||
use_quantized_matmul=use_quantized_matmul,
|
||||
)
|
||||
layer.weight.data = layer.sdnq_decompressor.pack_weight(layer.weight).to(return_device)
|
||||
layer.sdnq_decompressor = layer.sdnq_decompressor.to(return_device)
|
||||
layer.weight.data = layer.sdnq_dequantizer.pack_weight(layer.weight).to(return_device)
|
||||
layer.sdnq_dequantizer = layer.sdnq_dequantizer.to(return_device)
|
||||
|
||||
layer.forward = get_forward_func(layer_class_name, use_quantized_matmul, dtype_dict[weights_dtype]["is_integer"], use_tensorwise_fp8_matmul)
|
||||
layer.forward = layer.forward.__get__(layer, layer.__class__)
|
||||
@@ -193,7 +187,7 @@ def apply_sdnq_to_module(model, weights_dtype="int8", torch_dtype=None, group_si
|
||||
return model
|
||||
|
||||
|
||||
def get_scale_asymmetric(weight: torch.FloatTensor, reduction_axes: List[int], weights_dtype: str) -> Tuple[torch.FloatTensor, torch.FloatTensor]:
|
||||
def get_scale_asymmetric(weight: torch.FloatTensor, reduction_axes: Union[int, List[int]], weights_dtype: str) -> Tuple[torch.FloatTensor, torch.FloatTensor]:
|
||||
zero_point = torch.amin(weight, dim=reduction_axes, keepdims=True)
|
||||
scale = torch.amax(weight, dim=reduction_axes, keepdims=True).sub_(zero_point).div_(dtype_dict[weights_dtype]["max"] - dtype_dict[weights_dtype]["min"])
|
||||
eps = torch.finfo(scale.dtype).eps # prevent divison by 0
|
||||
@@ -203,22 +197,25 @@ def get_scale_asymmetric(weight: torch.FloatTensor, reduction_axes: List[int], w
|
||||
return scale, zero_point
|
||||
|
||||
|
||||
def get_scale_symmetric(weight: torch.FloatTensor, reduction_axes: List[int], weights_dtype: str) -> torch.FloatTensor:
|
||||
def get_scale_symmetric(weight: torch.FloatTensor, reduction_axes: Union[int, List[int]], weights_dtype: str) -> torch.FloatTensor:
|
||||
scale = torch.amax(weight.abs(), dim=reduction_axes, keepdims=True).div_(dtype_dict[weights_dtype]["max"])
|
||||
eps = torch.finfo(scale.dtype).eps # prevent divison by 0
|
||||
scale = torch.where(torch.abs(scale) < eps, eps, scale)
|
||||
return scale
|
||||
|
||||
|
||||
def quantize_weight(weight: torch.FloatTensor, scale: torch.FloatTensor, zero_point: torch.FloatTensor, weights_dtype: str) -> torch.Tensor:
|
||||
if zero_point is not None:
|
||||
compressed_weight = torch.sub(weight, zero_point).div_(scale)
|
||||
def quantize_weight(weight: torch.FloatTensor, reduction_axes: Union[int, List[int]], weights_dtype: str):
|
||||
if dtype_dict[weights_dtype]["is_unsigned"]:
|
||||
scale, zero_point = get_scale_asymmetric(weight, reduction_axes, weights_dtype)
|
||||
quantized_weight = torch.sub(weight, zero_point).div_(scale)
|
||||
else:
|
||||
compressed_weight = torch.div(weight, scale)
|
||||
scale = get_scale_symmetric(weight, reduction_axes, weights_dtype)
|
||||
quantized_weight = torch.div(weight, scale)
|
||||
zero_point = None
|
||||
if dtype_dict[weights_dtype]["is_integer"]:
|
||||
compressed_weight.round_()
|
||||
compressed_weight = compressed_weight.clamp_(dtype_dict[weights_dtype]["min"], dtype_dict[weights_dtype]["max"]).to(dtype_dict[weights_dtype]["torch_dtype"])
|
||||
return compressed_weight
|
||||
quantized_weight.round_()
|
||||
quantized_weight = quantized_weight.clamp_(dtype_dict[weights_dtype]["min"], dtype_dict[weights_dtype]["max"]).to(dtype_dict[weights_dtype]["torch_dtype"])
|
||||
return quantized_weight, scale, zero_point
|
||||
|
||||
|
||||
class QuantizationMethod(str, Enum):
|
||||
|
||||
@@ -7,14 +7,14 @@ from .common import dtype_dict
|
||||
from .packed_int import pack_int_symetric, unpack_int_symetric, packed_int_function_dict
|
||||
|
||||
|
||||
def decompress_asymmetric(input: torch.ByteTensor, scale: torch.FloatTensor, zero_point: torch.FloatTensor, dtype: torch.dtype, result_shape: torch.Size) -> torch.FloatTensor:
|
||||
def dequantize_asymmetric(input: torch.ByteTensor, scale: torch.FloatTensor, zero_point: torch.FloatTensor, dtype: torch.dtype, result_shape: torch.Size) -> torch.FloatTensor:
|
||||
result = torch.addcmul(zero_point, input.to(dtype=scale.dtype), scale).to(dtype=dtype)
|
||||
if result_shape is not None:
|
||||
result = result.reshape(result_shape)
|
||||
return result
|
||||
|
||||
|
||||
def decompress_symmetric(input: torch.CharTensor, scale: torch.FloatTensor, dtype: torch.dtype, result_shape: torch.Size, skip_quantized_matmul: bool = False) -> torch.FloatTensor:
|
||||
def dequantize_symmetric(input: torch.CharTensor, scale: torch.FloatTensor, dtype: torch.dtype, result_shape: torch.Size, skip_quantized_matmul: bool = False) -> torch.FloatTensor:
|
||||
if skip_quantized_matmul:
|
||||
result = input.transpose(0,1).to(dtype=scale.dtype).mul_(scale.transpose(0,1)).to(dtype=dtype)
|
||||
else:
|
||||
@@ -24,18 +24,18 @@ def decompress_symmetric(input: torch.CharTensor, scale: torch.FloatTensor, dtyp
|
||||
return result
|
||||
|
||||
|
||||
def decompress_packed_int_asymmetric(input: torch.ByteTensor, scale: torch.FloatTensor, zero_point: torch.FloatTensor, shape: torch.Size, dtype: torch.dtype, result_shape: torch.Size, weights_dtype: str) -> torch.FloatTensor:
|
||||
return decompress_asymmetric(packed_int_function_dict[weights_dtype]["unpack"](input, shape), scale, zero_point, dtype, result_shape)
|
||||
def dequantize_packed_int_asymmetric(input: torch.ByteTensor, scale: torch.FloatTensor, zero_point: torch.FloatTensor, shape: torch.Size, dtype: torch.dtype, result_shape: torch.Size, weights_dtype: str) -> torch.FloatTensor:
|
||||
return dequantize_asymmetric(packed_int_function_dict[weights_dtype]["unpack"](input, shape), scale, zero_point, dtype, result_shape)
|
||||
|
||||
|
||||
def decompress_packed_int_symmetric(input: torch.ByteTensor, scale: torch.FloatTensor, shape: torch.Size, dtype: torch.dtype, result_shape: torch.Size, weights_dtype: str, skip_quantized_matmul: bool = False) -> torch.FloatTensor:
|
||||
def dequantize_packed_int_symmetric(input: torch.ByteTensor, scale: torch.FloatTensor, shape: torch.Size, dtype: torch.dtype, result_shape: torch.Size, weights_dtype: str, skip_quantized_matmul: bool = False) -> torch.FloatTensor:
|
||||
if skip_quantized_matmul:
|
||||
return decompress_symmetric(unpack_int_symetric(input, shape, weights_dtype, dtype=scale.dtype), scale.transpose(0,1), dtype, result_shape)
|
||||
return dequantize_symmetric(unpack_int_symetric(input, shape, weights_dtype, dtype=scale.dtype), scale.transpose(0,1), dtype, result_shape)
|
||||
else:
|
||||
return decompress_symmetric(unpack_int_symetric(input, shape, weights_dtype, dtype=scale.dtype), scale, dtype, result_shape)
|
||||
return dequantize_symmetric(unpack_int_symetric(input, shape, weights_dtype, dtype=scale.dtype), scale, dtype, result_shape)
|
||||
|
||||
|
||||
class AsymmetricWeightsDecompressor(torch.nn.Module):
|
||||
class AsymmetricWeightsDequantizer(torch.nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
scale: torch.Tensor,
|
||||
@@ -57,10 +57,10 @@ class AsymmetricWeightsDecompressor(torch.nn.Module):
|
||||
return weight.to(dtype=dtype_dict[self.weights_dtype]["torch_dtype"])
|
||||
|
||||
def forward(self, weight, **kwargs): # pylint: disable=unused-argument
|
||||
return decompress_asymmetric(weight, self.scale, self.zero_point, self.result_dtype, self.result_shape)
|
||||
return dequantize_asymmetric(weight, self.scale, self.zero_point, self.result_dtype, self.result_shape)
|
||||
|
||||
|
||||
class SymmetricWeightsDecompressor(torch.nn.Module):
|
||||
class SymmetricWeightsDequantizer(torch.nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
scale: torch.Tensor,
|
||||
@@ -81,15 +81,15 @@ class SymmetricWeightsDecompressor(torch.nn.Module):
|
||||
return weight.to(dtype=dtype_dict[self.weights_dtype]["torch_dtype"])
|
||||
|
||||
def forward(self, weight, skip_quantized_matmul=False, **kwargs): # pylint: disable=unused-argument
|
||||
return decompress_symmetric(weight, self.scale, self.result_dtype, self.result_shape, skip_quantized_matmul=skip_quantized_matmul)
|
||||
return dequantize_symmetric(weight, self.scale, self.result_dtype, self.result_shape, skip_quantized_matmul=skip_quantized_matmul)
|
||||
|
||||
|
||||
class PackedINTAsymmetricWeightsDecompressor(torch.nn.Module):
|
||||
class PackedINTAsymmetricWeightsDequantizer(torch.nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
scale: torch.Tensor,
|
||||
zero_point: torch.Tensor,
|
||||
compressed_weight_shape: torch.Size,
|
||||
quantized_weight_shape: torch.Size,
|
||||
result_dtype: torch.dtype,
|
||||
result_shape: torch.Size,
|
||||
weights_dtype: str,
|
||||
@@ -98,7 +98,7 @@ class PackedINTAsymmetricWeightsDecompressor(torch.nn.Module):
|
||||
super().__init__()
|
||||
self.weights_dtype = weights_dtype
|
||||
self.use_quantized_matmul = False
|
||||
self.compressed_weight_shape = compressed_weight_shape
|
||||
self.quantized_weight_shape = quantized_weight_shape
|
||||
self.result_dtype = result_dtype
|
||||
self.result_shape = result_shape
|
||||
self.register_buffer("scale", scale)
|
||||
@@ -108,14 +108,14 @@ class PackedINTAsymmetricWeightsDecompressor(torch.nn.Module):
|
||||
return packed_int_function_dict[self.weights_dtype]["pack"](weight.to(dtype=dtype_dict[self.weights_dtype]["torch_dtype"]))
|
||||
|
||||
def forward(self, weight, **kwargs): # pylint: disable=unused-argument
|
||||
return decompress_packed_int_asymmetric(weight, self.scale, self.zero_point, self.compressed_weight_shape, self.result_dtype, self.result_shape, self.weights_dtype)
|
||||
return dequantize_packed_int_asymmetric(weight, self.scale, self.zero_point, self.quantized_weight_shape, self.result_dtype, self.result_shape, self.weights_dtype)
|
||||
|
||||
|
||||
class PackedINTSymmetricWeightsDecompressor(torch.nn.Module):
|
||||
class PackedINTSymmetricWeightsDequantizer(torch.nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
scale: torch.Tensor,
|
||||
compressed_weight_shape: torch.Size,
|
||||
quantized_weight_shape: torch.Size,
|
||||
result_dtype: torch.dtype,
|
||||
result_shape: torch.Size,
|
||||
weights_dtype: str,
|
||||
@@ -125,7 +125,7 @@ class PackedINTSymmetricWeightsDecompressor(torch.nn.Module):
|
||||
super().__init__()
|
||||
self.weights_dtype = weights_dtype
|
||||
self.use_quantized_matmul = use_quantized_matmul
|
||||
self.compressed_weight_shape = compressed_weight_shape
|
||||
self.quantized_weight_shape = quantized_weight_shape
|
||||
self.result_dtype = result_dtype
|
||||
self.result_shape = result_shape
|
||||
self.register_buffer("scale", scale)
|
||||
@@ -134,39 +134,39 @@ class PackedINTSymmetricWeightsDecompressor(torch.nn.Module):
|
||||
return pack_int_symetric(weight, self.weights_dtype)
|
||||
|
||||
def forward(self, weight, skip_quantized_matmul=False, **kwargs): # pylint: disable=unused-argument
|
||||
return decompress_packed_int_symmetric(weight, self.scale, self.compressed_weight_shape, self.result_dtype, self.result_shape, self.weights_dtype, skip_quantized_matmul=skip_quantized_matmul)
|
||||
return dequantize_packed_int_symmetric(weight, self.scale, self.quantized_weight_shape, self.result_dtype, self.result_shape, self.weights_dtype, skip_quantized_matmul=skip_quantized_matmul)
|
||||
|
||||
|
||||
decompressor_dict = {
|
||||
"int8": SymmetricWeightsDecompressor,
|
||||
"int7": PackedINTSymmetricWeightsDecompressor,
|
||||
"int6": PackedINTSymmetricWeightsDecompressor,
|
||||
"int5": PackedINTSymmetricWeightsDecompressor,
|
||||
"int4": PackedINTSymmetricWeightsDecompressor,
|
||||
"int3": PackedINTSymmetricWeightsDecompressor,
|
||||
"int2": PackedINTSymmetricWeightsDecompressor,
|
||||
"uint8": AsymmetricWeightsDecompressor,
|
||||
"uint7": PackedINTAsymmetricWeightsDecompressor,
|
||||
"uint6": PackedINTAsymmetricWeightsDecompressor,
|
||||
"uint5": PackedINTAsymmetricWeightsDecompressor,
|
||||
"uint4": PackedINTAsymmetricWeightsDecompressor,
|
||||
"uint3": PackedINTAsymmetricWeightsDecompressor,
|
||||
"uint2": PackedINTAsymmetricWeightsDecompressor,
|
||||
"uint1": AsymmetricWeightsDecompressor,
|
||||
"bool": AsymmetricWeightsDecompressor,
|
||||
"float8_e4m3fn": SymmetricWeightsDecompressor,
|
||||
"float8_e4m3fnuz": SymmetricWeightsDecompressor,
|
||||
"float8_e5m2": SymmetricWeightsDecompressor,
|
||||
"float8_e5m2fnuz": SymmetricWeightsDecompressor,
|
||||
dequantizer_dict = {
|
||||
"int8": SymmetricWeightsDequantizer,
|
||||
"int7": PackedINTSymmetricWeightsDequantizer,
|
||||
"int6": PackedINTSymmetricWeightsDequantizer,
|
||||
"int5": PackedINTSymmetricWeightsDequantizer,
|
||||
"int4": PackedINTSymmetricWeightsDequantizer,
|
||||
"int3": PackedINTSymmetricWeightsDequantizer,
|
||||
"int2": PackedINTSymmetricWeightsDequantizer,
|
||||
"uint8": AsymmetricWeightsDequantizer,
|
||||
"uint7": PackedINTAsymmetricWeightsDequantizer,
|
||||
"uint6": PackedINTAsymmetricWeightsDequantizer,
|
||||
"uint5": PackedINTAsymmetricWeightsDequantizer,
|
||||
"uint4": PackedINTAsymmetricWeightsDequantizer,
|
||||
"uint3": PackedINTAsymmetricWeightsDequantizer,
|
||||
"uint2": PackedINTAsymmetricWeightsDequantizer,
|
||||
"uint1": AsymmetricWeightsDequantizer,
|
||||
"bool": AsymmetricWeightsDequantizer,
|
||||
"float8_e4m3fn": SymmetricWeightsDequantizer,
|
||||
"float8_e4m3fnuz": SymmetricWeightsDequantizer,
|
||||
"float8_e5m2": SymmetricWeightsDequantizer,
|
||||
"float8_e5m2fnuz": SymmetricWeightsDequantizer,
|
||||
}
|
||||
|
||||
|
||||
if shared.opts.sdnq_decompress_compile:
|
||||
if shared.opts.sdnq_dequantize_compile:
|
||||
try:
|
||||
torch._dynamo.config.cache_size_limit = max(8192, torch._dynamo.config.cache_size_limit)
|
||||
decompress_asymmetric = torch.compile(decompress_asymmetric, fullgraph=True)
|
||||
decompress_symmetric = torch.compile(decompress_symmetric, fullgraph=True)
|
||||
decompress_packed_int_asymmetric = torch.compile(decompress_packed_int_asymmetric, fullgraph=True)
|
||||
decompress_packed_int_symmetric = torch.compile(decompress_packed_int_symmetric, fullgraph=True)
|
||||
dequantize_asymmetric = torch.compile(dequantize_asymmetric, fullgraph=True)
|
||||
dequantize_symmetric = torch.compile(dequantize_symmetric, fullgraph=True)
|
||||
dequantize_packed_int_asymmetric = torch.compile(dequantize_packed_int_asymmetric, fullgraph=True)
|
||||
dequantize_packed_int_symmetric = torch.compile(dequantize_packed_int_symmetric, fullgraph=True)
|
||||
except Exception as e:
|
||||
shared.log.warning(f"Quantization: type=sdnq Decompress using torch.compile is not available: {e}")
|
||||
shared.log.warning(f"Quantization: type=sdnq Dequantize using torch.compile is not available: {e}")
|
||||
+39
-39
@@ -5,7 +5,7 @@ import torch
|
||||
from modules import shared
|
||||
|
||||
from .common import conv_types, conv_transpose_types
|
||||
from .decompressor import decompress_symmetric
|
||||
from .dequantizer import dequantize_symmetric
|
||||
from .packed_int import unpack_int_symetric
|
||||
|
||||
|
||||
@@ -43,7 +43,7 @@ def get_forward_func(layer_class_name: str, use_quantized_matmul: bool, is_integ
|
||||
|
||||
def quantize_fp8_matmul_input(input: torch.FloatTensor) -> Tuple[torch.Tensor, torch.FloatTensor]:
|
||||
input = input.flatten(0,-2).contiguous()
|
||||
input_scale = torch.div(input.abs().amax(dim=-1, keepdims=True), 448)
|
||||
input_scale = torch.amax(input.abs(), dim=-1, keepdims=True).div_(448)
|
||||
input = torch.div(input, input_scale).clamp_(-448, 448).to(torch.float8_e4m3fn)
|
||||
input_scale = input_scale.to(torch.float32)
|
||||
return input, input_scale
|
||||
@@ -51,7 +51,7 @@ def quantize_fp8_matmul_input(input: torch.FloatTensor) -> Tuple[torch.Tensor, t
|
||||
|
||||
def quantize_fp8_matmul_input_tensorwise(input: torch.FloatTensor, scale: torch.FloatTensor) -> Tuple[torch.Tensor, torch.FloatTensor]:
|
||||
input = input.flatten(0,-2).contiguous()
|
||||
input_scale = torch.div(input.abs().amax(dim=-1, keepdims=True), 448)
|
||||
input_scale = torch.amax(input.abs(), dim=-1, keepdims=True).div_(448)
|
||||
input = torch.div(input, input_scale).clamp_(-448, 448).to(torch.float8_e4m3fn)
|
||||
scale = torch.mul(input_scale, scale)
|
||||
if scale.dtype == torch.float16: # fp16 will overflow
|
||||
@@ -61,7 +61,7 @@ def quantize_fp8_matmul_input_tensorwise(input: torch.FloatTensor, scale: torch.
|
||||
|
||||
def quantize_int8_matmul_input(input: torch.FloatTensor, scale: torch.FloatTensor) -> Tuple[torch.CharTensor, torch.FloatTensor]:
|
||||
input = input.flatten(0,-2).contiguous()
|
||||
input_scale = torch.div(input.abs().amax(dim=-1, keepdims=True), 127)
|
||||
input_scale = torch.amax(input.abs(), dim=-1, keepdims=True).div_(127)
|
||||
input = torch.div(input, input_scale).round_().clamp_(-128, 127).to(torch.int8)
|
||||
scale = torch.mul(input_scale, scale)
|
||||
if scale.dtype == torch.float16: # fp16 will overflow
|
||||
@@ -94,7 +94,7 @@ def fp8_matmul_tensorwise(
|
||||
output_shape[-1] = weight.shape[-1]
|
||||
dummy_input_scale = torch.ones(1, device=input.device, dtype=torch.float32)
|
||||
input, scale = quantize_fp8_matmul_input_tensorwise(input, scale)
|
||||
result = decompress_symmetric(torch._scaled_mm(input, weight, scale_a=dummy_input_scale, scale_b=dummy_input_scale, bias=None, out_dtype=scale.dtype), scale, return_dtype, output_shape)
|
||||
result = dequantize_symmetric(torch._scaled_mm(input, weight, scale_a=dummy_input_scale, scale_b=dummy_input_scale, bias=None, out_dtype=scale.dtype), scale, return_dtype, output_shape)
|
||||
if bias is not None:
|
||||
result.add_(bias)
|
||||
return result
|
||||
@@ -105,16 +105,16 @@ def int8_matmul(
|
||||
weight: torch.Tensor,
|
||||
bias: torch.FloatTensor,
|
||||
scale: torch.FloatTensor,
|
||||
compressed_weight_shape: torch.Size,
|
||||
quantized_weight_shape: torch.Size,
|
||||
weights_dtype: str,
|
||||
) -> torch.FloatTensor:
|
||||
if compressed_weight_shape is not None:
|
||||
weight = unpack_int_symetric(weight, compressed_weight_shape, weights_dtype, dtype=torch.int8, transpose=True)
|
||||
if quantized_weight_shape is not None:
|
||||
weight = unpack_int_symetric(weight, quantized_weight_shape, weights_dtype, dtype=torch.int8, transpose=True)
|
||||
return_dtype = input.dtype
|
||||
output_shape = list(input.shape)
|
||||
output_shape[-1] = weight.shape[-1]
|
||||
input, scale = quantize_int8_matmul_input(input, scale)
|
||||
result = decompress_symmetric(torch._int_mm(input, weight), scale, return_dtype, output_shape)
|
||||
result = dequantize_symmetric(torch._int_mm(input, weight), scale, return_dtype, output_shape)
|
||||
if bias is not None:
|
||||
result.add_(bias)
|
||||
return result
|
||||
@@ -224,14 +224,14 @@ def conv_fp8_matmul_tensorwise(
|
||||
dummy_input_scale = torch.ones(1, device=input.device, dtype=torch.float32)
|
||||
|
||||
if groups == 1:
|
||||
result = decompress_symmetric(torch._scaled_mm(input, weight, scale_a=dummy_input_scale, scale_b=dummy_input_scale, bias=None, out_dtype=scale.dtype), scale, return_dtype, mm_output_shape)
|
||||
result = dequantize_symmetric(torch._scaled_mm(input, weight, scale_a=dummy_input_scale, scale_b=dummy_input_scale, bias=None, out_dtype=scale.dtype), scale, return_dtype, mm_output_shape)
|
||||
else:
|
||||
weight = weight.reshape(weight.shape[0], groups, weight.shape[1] // groups).transpose(0,1)
|
||||
input = input.reshape(input.shape[0], groups, input.shape[1] // groups).transpose(0,1)
|
||||
result = []
|
||||
for i in range(groups):
|
||||
result.append(torch._scaled_mm(input[i], weight[i], scale_a=dummy_input_scale, scale_b=dummy_input_scale, bias=None, out_dtype=scale.dtype))
|
||||
result = decompress_symmetric(torch.cat(result, dim=-1), scale, return_dtype, mm_output_shape)
|
||||
result = dequantize_symmetric(torch.cat(result, dim=-1), scale, return_dtype, mm_output_shape)
|
||||
if bias is not None:
|
||||
result.add_(bias)
|
||||
|
||||
@@ -250,7 +250,7 @@ def conv_int8_matmul(
|
||||
bias: torch.FloatTensor,
|
||||
scale: torch.FloatTensor,
|
||||
result_shape: torch.Size,
|
||||
compressed_weight_shape: torch.Size,
|
||||
quantized_weight_shape: torch.Size,
|
||||
weights_dtype: str,
|
||||
reversed_padding_repeated_twice: List[int],
|
||||
padding_mode: str, conv_type: int,
|
||||
@@ -260,18 +260,18 @@ def conv_int8_matmul(
|
||||
return_dtype = input.dtype
|
||||
input, mm_output_shape = process_conv_input(conv_type, input, reversed_padding_repeated_twice, padding_mode, result_shape, stride, padding, dilation)
|
||||
input, scale = quantize_int8_matmul_input(input, scale)
|
||||
if compressed_weight_shape is not None:
|
||||
weight = unpack_int_symetric(weight, compressed_weight_shape, weights_dtype, dtype=torch.int8, transpose=True)
|
||||
if quantized_weight_shape is not None:
|
||||
weight = unpack_int_symetric(weight, quantized_weight_shape, weights_dtype, dtype=torch.int8, transpose=True)
|
||||
|
||||
if groups == 1:
|
||||
result = decompress_symmetric(torch._int_mm(input, weight), scale, return_dtype, mm_output_shape)
|
||||
result = dequantize_symmetric(torch._int_mm(input, weight), scale, return_dtype, mm_output_shape)
|
||||
else:
|
||||
weight = weight.reshape(weight.shape[0], groups, weight.shape[1] // groups).transpose(0,1)
|
||||
input = input.reshape(input.shape[0], groups, input.shape[1] // groups).transpose(0,1)
|
||||
result = []
|
||||
for i in range(groups):
|
||||
result.append(torch._int_mm(input[i], weight[i]))
|
||||
result = decompress_symmetric(torch.cat(result, dim=-1), scale, return_dtype, mm_output_shape)
|
||||
result = dequantize_symmetric(torch.cat(result, dim=-1), scale, return_dtype, mm_output_shape)
|
||||
if bias is not None:
|
||||
result.add_(bias)
|
||||
|
||||
@@ -286,24 +286,24 @@ def conv_int8_matmul(
|
||||
|
||||
def quantized_linear_forward_fp8_matmul(self, input: torch.FloatTensor) -> torch.FloatTensor:
|
||||
if torch.numel(input) / input.shape[-1] < 32:
|
||||
return torch.nn.functional.linear(input, self.sdnq_decompressor(self.weight, skip_quantized_matmul=True), self.bias)
|
||||
return fp8_matmul(input, self.weight, self.bias, self.sdnq_decompressor.scale)
|
||||
return torch.nn.functional.linear(input, self.sdnq_dequantizer(self.weight, skip_quantized_matmul=True), self.bias)
|
||||
return fp8_matmul(input, self.weight, self.bias, self.sdnq_dequantizer.scale)
|
||||
|
||||
|
||||
def quantized_linear_forward_fp8_matmul_tensorwise(self, input: torch.FloatTensor) -> torch.FloatTensor:
|
||||
if torch.numel(input) / input.shape[-1] < 32:
|
||||
return torch.nn.functional.linear(input, self.sdnq_decompressor(self.weight, skip_quantized_matmul=True), self.bias)
|
||||
return fp8_matmul_tensorwise(input, self.weight, self.bias, self.sdnq_decompressor.scale)
|
||||
return torch.nn.functional.linear(input, self.sdnq_dequantizer(self.weight, skip_quantized_matmul=True), self.bias)
|
||||
return fp8_matmul_tensorwise(input, self.weight, self.bias, self.sdnq_dequantizer.scale)
|
||||
|
||||
|
||||
def quantized_linear_forward_int8_matmul(self, input: torch.FloatTensor) -> torch.FloatTensor:
|
||||
if torch.numel(input) / input.shape[-1] < 32:
|
||||
return torch.nn.functional.linear(input, self.sdnq_decompressor(self.weight, skip_quantized_matmul=True), self.bias)
|
||||
return int8_matmul(input, self.weight, self.bias, self.sdnq_decompressor.scale, getattr(self.sdnq_decompressor, "compressed_weight_shape", None), self.sdnq_decompressor.weights_dtype)
|
||||
return torch.nn.functional.linear(input, self.sdnq_dequantizer(self.weight, skip_quantized_matmul=True), self.bias)
|
||||
return int8_matmul(input, self.weight, self.bias, self.sdnq_dequantizer.scale, getattr(self.sdnq_dequantizer, "quantized_weight_shape", None), self.sdnq_dequantizer.weights_dtype)
|
||||
|
||||
|
||||
def quantized_linear_forward(self, input: torch.FloatTensor) -> torch.FloatTensor:
|
||||
return torch.nn.functional.linear(input, self.sdnq_decompressor(self.weight), self.bias)
|
||||
return torch.nn.functional.linear(input, self.sdnq_dequantizer(self.weight), self.bias)
|
||||
|
||||
|
||||
def get_conv_args(input_ndim: int, stride, padding, dilation):
|
||||
@@ -328,12 +328,12 @@ def get_conv_args(input_ndim: int, stride, padding, dilation):
|
||||
|
||||
def quantized_conv_forward_fp8_matmul(self, input) -> torch.FloatTensor:
|
||||
if torch.numel(input) / input.shape[2] < 32:
|
||||
return self._conv_forward(input, self.sdnq_decompressor(self.weight, skip_quantized_matmul=True), self.bias)
|
||||
return self._conv_forward(input, self.sdnq_dequantizer(self.weight, skip_quantized_matmul=True), self.bias)
|
||||
conv_type, stride, padding, dilation = get_conv_args(input.ndim, self.stride, self.padding, self.dilation)
|
||||
return conv_fp8_matmul(
|
||||
input, self.weight, self.bias,
|
||||
self.sdnq_decompressor.scale,
|
||||
self.sdnq_decompressor.result_shape,
|
||||
self.sdnq_dequantizer.scale,
|
||||
self.sdnq_dequantizer.result_shape,
|
||||
self._reversed_padding_repeated_twice,
|
||||
self.padding_mode, conv_type,
|
||||
self.groups, stride, padding, dilation,
|
||||
@@ -342,12 +342,12 @@ def quantized_conv_forward_fp8_matmul(self, input) -> torch.FloatTensor:
|
||||
|
||||
def quantized_conv_forward_fp8_matmul_tensorwise(self, input) -> torch.FloatTensor:
|
||||
if torch.numel(input) / input.shape[2] < 32:
|
||||
return self._conv_forward(input, self.sdnq_decompressor(self.weight, skip_quantized_matmul=True), self.bias)
|
||||
return self._conv_forward(input, self.sdnq_dequantizer(self.weight, skip_quantized_matmul=True), self.bias)
|
||||
conv_type, stride, padding, dilation = get_conv_args(input.ndim, self.stride, self.padding, self.dilation)
|
||||
return conv_fp8_matmul_tensorwise(
|
||||
input, self.weight, self.bias,
|
||||
self.sdnq_decompressor.scale,
|
||||
self.sdnq_decompressor.result_shape,
|
||||
self.sdnq_dequantizer.scale,
|
||||
self.sdnq_dequantizer.result_shape,
|
||||
self._reversed_padding_repeated_twice,
|
||||
self.padding_mode, conv_type,
|
||||
self.groups, stride, padding, dilation,
|
||||
@@ -356,14 +356,14 @@ def quantized_conv_forward_fp8_matmul_tensorwise(self, input) -> torch.FloatTens
|
||||
|
||||
def quantized_conv_forward_int8_matmul(self, input) -> torch.FloatTensor:
|
||||
if torch.numel(input) / input.shape[2] < 32:
|
||||
return self._conv_forward(input, self.sdnq_decompressor(self.weight, skip_quantized_matmul=True), self.bias)
|
||||
return self._conv_forward(input, self.sdnq_dequantizer(self.weight, skip_quantized_matmul=True), self.bias)
|
||||
conv_type, stride, padding, dilation = get_conv_args(input.ndim, self.stride, self.padding, self.dilation)
|
||||
return conv_int8_matmul(
|
||||
input, self.weight, self.bias,
|
||||
self.sdnq_decompressor.scale,
|
||||
self.sdnq_decompressor.result_shape,
|
||||
getattr(self.sdnq_decompressor, "compressed_weight_shape", None),
|
||||
self.sdnq_decompressor.weights_dtype,
|
||||
self.sdnq_dequantizer.scale,
|
||||
self.sdnq_dequantizer.result_shape,
|
||||
getattr(self.sdnq_dequantizer, "quantized_weight_shape", None),
|
||||
self.sdnq_dequantizer.weights_dtype,
|
||||
self._reversed_padding_repeated_twice,
|
||||
self.padding_mode, conv_type,
|
||||
self.groups, stride, padding, dilation,
|
||||
@@ -371,25 +371,25 @@ def quantized_conv_forward_int8_matmul(self, input) -> torch.FloatTensor:
|
||||
|
||||
|
||||
def quantized_conv_forward(self, input) -> torch.FloatTensor:
|
||||
return self._conv_forward(input, self.sdnq_decompressor(self.weight), self.bias)
|
||||
return self._conv_forward(input, self.sdnq_dequantizer(self.weight), self.bias)
|
||||
|
||||
|
||||
def quantized_conv_transpose_1d_forward(self, input: torch.FloatTensor, output_size: Optional[list[int]] = None) -> torch.FloatTensor:
|
||||
output_padding = self._output_padding(input, output_size, self.stride, self.padding, self.kernel_size, 1, self.dilation)
|
||||
return torch.nn.functional.conv_transpose1d(input, self.sdnq_decompressor(self.weight), self.bias, self.stride, self.padding, output_padding, self.groups, self.dilation)
|
||||
return torch.nn.functional.conv_transpose1d(input, self.sdnq_dequantizer(self.weight), self.bias, self.stride, self.padding, output_padding, self.groups, self.dilation)
|
||||
|
||||
|
||||
def quantized_conv_transpose_2d_forward(self, input: torch.FloatTensor, output_size: Optional[list[int]] = None) -> torch.FloatTensor:
|
||||
output_padding = self._output_padding(input, output_size, self.stride, self.padding, self.kernel_size, 2, self.dilation)
|
||||
return torch.nn.functional.conv_transpose2d(input, self.sdnq_decompressor(self.weight), self.bias, self.stride, self.padding, output_padding, self.groups, self.dilation)
|
||||
return torch.nn.functional.conv_transpose2d(input, self.sdnq_dequantizer(self.weight), self.bias, self.stride, self.padding, output_padding, self.groups, self.dilation)
|
||||
|
||||
|
||||
def quantized_conv_transpose_3d_forward(self, input: torch.FloatTensor, output_size: Optional[list[int]] = None) -> torch.FloatTensor:
|
||||
output_padding = self._output_padding(input, output_size, self.stride, self.padding, self.kernel_size, 3, self.dilation)
|
||||
return torch.nn.functional.conv_transpose3d(input, self.sdnq_decompressor(self.weight), self.bias, self.stride, self.padding, output_padding, self.groups, self.dilation)
|
||||
return torch.nn.functional.conv_transpose3d(input, self.sdnq_dequantizer(self.weight), self.bias, self.stride, self.padding, output_padding, self.groups, self.dilation)
|
||||
|
||||
|
||||
if shared.opts.sdnq_decompress_compile:
|
||||
if shared.opts.sdnq_dequantize_compile:
|
||||
try:
|
||||
torch._dynamo.config.cache_size_limit = max(8192, torch._dynamo.config.cache_size_limit)
|
||||
int8_matmul = torch.compile(int8_matmul, fullgraph=True)
|
||||
|
||||
+2
-2
@@ -521,11 +521,11 @@ options_templates.update(options_section(("quantization", "Quantization Settings
|
||||
"sdnq_quantize_weights_mode": OptionInfo("int8", "Quantization type", gr.Dropdown, {"choices": ["int8", "float8_e4m3fn", "int7", "int6", "int5", "uint4", "uint3", "uint2", "float8_e5m2", "float8_e4m3fnuz", "float8_e5m2fnuz", "uint8", "uint7", "uint6", "uint5", "int4", "int3", "int2", "uint1"], "visible": native}),
|
||||
"sdnq_quantize_weights_group_size": OptionInfo(0, "Group size", gr.Slider, {"minimum": -1, "maximum": 4096, "step": 1, "visible": native}),
|
||||
"sdnq_quantize_conv_layers": OptionInfo(False, "Quantize the convolutional layers", gr.Checkbox, {"visible": native}),
|
||||
"sdnq_decompress_compile": OptionInfo(devices.has_triton(), "Decompress using torch.compile", gr.Checkbox, {"visible": native}),
|
||||
"sdnq_dequantize_compile": OptionInfo(devices.has_triton(), "Dequantize using torch.compile", gr.Checkbox, {"visible": native}),
|
||||
"sdnq_use_quantized_matmul": OptionInfo(False, "Use Quantized MatMul", gr.Checkbox, {"visible": native}),
|
||||
"sdnq_use_quantized_matmul_conv": OptionInfo(False, "Use Quantized MatMul with convolutional layers", gr.Checkbox, {"visible": native}),
|
||||
"sdnq_quantize_with_gpu": OptionInfo(True, "Quantize with the GPU", gr.Checkbox, {"visible": native}),
|
||||
"sdnq_decompress_fp32": OptionInfo(False, "Decompress using full precision", gr.Checkbox, {"visible": native}),
|
||||
"sdnq_dequantize_fp32": OptionInfo(False, "Dequantize using full precision", gr.Checkbox, {"visible": native}),
|
||||
"sdnq_quantize_shuffle_weights": OptionInfo(False, "Shuffle weights in post mode", gr.Checkbox, {"visible": native}),
|
||||
|
||||
"bnb_quantization_sep": OptionInfo("<h2>BitsAndBytes</h2>", "", gr.HTML),
|
||||
|
||||
+1
-1
Submodule wiki updated: 70ea13a0c1...04cfb75b89
Reference in New Issue
Block a user