mirror of
https://github.com/vladmandic/automatic
synced 2026-09-19 09:14:35 +02:00
SDNQ add SVDQuant quantization method
This commit is contained in:
Disty0
@@ -52,6 +52,12 @@ def network_backup_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.n
|
||||
self.sdnq_zero_point_backup = self.zero_point.clone().to(devices.cpu)
|
||||
else:
|
||||
self.sdnq_zero_point_backup = None
|
||||
if self.svd_up is not None:
|
||||
self.sdnq_svd_up_backup = self.svd_up.clone().to(devices.cpu)
|
||||
self.sdnq_svd_down_backup = self.svd_down.clone().to(devices.cpu)
|
||||
else:
|
||||
self.sdnq_svd_up_backup = None
|
||||
self.sdnq_svd_down_backup = None
|
||||
|
||||
if bias_backup is None:
|
||||
if getattr(self, 'bias', None) is not None:
|
||||
@@ -85,9 +91,23 @@ def network_calc_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn.
|
||||
try:
|
||||
t0 = time.time()
|
||||
if hasattr(self, "sdnq_dequantizer_backup"):
|
||||
weight = self.sdnq_dequantizer_backup.to(devices.device)(self.weight.to(devices.device), self.sdnq_scale_backup.to(devices.device), self.sdnq_zero_point_backup.to(devices.device) if self.sdnq_zero_point_backup is not None else None, skip_quantized_matmul=self.sdnq_dequantizer_backup.use_quantized_matmul)
|
||||
weight = self.sdnq_dequantizer_backup.to(devices.device)(
|
||||
self.weight.to(devices.device),
|
||||
self.sdnq_scale_backup.to(devices.device),
|
||||
self.sdnq_zero_point_backup.to(devices.device) if self.sdnq_zero_point_backup is not None else None,
|
||||
self.sdnq_svd_up_backup.to(devices.device) if self.sdnq_svd_up_backup is not None else None,
|
||||
self.sdnq_svd_down_backup.to(devices.device) if self.sdnq_svd_down_backup is not None else None,
|
||||
skip_quantized_matmul=self.sdnq_dequantizer_backup.use_quantized_matmul
|
||||
)
|
||||
elif hasattr(self, "sdnq_dequantizer"):
|
||||
weight = self.sdnq_dequantizer.to(devices.device)(self.weight.to(devices.device), self.scale.to(devices.device), self.zero_point.to(devices.device) if self.zero_point is not None else None, skip_quantized_matmul=self.sdnq_dequantizer.use_quantized_matmul)
|
||||
weight = self.sdnq_dequantizer.to(devices.device)(
|
||||
self.weight.to(devices.device),
|
||||
self.scale.to(devices.device),
|
||||
self.zero_point.to(devices.device) if self.zero_point is not None else None,
|
||||
self.svd_up.to(devices.device) if self.svd_up is not None else None,
|
||||
self.svd_down.to(devices.device) if self.svd_down is not None else None,
|
||||
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)
|
||||
@@ -148,19 +168,35 @@ def network_add_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn.G
|
||||
from modules.sdnq import sdnq_quantize_layer
|
||||
if hasattr(self, "sdnq_dequantizer_backup"):
|
||||
weights_dtype = self.sdnq_dequantizer_backup.weights_dtype
|
||||
dequant_weight = self.sdnq_dequantizer_backup.to(devices.device)(model_weights.to(devices.device), self.sdnq_scale_backup.to(devices.device), self.sdnq_zero_point_backup.to(devices.device) if self.sdnq_zero_point_backup is not None else None, skip_quantized_matmul=self.sdnq_dequantizer_backup.use_quantized_matmul)
|
||||
dequant_weight = self.sdnq_dequantizer_backup.to(devices.device)(
|
||||
model_weights.to(devices.device),
|
||||
self.sdnq_scale_backup.to(devices.device),
|
||||
self.sdnq_zero_point_backup.to(devices.device) if self.sdnq_zero_point_backup is not None else None,
|
||||
self.sdnq_svd_up_backup.to(devices.device) if self.sdnq_svd_up_backup is not None else None,
|
||||
self.sdnq_svd_down_backup.to(devices.device) if self.sdnq_svd_down_backup is not None else None,
|
||||
skip_quantized_matmul=self.sdnq_dequantizer_backup.use_quantized_matmul
|
||||
)
|
||||
elif hasattr(self, "sdnq_dequantizer"):
|
||||
weights_dtype = self.sdnq_dequantizer.weights_dtype
|
||||
dequant_weight = self.sdnq_dequantizer.to(devices.device)(model_weights.to(devices.device), self.scale.to(devices.device), self.zero_point.to(devices.device) if self.zero_point is not None else None, skip_quantized_matmul=self.sdnq_dequantizer.use_quantized_matmul)
|
||||
dequant_weight = self.sdnq_dequantizer.to(devices.device)(
|
||||
model_weights.to(devices.device),
|
||||
self.scale.to(devices.device),
|
||||
self.zero_point.to(devices.device) if self.zero_point is not None else None,
|
||||
self.svd_up.to(devices.device) if self.svd_up is not None else None,
|
||||
self.svd_down.to(devices.device) if self.svd_down is not None else None,
|
||||
skip_quantized_matmul=self.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)
|
||||
del self.sdnq_dequantizer, self.scale, self.zero_point
|
||||
del self.sdnq_dequantizer, self.scale, self.zero_point, self.svd_up, self.svd_down
|
||||
self = sdnq_quantize_layer(
|
||||
self,
|
||||
weights_dtype=weights_dtype,
|
||||
torch_dtype=devices.dtype,
|
||||
group_size=shared.opts.sdnq_quantize_weights_group_size,
|
||||
svd_rank=shared.opts.sdnq_svd_rank,
|
||||
use_svd=shared.opts.sdnq_use_svd,
|
||||
quant_conv=shared.opts.sdnq_quantize_conv_layers,
|
||||
use_quantized_matmul=shared.opts.sdnq_use_quantized_matmul,
|
||||
use_quantized_matmul_conv=shared.opts.sdnq_use_quantized_matmul_conv,
|
||||
@@ -243,7 +279,12 @@ def network_apply_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn
|
||||
self.zero_point = torch.nn.Parameter(self.sdnq_zero_point_backup.to(device), requires_grad=False)
|
||||
else:
|
||||
self.zero_point = None
|
||||
del self.sdnq_dequantizer_backup, self.sdnq_scale_backup, self.sdnq_zero_point_backup
|
||||
if self.sdnq_svd_up_backup is not None:
|
||||
self.svd_up = torch.nn.Parameter(self.sdnq_svd_up_backup.to(device), requires_grad=False)
|
||||
self.svd_down = torch.nn.Parameter(self.sdnq_svd_down_backup.to(device), requires_grad=False)
|
||||
else:
|
||||
self.svd_up, self.svd_down = None, None
|
||||
del self.sdnq_dequantizer_backup, self.sdnq_scale_backup, self.sdnq_zero_point_backup, self.sdnq_svd_up_backup, self.sdnq_svd_down_backup
|
||||
|
||||
if bias_backup is not None:
|
||||
self.bias = None
|
||||
|
||||
@@ -205,6 +205,8 @@ def create_sdnq_config(kwargs = None, allow: bool = True, module: str = 'Model',
|
||||
sdnq_config = SDNQConfig(
|
||||
weights_dtype=weights_dtype,
|
||||
group_size=shared.opts.sdnq_quantize_weights_group_size,
|
||||
svd_rank=shared.opts.sdnq_svd_rank,
|
||||
use_svd=shared.opts.sdnq_use_svd,
|
||||
quant_conv=shared.opts.sdnq_quantize_conv_layers,
|
||||
use_quantized_matmul=shared.opts.sdnq_use_quantized_matmul,
|
||||
use_quantized_matmul_conv=shared.opts.sdnq_use_quantized_matmul_conv,
|
||||
@@ -215,7 +217,7 @@ def create_sdnq_config(kwargs = None, allow: bool = True, module: str = 'Model',
|
||||
modules_to_not_convert=modules_to_not_convert,
|
||||
modules_dtype_dict=modules_dtype_dict.copy(),
|
||||
)
|
||||
log.debug(f'Quantization: module="{module}" type=sdnq mode=pre dtype={weights_dtype} matmul={shared.opts.sdnq_use_quantized_matmul} group_size={shared.opts.sdnq_quantize_weights_group_size} quant_conv={shared.opts.sdnq_quantize_conv_layers} matmul_conv={shared.opts.sdnq_use_quantized_matmul_conv} dequantize_fp32={shared.opts.sdnq_dequantize_fp32} quantize_with_gpu={shared.opts.sdnq_quantize_with_gpu} quantization_device={quantization_device} return_device={return_device} device_map={shared.opts.device_map} offload_mode={shared.opts.diffusers_offload_mode} non_blocking={shared.opts.diffusers_offload_nonblocking} modules_to_not_convert={modules_to_not_convert} modules_dtype_dict={modules_dtype_dict}')
|
||||
log.debug(f'Quantization: module="{module}" type=sdnq mode=pre dtype={weights_dtype} matmul={shared.opts.sdnq_use_quantized_matmul} group_size={shared.opts.sdnq_quantize_weights_group_size} svd_rank={shared.opts.sdnq_svd_rank} use_svd={shared.opts.sdnq_use_svd} quant_conv={shared.opts.sdnq_quantize_conv_layers} matmul_conv={shared.opts.sdnq_use_quantized_matmul_conv} dequantize_fp32={shared.opts.sdnq_dequantize_fp32} quantize_with_gpu={shared.opts.sdnq_quantize_with_gpu} quantization_device={quantization_device} return_device={return_device} device_map={shared.opts.device_map} offload_mode={shared.opts.diffusers_offload_mode} non_blocking={shared.opts.diffusers_offload_nonblocking} modules_to_not_convert={modules_to_not_convert} modules_dtype_dict={modules_dtype_dict}')
|
||||
if kwargs is None:
|
||||
return sdnq_config
|
||||
else:
|
||||
@@ -533,6 +535,8 @@ def sdnq_quantize_model(model, op=None, sd_model=None, do_gc: bool = True, weigh
|
||||
weights_dtype=weights_dtype,
|
||||
torch_dtype=devices.dtype,
|
||||
group_size=shared.opts.sdnq_quantize_weights_group_size,
|
||||
svd_rank=shared.opts.sdnq_svd_rank,
|
||||
use_svd=shared.opts.sdnq_use_svd,
|
||||
quant_conv=shared.opts.sdnq_quantize_conv_layers,
|
||||
use_quantized_matmul=shared.opts.sdnq_use_quantized_matmul,
|
||||
use_quantized_matmul_conv=shared.opts.sdnq_use_quantized_matmul_conv,
|
||||
@@ -572,7 +576,7 @@ def sdnq_quantize_model(model, op=None, sd_model=None, do_gc: bool = True, weigh
|
||||
if do_gc:
|
||||
devices.torch_gc(force=True, reason='sdnq')
|
||||
|
||||
log.debug(f'Quantization: module="{op if op is not None else model.__class__}" type=sdnq mode=post dtype={weights_dtype} matmul={shared.opts.sdnq_use_quantized_matmul} group_size={shared.opts.sdnq_quantize_weights_group_size} quant_conv={shared.opts.sdnq_quantize_conv_layers} matmul_conv={shared.opts.sdnq_use_quantized_matmul_conv} dequantize_fp32={shared.opts.sdnq_dequantize_fp32} quantize_with_gpu={shared.opts.sdnq_quantize_with_gpu} quantization_device={quantization_device} return_device={return_device} device_map={shared.opts.device_map} offload_mode={shared.opts.diffusers_offload_mode} non_blocking={shared.opts.diffusers_offload_nonblocking} modules_to_not_convert={modules_to_not_convert} modules_dtype_dict={modules_dtype_dict}')
|
||||
log.debug(f'Quantization: module="{op if op is not None else model.__class__}" type=sdnq mode=post dtype={weights_dtype} matmul={shared.opts.sdnq_use_quantized_matmul} group_size={shared.opts.sdnq_quantize_weights_group_size} svd_rank={shared.opts.sdnq_svd_rank} use_svd={shared.opts.sdnq_use_svd} quant_conv={shared.opts.sdnq_quantize_conv_layers} matmul_conv={shared.opts.sdnq_use_quantized_matmul_conv} dequantize_fp32={shared.opts.sdnq_dequantize_fp32} quantize_with_gpu={shared.opts.sdnq_quantize_with_gpu} quantization_device={quantization_device} return_device={return_device} device_map={shared.opts.device_map} offload_mode={shared.opts.diffusers_offload_mode} non_blocking={shared.opts.diffusers_offload_nonblocking} modules_to_not_convert={modules_to_not_convert} modules_dtype_dict={modules_dtype_dict}')
|
||||
return model
|
||||
|
||||
|
||||
@@ -580,7 +584,7 @@ def sdnq_quantize_weights(sd_model):
|
||||
try:
|
||||
t0 = time.time()
|
||||
from modules import shared, devices, sd_models
|
||||
log.debug(f"Quantization: type=SDNQ modules={shared.opts.sdnq_quantize_weights} dtype={shared.opts.sdnq_quantize_weights_mode} dtype_te={shared.opts.sdnq_quantize_weights_mode_te} matmul={shared.opts.sdnq_use_quantized_matmul} group_size={shared.opts.sdnq_quantize_weights_group_size} quant_conv={shared.opts.sdnq_quantize_conv_layers} matmul_conv={shared.opts.sdnq_use_quantized_matmul_conv} quantize_with_gpu={shared.opts.sdnq_quantize_with_gpu} dequantize_fp32={shared.opts.sdnq_dequantize_fp32} pre_forward={shared.opts.diffusers_offload_pre}")
|
||||
log.debug(f"Quantization: type=SDNQ modules={shared.opts.sdnq_quantize_weights} dtype={shared.opts.sdnq_quantize_weights_mode} dtype_te={shared.opts.sdnq_quantize_weights_mode_te} matmul={shared.opts.sdnq_use_quantized_matmul} svd_rank={shared.opts.sdnq_svd_rank} use_svd={shared.opts.sdnq_use_svd} group_size={shared.opts.sdnq_quantize_weights_group_size} quant_conv={shared.opts.sdnq_quantize_conv_layers} matmul_conv={shared.opts.sdnq_use_quantized_matmul_conv} quantize_with_gpu={shared.opts.sdnq_quantize_with_gpu} dequantize_fp32={shared.opts.sdnq_dequantize_fp32} pre_forward={shared.opts.diffusers_offload_pre}")
|
||||
global quant_last_model_name, quant_last_model_device # pylint: disable=global-statement
|
||||
|
||||
sd_model = sd_models.apply_function_to_model(sd_model, sdnq_quantize_model, shared.opts.sdnq_quantize_weights, op="sdnq")
|
||||
|
||||
+58
-17
@@ -48,8 +48,15 @@ def quantize_weight(weight: torch.FloatTensor, reduction_axes: Union[int, List[i
|
||||
return quantized_weight, scale, zero_point
|
||||
|
||||
|
||||
def apply_svdquant(weight: torch.FloatTensor, rank: int = 128) -> Tuple[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]:
|
||||
U, S, svd_down = torch.svd_lowrank(weight, q=rank)
|
||||
svd_up = torch.mul(U, S.unsqueeze(0))
|
||||
svd_down = svd_down.t_()
|
||||
return weight.sub_(torch.mm(svd_up, svd_down)), svd_up, svd_down
|
||||
|
||||
|
||||
@devices.inference_context()
|
||||
def sdnq_quantize_layer(layer, weights_dtype="int8", torch_dtype=None, group_size=0, quant_conv=False, use_quantized_matmul=False, use_quantized_matmul_conv=False, dequantize_fp32=False, non_blocking=False, quantization_device=None, return_device=None, param_name=None): # pylint: disable=unused-argument
|
||||
def sdnq_quantize_layer(layer, weights_dtype="int8", torch_dtype=None, group_size=0, svd_rank=128, use_svd=False, quant_conv=False, use_quantized_matmul=False, use_quantized_matmul_conv=False, dequantize_fp32=False, non_blocking=False, quantization_device=None, return_device=None, param_name=None): # pylint: disable=unused-argument
|
||||
layer_class_name = layer.__class__.__name__
|
||||
if layer_class_name in allowed_types:
|
||||
num_of_groups = 1
|
||||
@@ -103,13 +110,30 @@ def sdnq_quantize_layer(layer, weights_dtype="int8", torch_dtype=None, group_siz
|
||||
else:
|
||||
use_quantized_matmul = output_channel_size % 16 == 0 and channel_size % 16 == 0
|
||||
|
||||
layer.weight.requires_grad = False
|
||||
if return_device is None:
|
||||
return_device = layer.weight.device
|
||||
if quantization_device is not None:
|
||||
layer.weight.data = layer.weight.to(quantization_device, non_blocking=non_blocking)
|
||||
if layer.weight.dtype != torch.float32:
|
||||
layer.weight.data = layer.weight.to(dtype=torch.float32)
|
||||
|
||||
if use_svd and is_linear_type:
|
||||
layer.weight.data, svd_up, svd_down = apply_svdquant(layer.weight, rank=svd_rank)
|
||||
if use_quantized_matmul:
|
||||
svd_up = svd_up.t_()
|
||||
svd_down = svd_down.t_()
|
||||
else:
|
||||
use_svd = False
|
||||
svd_up, svd_down = None, None
|
||||
|
||||
if group_size == 0:
|
||||
if use_quantized_matmul and dtype_dict[weights_dtype]["num_bits"] >= 6:
|
||||
group_size = -1
|
||||
elif is_linear_type:
|
||||
group_size = 2 ** (2 + dtype_dict[weights_dtype]["num_bits"])
|
||||
group_size = 2 ** ((2 if not use_svd else 3) + dtype_dict[weights_dtype]["num_bits"])
|
||||
else:
|
||||
group_size = 2 ** (1 + dtype_dict[weights_dtype]["num_bits"])
|
||||
group_size = 2 ** ((1 if not use_svd else 2) + dtype_dict[weights_dtype]["num_bits"])
|
||||
elif use_quantized_matmul and dtype_dict[weights_dtype]["num_bits"] == 8:
|
||||
group_size = -1 # override user value, re-quantizing 8bit into 8bit is pointless
|
||||
elif group_size != -1 and not is_linear_type:
|
||||
@@ -154,24 +178,19 @@ def sdnq_quantize_layer(layer, weights_dtype="int8", torch_dtype=None, group_siz
|
||||
new_shape[last_dim_index - 1 : last_dim_index] = (num_of_groups, group_size)
|
||||
layer.weight.data = layer.weight.reshape(new_shape)
|
||||
|
||||
layer.weight.requires_grad = False
|
||||
if return_device is None:
|
||||
return_device = layer.weight.device
|
||||
if quantization_device is not None:
|
||||
layer.weight.data = layer.weight.to(quantization_device, non_blocking=non_blocking)
|
||||
if layer.weight.dtype != torch.float32:
|
||||
layer.weight.data = layer.weight.to(dtype=torch.float32)
|
||||
|
||||
layer.weight.data, scale, zero_point = quantize_weight(layer.weight, reduction_axes, weights_dtype)
|
||||
if not 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)
|
||||
scale = scale.to(dtype=torch_dtype)
|
||||
if zero_point is not None:
|
||||
zero_point = zero_point.to(torch_dtype)
|
||||
zero_point = zero_point.to(dtype=torch_dtype)
|
||||
if svd_up is not None:
|
||||
svd_up = svd_up.to(dtype=torch_dtype)
|
||||
svd_down = svd_down.to(dtype=torch_dtype)
|
||||
|
||||
re_quantize_for_matmul = (num_of_groups > 1 or zero_point is not None)
|
||||
if use_quantized_matmul and not re_quantize_for_matmul:
|
||||
scale.transpose_(0,1)
|
||||
layer.weight.transpose_(0,1)
|
||||
scale.t_()
|
||||
layer.weight.t_()
|
||||
if use_contiguous_mm:
|
||||
layer.weight.data = layer.weight.contiguous()
|
||||
elif layer.weight.is_contiguous():
|
||||
@@ -186,6 +205,13 @@ def sdnq_quantize_layer(layer, weights_dtype="int8", torch_dtype=None, group_siz
|
||||
layer.zero_point = torch.nn.Parameter(zero_point, requires_grad=False)
|
||||
else:
|
||||
layer.zero_point = None
|
||||
if svd_up is not None:
|
||||
svd_up = svd_up.to(return_device, non_blocking=non_blocking)
|
||||
svd_down = svd_down.to(return_device, non_blocking=non_blocking)
|
||||
layer.svd_up = torch.nn.Parameter(svd_up, requires_grad=False)
|
||||
layer.svd_down = torch.nn.Parameter(svd_down, requires_grad=False)
|
||||
else:
|
||||
layer.svd_up, layer.svd_down = None, None
|
||||
|
||||
layer.sdnq_dequantizer = dequantizer_dict[weights_dtype](
|
||||
quantized_weight_shape=layer.weight.shape,
|
||||
@@ -204,7 +230,7 @@ def sdnq_quantize_layer(layer, weights_dtype="int8", torch_dtype=None, group_siz
|
||||
return layer
|
||||
|
||||
|
||||
def apply_sdnq_to_module(model, weights_dtype="int8", torch_dtype=None, group_size=0, quant_conv=False, use_quantized_matmul=False, use_quantized_matmul_conv=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, op=None): # pylint: disable=unused-argument
|
||||
def apply_sdnq_to_module(model, weights_dtype="int8", torch_dtype=None, group_size=0, svd_rank=128, use_svd=False, quant_conv=False, use_quantized_matmul=False, use_quantized_matmul_conv=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, op=None): # pylint: disable=unused-argument
|
||||
has_children = list(model.children())
|
||||
if not has_children:
|
||||
return model
|
||||
@@ -238,6 +264,8 @@ def apply_sdnq_to_module(model, weights_dtype="int8", torch_dtype=None, group_si
|
||||
weights_dtype=weights_dtype,
|
||||
torch_dtype=torch_dtype,
|
||||
group_size=group_size,
|
||||
svd_rank=svd_rank,
|
||||
use_svd=use_svd,
|
||||
quant_conv=quant_conv,
|
||||
use_quantized_matmul=use_quantized_matmul,
|
||||
use_quantized_matmul_conv=use_quantized_matmul_conv,
|
||||
@@ -252,6 +280,8 @@ def apply_sdnq_to_module(model, weights_dtype="int8", torch_dtype=None, group_si
|
||||
weights_dtype=weights_dtype,
|
||||
torch_dtype=torch_dtype,
|
||||
group_size=group_size,
|
||||
svd_rank=svd_rank,
|
||||
use_svd=use_svd,
|
||||
quant_conv=quant_conv,
|
||||
use_quantized_matmul=use_quantized_matmul,
|
||||
use_quantized_matmul_conv=use_quantized_matmul_conv,
|
||||
@@ -371,6 +401,8 @@ class SDNQQuantizer(DiffusersQuantizer):
|
||||
weights_dtype=weights_dtype,
|
||||
torch_dtype=self.torch_dtype,
|
||||
group_size=self.quantization_config.group_size,
|
||||
svd_rank=self.quantization_config.svd_rank,
|
||||
use_svd=self.quantization_config.use_svd,
|
||||
quant_conv=self.quantization_config.quant_conv,
|
||||
use_quantized_matmul=self.quantization_config.use_quantized_matmul,
|
||||
use_quantized_matmul_conv=self.quantization_config.use_quantized_matmul_conv,
|
||||
@@ -497,8 +529,13 @@ class SDNQConfig(QuantizationConfigMixin):
|
||||
weights_dtype (`str`, *optional*, defaults to `"int8"`):
|
||||
The target dtype for the weights after quantization. Supported values are:
|
||||
("int8", "int7", "int6", "int5", "int4", "int3", "int2", "uint8", "uint7", "uint6", "uint5", "uint4", "uint3", "uint2", "uint1", "bool", "float8_e4m3fn", "float8_e4m3fnuz", "float8_e5m2", "float8_e5m2fnuz")
|
||||
weights_dtype (`int`, *optional*, defaults to `0`):
|
||||
group_size (`int`, *optional*, defaults to `0`):
|
||||
Used to decide how many elements of a tensor will share the same quantization group.
|
||||
group_size = 0 will automatically select a group size based on weights_dtype.
|
||||
svd_rank (`int`, *optional*, defaults to `128`):
|
||||
The rank size used for the SVDQuant algorithm.
|
||||
use_svd (`bool`, *optional*, defaults to `False`):
|
||||
Enabling this option will use SVDQuant algorithm.
|
||||
quant_conv (`bool`, *optional*, defaults to `False`):
|
||||
Enabling this option will quantize the convolutional layers in UNet models too.
|
||||
use_quantized_matmul (`bool`, *optional*, defaults to `False`):
|
||||
@@ -524,6 +561,8 @@ class SDNQConfig(QuantizationConfigMixin):
|
||||
self,
|
||||
weights_dtype: str = "int8",
|
||||
group_size: int = 0,
|
||||
svd_rank: int = 128,
|
||||
use_svd: bool = False,
|
||||
quant_conv: bool = False,
|
||||
use_quantized_matmul: bool = False,
|
||||
use_quantized_matmul_conv: bool = False,
|
||||
@@ -538,6 +577,8 @@ class SDNQConfig(QuantizationConfigMixin):
|
||||
self.weights_dtype = weights_dtype
|
||||
self.quant_method = QuantizationMethod.SDNQ
|
||||
self.group_size = group_size
|
||||
self.svd_rank = svd_rank
|
||||
self.use_svd = use_svd
|
||||
self.quant_conv = quant_conv
|
||||
self.use_quantized_matmul = use_quantized_matmul
|
||||
self.use_quantized_matmul_conv = use_quantized_matmul_conv
|
||||
|
||||
+45
-33
@@ -1,6 +1,6 @@
|
||||
# pylint: disable=redefined-builtin,no-member,protected-access
|
||||
|
||||
from typing import Tuple
|
||||
from typing import Tuple, Optional
|
||||
|
||||
import torch
|
||||
|
||||
@@ -8,19 +8,31 @@ from .common import dtype_dict, compile_func
|
||||
from .packed_int import pack_int_symetric, unpack_int_symetric, pack_int_asymetric, unpack_int_asymetric
|
||||
|
||||
|
||||
def dequantize_asymmetric(weight: torch.ByteTensor, scale: torch.FloatTensor, zero_point: torch.FloatTensor, dtype: torch.dtype, result_shape: torch.Size) -> torch.FloatTensor:
|
||||
result = torch.addcmul(zero_point, weight.to(dtype=scale.dtype), scale).to(dtype=dtype)
|
||||
def dequantize_asymmetric(weight: torch.ByteTensor, scale: torch.FloatTensor, zero_point: torch.FloatTensor, dtype: torch.dtype, result_shape: torch.Size, svd_up: Optional[torch.FloatTensor] = None, svd_down: Optional[torch.FloatTensor] = None, skip_quantized_matmul: bool = False) -> torch.FloatTensor:
|
||||
result = torch.addcmul(zero_point, weight.to(dtype=scale.dtype), scale)
|
||||
if result_shape is not None:
|
||||
result = result.view(result_shape)
|
||||
if svd_up is not None:
|
||||
if skip_quantized_matmul:
|
||||
svd_up, svd_down = svd_up.t(), svd_down.t()
|
||||
result = torch.addmm(result, svd_up, svd_down)
|
||||
if dtype is not None:
|
||||
result = result.to(dtype=dtype)
|
||||
return result
|
||||
|
||||
|
||||
def dequantize_symmetric(weight: torch.CharTensor, scale: torch.FloatTensor, dtype: torch.dtype, result_shape: torch.Size, skip_quantized_matmul: bool = False) -> torch.FloatTensor:
|
||||
result = weight.to(dtype=scale.dtype).mul_(scale).to(dtype=dtype)
|
||||
def dequantize_symmetric(weight: torch.CharTensor, scale: torch.FloatTensor, dtype: torch.dtype, result_shape: torch.Size, svd_up: Optional[torch.FloatTensor] = None, svd_down: Optional[torch.FloatTensor] = None, skip_quantized_matmul: bool = False) -> torch.FloatTensor:
|
||||
result = weight.to(dtype=scale.dtype).mul_(scale)
|
||||
if skip_quantized_matmul:
|
||||
result.t_()
|
||||
if result_shape is not None:
|
||||
result = result.view(result_shape)
|
||||
if svd_up is not None:
|
||||
if skip_quantized_matmul:
|
||||
svd_up, svd_down = svd_up.t(), svd_down.t()
|
||||
result = torch.addmm(result, svd_up, svd_down)
|
||||
if dtype is not None:
|
||||
result = result.to(dtype=dtype)
|
||||
return result
|
||||
|
||||
|
||||
@@ -28,12 +40,12 @@ def dequantize_symmetric_with_bias(weight: torch.CharTensor, scale: torch.FloatT
|
||||
return torch.addcmul(bias, weight.to(dtype=scale.dtype), scale).to(dtype=dtype).view(result_shape)
|
||||
|
||||
|
||||
def dequantize_packed_int_asymmetric(weight: 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(unpack_int_asymetric(weight, shape, weights_dtype), scale, zero_point, dtype, result_shape)
|
||||
def dequantize_packed_int_asymmetric(weight: torch.ByteTensor, scale: torch.FloatTensor, zero_point: torch.FloatTensor, shape: torch.Size, dtype: torch.dtype, result_shape: torch.Size, weights_dtype: str, svd_up: Optional[torch.FloatTensor] = None, svd_down: Optional[torch.FloatTensor] = None, skip_quantized_matmul: bool = False) -> torch.FloatTensor:
|
||||
return dequantize_asymmetric(unpack_int_asymetric(weight, shape, weights_dtype), scale, zero_point, dtype, result_shape, svd_up=svd_up, svd_down=svd_down, skip_quantized_matmul=skip_quantized_matmul)
|
||||
|
||||
|
||||
def dequantize_packed_int_symmetric(weight: 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:
|
||||
return dequantize_symmetric(unpack_int_symetric(weight, shape, weights_dtype, dtype=scale.dtype), scale, dtype, result_shape, skip_quantized_matmul=skip_quantized_matmul)
|
||||
def dequantize_packed_int_symmetric(weight: torch.ByteTensor, scale: torch.FloatTensor, shape: torch.Size, dtype: torch.dtype, result_shape: torch.Size, weights_dtype: str, svd_up: Optional[torch.FloatTensor] = None, svd_down: Optional[torch.FloatTensor] = None, skip_quantized_matmul: bool = False) -> torch.FloatTensor:
|
||||
return dequantize_symmetric(unpack_int_symetric(weight, shape, weights_dtype, dtype=scale.dtype), scale, dtype, result_shape, svd_up=svd_up, svd_down=svd_down, skip_quantized_matmul=skip_quantized_matmul)
|
||||
|
||||
|
||||
def quantize_int8(input: torch.FloatTensor, dim: int = -1) -> Tuple[torch.CharTensor, torch.FloatTensor]:
|
||||
@@ -55,20 +67,20 @@ def re_quantize_int8(weight: torch.FloatTensor) -> Tuple[torch.CharTensor, torch
|
||||
return weight, scale
|
||||
|
||||
|
||||
def re_quantize_matmul_asymmetric(weight: torch.ByteTensor, scale: torch.FloatTensor, zero_point: torch.FloatTensor, result_shape: torch.Size) -> Tuple[torch.CharTensor, torch.FloatTensor]:
|
||||
return re_quantize_int8(dequantize_asymmetric(weight, scale, zero_point, scale.dtype, result_shape))
|
||||
def re_quantize_matmul_asymmetric(weight: torch.ByteTensor, scale: torch.FloatTensor, zero_point: torch.FloatTensor, result_shape: torch.Size, svd_up: Optional[torch.FloatTensor] = None, svd_down: Optional[torch.FloatTensor] = None) -> Tuple[torch.CharTensor, torch.FloatTensor]:
|
||||
return re_quantize_int8(dequantize_asymmetric(weight, scale, zero_point, scale.dtype, result_shape, svd_up=svd_up, svd_down=svd_down))
|
||||
|
||||
|
||||
def re_quantize_matmul_symmetric(weight: torch.CharTensor, scale: torch.FloatTensor, result_shape: torch.Size) -> Tuple[torch.CharTensor, torch.FloatTensor]:
|
||||
return re_quantize_int8(dequantize_symmetric(weight, scale, scale.dtype, result_shape))
|
||||
def re_quantize_matmul_symmetric(weight: torch.CharTensor, scale: torch.FloatTensor, result_shape: torch.Size, svd_up: Optional[torch.FloatTensor] = None, svd_down: Optional[torch.FloatTensor] = None) -> Tuple[torch.CharTensor, torch.FloatTensor]:
|
||||
return re_quantize_int8(dequantize_symmetric(weight, scale, scale.dtype, result_shape, svd_up=svd_up, svd_down=svd_down))
|
||||
|
||||
|
||||
def re_quantize_matmul_packed_int_asymmetric(weight: torch.ByteTensor, scale: torch.FloatTensor, zero_point: torch.FloatTensor, shape: torch.Size, result_shape: torch.Size, weights_dtype: str) -> torch.FloatTensor:
|
||||
return re_quantize_matmul_asymmetric(unpack_int_asymetric(weight, shape, weights_dtype), scale, zero_point, result_shape)
|
||||
def re_quantize_matmul_packed_int_asymmetric(weight: torch.ByteTensor, scale: torch.FloatTensor, zero_point: torch.FloatTensor, shape: torch.Size, result_shape: torch.Size, weights_dtype: str, svd_up: Optional[torch.FloatTensor] = None, svd_down: Optional[torch.FloatTensor] = None) -> torch.FloatTensor:
|
||||
return re_quantize_matmul_asymmetric(unpack_int_asymetric(weight, shape, weights_dtype), scale, zero_point, result_shape, svd_up=svd_up, svd_down=svd_down)
|
||||
|
||||
|
||||
def re_quantize_matmul_packed_int_symmetric(weight: torch.ByteTensor, scale: torch.FloatTensor, shape: torch.Size, result_shape: torch.Size, weights_dtype: str) -> torch.FloatTensor:
|
||||
return re_quantize_matmul_symmetric(unpack_int_symetric(weight, shape, weights_dtype, dtype=scale.dtype), scale, result_shape)
|
||||
def re_quantize_matmul_packed_int_symmetric(weight: torch.ByteTensor, scale: torch.FloatTensor, shape: torch.Size, result_shape: torch.Size, weights_dtype: str, svd_up: Optional[torch.FloatTensor] = None, svd_down: Optional[torch.FloatTensor] = None) -> torch.FloatTensor:
|
||||
return re_quantize_matmul_symmetric(unpack_int_symetric(weight, shape, weights_dtype, dtype=scale.dtype), scale, result_shape, svd_up=svd_up, svd_down=svd_down)
|
||||
|
||||
|
||||
class AsymmetricWeightsDequantizer(torch.nn.Module):
|
||||
@@ -92,11 +104,11 @@ class AsymmetricWeightsDequantizer(torch.nn.Module):
|
||||
def pack_weight(self, weight: torch.Tensor) -> torch.Tensor:
|
||||
return weight.to(dtype=dtype_dict[self.weights_dtype]["torch_dtype"])
|
||||
|
||||
def re_quantize_matmul(self, weight, scale, zero_point, **kwargs): # pylint: disable=unused-argument
|
||||
return re_quantize_matmul_asymmetric_compiled(weight, scale, zero_point, self.result_shape)
|
||||
def re_quantize_matmul(self, weight, scale, zero_point, svd_up, svd_down, **kwargs): # pylint: disable=unused-argument
|
||||
return re_quantize_matmul_asymmetric_compiled(weight, scale, zero_point, self.result_shape, svd_up=svd_up, svd_down=svd_down)
|
||||
|
||||
def forward(self, weight, scale, zero_point, **kwargs): # pylint: disable=unused-argument
|
||||
return dequantize_asymmetric_compiled(weight, scale, zero_point, self.result_dtype, self.result_shape)
|
||||
def forward(self, weight, scale, zero_point, svd_up, svd_down, skip_quantized_matmul=False): # pylint: disable=unused-argument
|
||||
return dequantize_asymmetric_compiled(weight, scale, zero_point, self.result_dtype, self.result_shape, svd_up=svd_up, svd_down=svd_down, skip_quantized_matmul=skip_quantized_matmul)
|
||||
|
||||
|
||||
class SymmetricWeightsDequantizer(torch.nn.Module):
|
||||
@@ -121,12 +133,12 @@ class SymmetricWeightsDequantizer(torch.nn.Module):
|
||||
def pack_weight(self, weight: torch.Tensor) -> torch.Tensor:
|
||||
return weight.to(dtype=dtype_dict[self.weights_dtype]["torch_dtype"])
|
||||
|
||||
def re_quantize_matmul(self, weight, scale, zero_point, **kwargs): # pylint: disable=unused-argument
|
||||
return re_quantize_matmul_symmetric_compiled(weight, scale, self.result_shape)
|
||||
def re_quantize_matmul(self, weight, scale, zero_point, svd_up, svd_down, **kwargs): # pylint: disable=unused-argument
|
||||
return re_quantize_matmul_symmetric_compiled(weight, scale, self.result_shape, svd_up=svd_up, svd_down=svd_down)
|
||||
|
||||
def forward(self, weight, scale, zero_point, skip_quantized_matmul=False, **kwargs): # pylint: disable=unused-argument
|
||||
def forward(self, weight, scale, zero_point, svd_up, svd_down, skip_quantized_matmul=False): # pylint: disable=unused-argument
|
||||
skip_quantized_matmul = skip_quantized_matmul and not self.re_quantize_for_matmul
|
||||
return dequantize_symmetric_compiled(weight, scale, self.result_dtype, self.result_shape, skip_quantized_matmul=skip_quantized_matmul)
|
||||
return dequantize_symmetric_compiled(weight, scale, self.result_dtype, self.result_shape, svd_up=svd_up, svd_down=svd_down, skip_quantized_matmul=skip_quantized_matmul)
|
||||
|
||||
|
||||
class PackedINTAsymmetricWeightsDequantizer(torch.nn.Module):
|
||||
@@ -152,11 +164,11 @@ class PackedINTAsymmetricWeightsDequantizer(torch.nn.Module):
|
||||
def pack_weight(self, weight: torch.Tensor) -> torch.Tensor:
|
||||
return pack_int_asymetric(weight, self.weights_dtype)
|
||||
|
||||
def re_quantize_matmul(self, weight, scale, zero_point, **kwargs): # pylint: disable=unused-argument
|
||||
return re_quantize_matmul_packed_int_asymmetric_compiled(weight, scale, zero_point, self.quantized_weight_shape, self.result_shape, self.weights_dtype)
|
||||
def re_quantize_matmul(self, weight, scale, zero_point, svd_up, svd_down, **kwargs): # pylint: disable=unused-argument
|
||||
return re_quantize_matmul_packed_int_asymmetric_compiled(weight, scale, zero_point, self.quantized_weight_shape, self.result_shape, self.weights_dtype, svd_up=svd_up, svd_down=svd_down)
|
||||
|
||||
def forward(self, weight, scale, zero_point, **kwargs): # pylint: disable=unused-argument
|
||||
return dequantize_packed_int_asymmetric_compiled(weight, scale, zero_point, self.quantized_weight_shape, self.result_dtype, self.result_shape, self.weights_dtype)
|
||||
def forward(self, weight, scale, zero_point, svd_up, svd_down, skip_quantized_matmul=False): # pylint: disable=unused-argument
|
||||
return dequantize_packed_int_asymmetric_compiled(weight, scale, zero_point, self.quantized_weight_shape, self.result_dtype, self.result_shape, self.weights_dtype, svd_up=svd_up, svd_down=svd_down, skip_quantized_matmul=skip_quantized_matmul)
|
||||
|
||||
|
||||
class PackedINTSymmetricWeightsDequantizer(torch.nn.Module):
|
||||
@@ -183,12 +195,12 @@ class PackedINTSymmetricWeightsDequantizer(torch.nn.Module):
|
||||
def pack_weight(self, weight: torch.Tensor) -> torch.Tensor:
|
||||
return pack_int_symetric(weight, self.weights_dtype)
|
||||
|
||||
def re_quantize_matmul(self, weight, scale, zero_point, **kwargs): # pylint: disable=unused-argument
|
||||
return re_quantize_matmul_packed_int_symmetric_compiled(weight, scale, self.quantized_weight_shape, self.result_shape, self.weights_dtype)
|
||||
def re_quantize_matmul(self, weight, scale, zero_point, svd_up, svd_down, **kwargs): # pylint: disable=unused-argument
|
||||
return re_quantize_matmul_packed_int_symmetric_compiled(weight, scale, self.quantized_weight_shape, self.result_shape, self.weights_dtype, svd_up=svd_up, svd_down=svd_down)
|
||||
|
||||
def forward(self, weight, scale, zero_point, skip_quantized_matmul=False, **kwargs): # pylint: disable=unused-argument
|
||||
def forward(self, weight, scale, zero_point, svd_up, svd_down, skip_quantized_matmul=False): # pylint: disable=unused-argument
|
||||
skip_quantized_matmul = skip_quantized_matmul and not self.re_quantize_for_matmul
|
||||
return dequantize_packed_int_symmetric_compiled(weight, scale, self.quantized_weight_shape, self.result_dtype, self.result_shape, self.weights_dtype, skip_quantized_matmul=skip_quantized_matmul)
|
||||
return dequantize_packed_int_symmetric_compiled(weight, scale, self.quantized_weight_shape, self.result_dtype, self.result_shape, self.weights_dtype, svd_up=svd_up, svd_down=svd_down, skip_quantized_matmul=skip_quantized_matmul)
|
||||
|
||||
|
||||
dequantizer_dict = {
|
||||
|
||||
@@ -15,6 +15,8 @@ def conv_fp8_matmul(
|
||||
weight: torch.Tensor,
|
||||
bias: torch.FloatTensor,
|
||||
scale: torch.FloatTensor,
|
||||
svd_up: torch.FloatTensor,
|
||||
svd_down: torch.FloatTensor,
|
||||
result_shape: torch.Size,
|
||||
reversed_padding_repeated_twice: List[int],
|
||||
padding_mode: str, conv_type: int,
|
||||
@@ -23,6 +25,9 @@ def conv_fp8_matmul(
|
||||
) -> torch.FloatTensor:
|
||||
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)
|
||||
if svd_up is not None:
|
||||
svd_bias = torch.mm(torch.mm(input.flatten(0,-2).to(dtype=svd_down.dtype), svd_down), svd_up)
|
||||
|
||||
input, input_scale = quantize_fp8_matmul_input(input)
|
||||
input, weight = check_mats(input, weight)
|
||||
|
||||
@@ -46,6 +51,8 @@ def conv_fp8_matmul(
|
||||
for i in range(groups):
|
||||
result.append(torch._scaled_mm(input[:, i], weight[:, i], scale_a=input_scale[i], scale_b=scale[i], bias=None, out_dtype=torch.bfloat16))
|
||||
result = torch.cat(result, dim=-1).view(mm_output_shape).to(return_dtype)
|
||||
if svd_up is not None:
|
||||
result.add_(svd_bias)
|
||||
|
||||
if conv_type == 1:
|
||||
result = result.transpose_(1,2)
|
||||
@@ -58,10 +65,11 @@ def conv_fp8_matmul(
|
||||
|
||||
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_dequantizer(self.weight, self.scale, self.zero_point, skip_quantized_matmul=True), self.bias)
|
||||
return self._conv_forward(input, self.sdnq_dequantizer(self.weight, self.scale, self.zero_point, self.svd_up, self.svd_down, 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.scale,
|
||||
input, self.weight, self.bias,
|
||||
self.scale, self.svd_up, self.svd_down,
|
||||
self.sdnq_dequantizer.result_shape,
|
||||
self._reversed_padding_repeated_twice,
|
||||
self.padding_mode, conv_type,
|
||||
|
||||
@@ -16,6 +16,8 @@ def conv_fp8_matmul_tensorwise(
|
||||
weight: torch.Tensor,
|
||||
bias: torch.FloatTensor,
|
||||
scale: torch.FloatTensor,
|
||||
svd_up: torch.FloatTensor,
|
||||
svd_down: torch.FloatTensor,
|
||||
result_shape: torch.Size,
|
||||
reversed_padding_repeated_twice: List[int],
|
||||
padding_mode: str, conv_type: int,
|
||||
@@ -24,6 +26,12 @@ def conv_fp8_matmul_tensorwise(
|
||||
) -> torch.FloatTensor:
|
||||
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)
|
||||
if svd_up is not None:
|
||||
if bias is not None:
|
||||
bias = torch.addmm(bias, torch.mm(input.flatten(0,-2).to(dtype=svd_down.dtype), svd_down), svd_up)
|
||||
else:
|
||||
bias = torch.mm(torch.mm(input.flatten(0,-2).to(dtype=svd_down.dtype), svd_down), svd_up)
|
||||
|
||||
input, scale = quantize_fp8_matmul_input_tensorwise(input, scale)
|
||||
input, weight = check_mats(input, weight)
|
||||
dummy_input_scale = torch.ones(1, device=input.device, dtype=torch.float32)
|
||||
@@ -53,10 +61,11 @@ def conv_fp8_matmul_tensorwise(
|
||||
|
||||
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_dequantizer(self.weight, self.scale, self.zero_point, skip_quantized_matmul=True), self.bias)
|
||||
return self._conv_forward(input, self.sdnq_dequantizer(self.weight, self.scale, self.zero_point, self.svd_up, self.svd_down, 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.scale,
|
||||
input, self.weight, self.bias,
|
||||
self.scale, self.svd_up, self.svd_down,
|
||||
self.sdnq_dequantizer.result_shape,
|
||||
self._reversed_padding_repeated_twice,
|
||||
self.padding_mode, conv_type,
|
||||
|
||||
@@ -17,6 +17,8 @@ def conv_int8_matmul(
|
||||
weight: torch.CharTensor,
|
||||
bias: torch.FloatTensor,
|
||||
scale: torch.FloatTensor,
|
||||
svd_up: torch.FloatTensor,
|
||||
svd_down: torch.FloatTensor,
|
||||
quantized_weight_shape: torch.Size,
|
||||
result_shape: torch.Size,
|
||||
weights_dtype: str,
|
||||
@@ -27,6 +29,12 @@ def conv_int8_matmul(
|
||||
) -> torch.FloatTensor:
|
||||
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)
|
||||
if svd_up is not None:
|
||||
if bias is not None:
|
||||
bias = torch.addmm(bias, torch.mm(input.flatten(0,-2).to(dtype=svd_down.dtype), svd_down), svd_up)
|
||||
else:
|
||||
bias = torch.mm(torch.mm(input.flatten(0,-2).to(dtype=svd_down.dtype), svd_down), svd_up)
|
||||
|
||||
input, scale = quantize_int8_matmul_input(input, scale)
|
||||
if quantized_weight_shape is not None:
|
||||
weight = unpack_int_symetric(weight, quantized_weight_shape, weights_dtype, dtype=torch.int8)
|
||||
@@ -57,10 +65,10 @@ def conv_int8_matmul(
|
||||
|
||||
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_dequantizer(self.weight, self.scale, self.zero_point, skip_quantized_matmul=True), self.bias)
|
||||
return self._conv_forward(input, self.sdnq_dequantizer(self.weight, self.scale, self.zero_point, self.svd_up, self.svd_down, skip_quantized_matmul=True), self.bias)
|
||||
conv_type, stride, padding, dilation = get_conv_args(input.ndim, self.stride, self.padding, self.dilation)
|
||||
if self.sdnq_dequantizer.re_quantize_for_matmul:
|
||||
weight, scale = self.sdnq_dequantizer.re_quantize_matmul(self.weight, self.scale, self.zero_point)
|
||||
weight, scale = self.sdnq_dequantizer.re_quantize_matmul(self.weight, self.scale, self.zero_point, None, None)
|
||||
quantized_weight_shape = None
|
||||
else:
|
||||
weight = self.weight
|
||||
@@ -68,7 +76,8 @@ def quantized_conv_forward_int8_matmul(self, input) -> torch.FloatTensor:
|
||||
quantized_weight_shape = getattr(self.sdnq_dequantizer, "quantized_weight_shape", None)
|
||||
return conv_int8_matmul(
|
||||
input, weight, self.bias,
|
||||
scale, quantized_weight_shape,
|
||||
scale, self.svd_up, self.svd_down,
|
||||
quantized_weight_shape,
|
||||
self.sdnq_dequantizer.result_shape,
|
||||
self.sdnq_dequantizer.weights_dtype,
|
||||
self._reversed_padding_repeated_twice,
|
||||
|
||||
@@ -75,19 +75,19 @@ def process_conv_input(conv_type, input, reversed_padding_repeated_twice, paddin
|
||||
|
||||
|
||||
def quantized_conv_forward(self, input) -> torch.FloatTensor:
|
||||
return self._conv_forward(input, self.sdnq_dequantizer(self.weight, self.scale, self.zero_point), self.bias)
|
||||
return self._conv_forward(input, self.sdnq_dequantizer(self.weight, self.scale, self.zero_point, self.svd_up, self.svd_down), 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_dequantizer(self.weight, self.scale, self.zero_point), 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.scale, self.zero_point, self.svd_up, self.svd_down), 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_dequantizer(self.weight, self.scale, self.zero_point), 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.scale, self.zero_point, self.svd_up, self.svd_down), 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_dequantizer(self.weight, self.scale, self.zero_point), 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.scale, self.zero_point, self.svd_up, self.svd_down), self.bias, self.stride, self.padding, output_padding, self.groups, self.dilation)
|
||||
|
||||
@@ -17,4 +17,4 @@ def check_mats(input: torch.Tensor, weight: torch.Tensor) -> Tuple[torch.Tensor,
|
||||
|
||||
|
||||
def quantized_linear_forward(self, input: torch.FloatTensor) -> torch.FloatTensor:
|
||||
return torch.nn.functional.linear(input, self.sdnq_dequantizer(self.weight, self.scale, self.zero_point), self.bias)
|
||||
return torch.nn.functional.linear(input, self.sdnq_dequantizer(self.weight, self.scale, self.zero_point, self.svd_up, self.svd_down), self.bias)
|
||||
|
||||
@@ -20,20 +20,27 @@ def fp8_matmul(
|
||||
weight: torch.Tensor,
|
||||
bias: torch.FloatTensor,
|
||||
scale: torch.FloatTensor,
|
||||
svd_up: torch.FloatTensor,
|
||||
svd_down: torch.FloatTensor,
|
||||
) -> torch.FloatTensor:
|
||||
return_dtype = input.dtype
|
||||
output_shape = (*input.shape[:-1], weight.shape[-1])
|
||||
if svd_up is not None:
|
||||
svd_bias = torch.mm(torch.mm(input.flatten(0,-2).to(dtype=svd_down.dtype), svd_down), svd_up)
|
||||
input, input_scale = quantize_fp8_matmul_input(input)
|
||||
input, weight = check_mats(input, weight)
|
||||
if bias is not None and bias.dtype != torch.bfloat16:
|
||||
bias = bias.to(dtype=torch.bfloat16)
|
||||
return torch._scaled_mm(input, weight, scale_a=input_scale, scale_b=scale, bias=bias, out_dtype=torch.bfloat16).view(output_shape).to(return_dtype)
|
||||
result = torch._scaled_mm(input, weight, scale_a=input_scale, scale_b=scale, bias=bias, out_dtype=torch.bfloat16).view(output_shape).to(return_dtype)
|
||||
if svd_up is not None:
|
||||
result.add_(svd_bias)
|
||||
return result
|
||||
|
||||
|
||||
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_dequantizer(self.weight, self.scale, self.zero_point, skip_quantized_matmul=True), self.bias)
|
||||
return fp8_matmul(input, self.weight, self.bias, self.scale)
|
||||
return torch.nn.functional.linear(input, self.sdnq_dequantizer(self.weight, self.scale, self.zero_point, self.svd_up, self.svd_down, skip_quantized_matmul=True), self.bias)
|
||||
return fp8_matmul(input, self.weight, self.bias, self.scale, self.svd_up, self.svd_down)
|
||||
|
||||
|
||||
fp8_matmul = compile_func(fp8_matmul)
|
||||
|
||||
@@ -23,10 +23,17 @@ def fp8_matmul_tensorwise(
|
||||
weight: torch.Tensor,
|
||||
bias: torch.FloatTensor,
|
||||
scale: torch.FloatTensor,
|
||||
svd_up: torch.FloatTensor,
|
||||
svd_down: torch.FloatTensor,
|
||||
) -> torch.FloatTensor:
|
||||
return_dtype = input.dtype
|
||||
output_shape = (*input.shape[:-1], weight.shape[-1])
|
||||
dummy_input_scale = torch.ones(1, device=input.device, dtype=torch.float32)
|
||||
if svd_up is not None:
|
||||
if bias is not None:
|
||||
bias = torch.addmm(bias, torch.mm(input.flatten(0,-2).to(dtype=svd_down.dtype), svd_down), svd_up)
|
||||
else:
|
||||
bias = torch.mm(torch.mm(input.flatten(0,-2).to(dtype=svd_down.dtype), svd_down), svd_up)
|
||||
input, scale = quantize_fp8_matmul_input_tensorwise(input, scale)
|
||||
input, weight = check_mats(input, weight)
|
||||
if bias is not None:
|
||||
@@ -37,8 +44,8 @@ def fp8_matmul_tensorwise(
|
||||
|
||||
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_dequantizer(self.weight, self.scale, self.zero_point, skip_quantized_matmul=True), self.bias)
|
||||
return fp8_matmul_tensorwise(input, self.weight, self.bias, self.scale)
|
||||
return torch.nn.functional.linear(input, self.sdnq_dequantizer(self.weight, self.scale, self.zero_point, self.svd_up, self.svd_down, skip_quantized_matmul=True), self.bias)
|
||||
return fp8_matmul_tensorwise(input, self.weight, self.bias, self.scale, self.svd_up, self.svd_down)
|
||||
|
||||
|
||||
fp8_matmul_tensorwise = compile_func(fp8_matmul_tensorwise)
|
||||
|
||||
@@ -24,6 +24,8 @@ def int8_matmul(
|
||||
weight: torch.Tensor,
|
||||
bias: torch.FloatTensor,
|
||||
scale: torch.FloatTensor,
|
||||
svd_up: torch.FloatTensor,
|
||||
svd_down: torch.FloatTensor,
|
||||
quantized_weight_shape: torch.Size,
|
||||
weights_dtype: str,
|
||||
) -> torch.FloatTensor:
|
||||
@@ -31,6 +33,11 @@ def int8_matmul(
|
||||
weight = unpack_int_symetric(weight, quantized_weight_shape, weights_dtype, dtype=torch.int8)
|
||||
return_dtype = input.dtype
|
||||
output_shape = (*input.shape[:-1], weight.shape[-1])
|
||||
if svd_up is not None:
|
||||
if bias is not None:
|
||||
bias = torch.addmm(bias, torch.mm(input.flatten(0,-2).to(dtype=svd_down.dtype), svd_down), svd_up)
|
||||
else:
|
||||
bias = torch.mm(torch.mm(input.flatten(0,-2).to(dtype=svd_down.dtype), svd_down), svd_up)
|
||||
input, scale = quantize_int8_matmul_input(input, scale)
|
||||
input, weight = check_mats(input, weight)
|
||||
if bias is not None:
|
||||
@@ -41,15 +48,15 @@ def int8_matmul(
|
||||
|
||||
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_dequantizer(self.weight, self.scale, self.zero_point, skip_quantized_matmul=True), self.bias)
|
||||
return torch.nn.functional.linear(input, self.sdnq_dequantizer(self.weight, self.scale, self.zero_point, self.svd_up, self.svd_down, skip_quantized_matmul=True), self.bias)
|
||||
if self.sdnq_dequantizer.re_quantize_for_matmul:
|
||||
weight, scale = self.sdnq_dequantizer.re_quantize_matmul(self.weight, self.scale, self.zero_point)
|
||||
weight, scale = self.sdnq_dequantizer.re_quantize_matmul(self.weight, self.scale, self.zero_point, None, None)
|
||||
quantized_weight_shape = None
|
||||
else:
|
||||
weight = self.weight
|
||||
scale = self.scale
|
||||
quantized_weight_shape = getattr(self.sdnq_dequantizer, "quantized_weight_shape", None)
|
||||
return int8_matmul(input, weight, self.bias, scale, quantized_weight_shape, self.sdnq_dequantizer.weights_dtype)
|
||||
return int8_matmul(input, weight, self.bias, scale, self.svd_up, self.svd_down, quantized_weight_shape, self.sdnq_dequantizer.weights_dtype)
|
||||
|
||||
|
||||
int8_matmul = compile_func(int8_matmul)
|
||||
|
||||
@@ -201,6 +201,8 @@ options_templates.update(options_section(("quantization", "Model Quantization"),
|
||||
"sdnq_modules_to_not_convert": OptionInfo("", "Modules to not convert"),
|
||||
"sdnq_modules_dtype_dict": OptionInfo("{}", "Modules dtype dict"),
|
||||
"sdnq_quantize_weights_group_size": OptionInfo(0, "Group size", gr.Slider, {"minimum": -1, "maximum": 4096, "step": 1}),
|
||||
"sdnq_svd_rank": OptionInfo(128, "SVDQuant Rank size", gr.Slider, {"minimum": 1, "maximum": 512, "step": 1}),
|
||||
"sdnq_use_svd": OptionInfo(False, "Use SVDQuant quantization", gr.Checkbox),
|
||||
"sdnq_quantize_conv_layers": OptionInfo(False, "Quantize convolutional layers", gr.Checkbox),
|
||||
"sdnq_dequantize_compile": OptionInfo(devices.has_triton(), "Dequantize using torch.compile", gr.Checkbox),
|
||||
"sdnq_use_quantized_matmul": OptionInfo(False, "Use quantized MatMul", gr.Checkbox),
|
||||
|
||||
Reference in New Issue
Block a user