mirror of
https://github.com/vladmandic/automatic
synced 2026-08-27 23:51:00 +02:00
Vladimir Mandic
1015 lines
49 KiB
Python
1015 lines
49 KiB
Python
# pylint: disable=redefined-builtin,no-member,protected-access
|
|
|
|
from dataclasses import dataclass
|
|
from enum import Enum
|
|
|
|
import torch
|
|
|
|
from transformers.quantizers import HfQuantizer
|
|
from diffusers.quantizers.base import DiffusersQuantizer
|
|
from diffusers.quantizers.quantization_config import QuantizationConfigMixin
|
|
|
|
from diffusers.utils import get_module_from_name
|
|
from accelerate import init_empty_weights
|
|
|
|
from modules import devices, shared
|
|
from .common import (
|
|
sdnq_version,
|
|
sdnq_keys,
|
|
dtype_dict,
|
|
accepted_weight_dtypes,
|
|
accepted_matmul_dtypes,
|
|
allowed_types,
|
|
linear_types,
|
|
embedding_types,
|
|
conv_types,
|
|
conv_transpose_types,
|
|
weights_dtype_order,
|
|
is_fp8_mm_supported,
|
|
use_tensorwise_fp8_matmul,
|
|
check_torch_compile,
|
|
compile_func,
|
|
)
|
|
from .dequantizer import SDNQDequantizer, dequantize_sdnq_model
|
|
from .packed_int import pack_int
|
|
from .packed_float import pack_float
|
|
from .forward import get_forward_func
|
|
from .layers import get_sdnq_wrapper_class
|
|
|
|
from .quant_utils import quantize_weight, apply_svdquant, apply_hadamard, prepare_weight_for_matmul, prepare_svd_for_matmul
|
|
from .utils import check_param_name_in, get_quant_args_from_config, get_quant_kwargs, add_module_skip_keys
|
|
|
|
|
|
class QuantizationMethod(str, Enum):
|
|
SDNQ = "sdnq"
|
|
SDNQ_TRAINING = "sdnq_training"
|
|
|
|
|
|
@devices.inference_context()
|
|
def sdnq_quantize_layer_weight(
|
|
weight: torch.FloatTensor,
|
|
layer_class_name: str | None = None,
|
|
weights_dtype: str = "int8",
|
|
quantized_matmul_dtype: str | None = None,
|
|
group_size: int = 0,
|
|
hadamard_group_size: int = 256,
|
|
svd_rank: int = 32,
|
|
svd_steps: int = 8,
|
|
use_svd: bool = False,
|
|
use_hadamard: bool = False,
|
|
use_quantized_matmul: bool = False,
|
|
use_stochastic_rounding: bool = False,
|
|
dequantize_fp32: bool = True,
|
|
hadamard: torch.FloatTensor | None = None,
|
|
using_pre_calculated_svd: bool = False,
|
|
using_pre_rotated_hadamard: bool = False,
|
|
skip_sr: bool = False,
|
|
param_name: str | None = None, # pylint: disable=unused-argument
|
|
torch_dtype: torch.dtype | None = None,
|
|
):
|
|
num_of_groups = 1
|
|
is_conv_type = False
|
|
is_conv_transpose_type = False
|
|
is_linear_type = False
|
|
result_shape = None
|
|
scale_dtype = None
|
|
|
|
weight = weight.detach()
|
|
original_shape = weight.shape
|
|
original_stride = weight.stride()
|
|
|
|
if torch_dtype is None:
|
|
torch_dtype = weight.dtype
|
|
if quantized_matmul_dtype is None:
|
|
if dtype_dict[weights_dtype]["is_integer"]:
|
|
quantized_matmul_dtype = "int8"
|
|
elif dtype_dict[weights_dtype]["num_bits"] < 16:
|
|
quantized_matmul_dtype = "float8_e4m3fn"
|
|
else:
|
|
quantized_matmul_dtype = "float16"
|
|
|
|
re_quantize_for_matmul = bool(
|
|
dtype_dict[weights_dtype]["is_unsigned"]
|
|
or dtype_dict[weights_dtype]["is_integer"] != dtype_dict[quantized_matmul_dtype]["is_integer"]
|
|
or dtype_dict[weights_dtype]["num_bits"] > dtype_dict[quantized_matmul_dtype]["num_bits"]
|
|
or (
|
|
dtype_dict[weights_dtype]["is_packed"]
|
|
and not dtype_dict[weights_dtype]["is_integer"]
|
|
and not dtype_dict[quantized_matmul_dtype]["is_integer"]
|
|
and (
|
|
dtype_dict[weights_dtype]["num_bits"] >= dtype_dict[quantized_matmul_dtype]["num_bits"]
|
|
or dtype_dict[weights_dtype]["max"] > dtype_dict[quantized_matmul_dtype]["max"]
|
|
)
|
|
)
|
|
)
|
|
|
|
if layer_class_name in conv_types:
|
|
is_conv_type = True
|
|
reduction_axes = 1
|
|
output_channel_size, channel_size = weight.shape[:2]
|
|
use_quantized_matmul = use_quantized_matmul and channel_size >= 32 and output_channel_size >= 32 and output_channel_size % 16 == 0 and channel_size % 16 == 0
|
|
if use_quantized_matmul and not re_quantize_for_matmul and not dtype_dict[weights_dtype]["is_packed"]:
|
|
result_shape = weight.shape
|
|
weight = weight.flatten(1,-1)
|
|
reduction_axes = -1
|
|
elif layer_class_name in conv_transpose_types:
|
|
is_conv_transpose_type = True
|
|
reduction_axes = 0
|
|
channel_size, output_channel_size = weight.shape[:2]
|
|
use_quantized_matmul = False
|
|
elif layer_class_name in linear_types:
|
|
is_linear_type = True
|
|
reduction_axes = -1
|
|
output_channel_size, channel_size = weight.shape
|
|
use_quantized_matmul = use_quantized_matmul and channel_size >= 32 and output_channel_size >= 32 and output_channel_size % 16 == 0 and channel_size % 16 == 0
|
|
else:
|
|
if weight.ndim > 1:
|
|
output_channel_size, channel_size = weight.shape[-2:]
|
|
else:
|
|
output_channel_size, channel_size = 1, weight.shape[-1]
|
|
reduction_axes = -1
|
|
use_quantized_matmul = False
|
|
|
|
if (
|
|
not dequantize_fp32
|
|
and dtype_dict[weights_dtype]["max"] <= 16384 # 1/fp16_min_normal
|
|
and not (
|
|
use_quantized_matmul
|
|
and not dtype_dict[quantized_matmul_dtype]["is_integer"]
|
|
and (not use_tensorwise_fp8_matmul or dtype_dict[quantized_matmul_dtype]["num_bits"] == 16)
|
|
)
|
|
):
|
|
scale_dtype = torch_dtype
|
|
|
|
if use_hadamard:
|
|
weight, use_hadamard, hadamard_group_size = apply_hadamard(weight, group_size=hadamard_group_size, hadamard=hadamard, layer_class_name=layer_class_name)
|
|
|
|
if use_svd:
|
|
try:
|
|
weight, svd_up, svd_down = apply_svdquant(weight, rank=svd_rank, niter=svd_steps, dtype=torch_dtype)
|
|
if use_quantized_matmul:
|
|
svd_up = svd_up.t_()
|
|
svd_down = svd_down.t_()
|
|
svd_up, svd_down = prepare_svd_for_matmul(svd_up, svd_down, use_quantized_matmul)
|
|
except Exception:
|
|
svd_up, svd_down = None, None
|
|
else:
|
|
svd_up, svd_down = None, None
|
|
|
|
if group_size == 0:
|
|
if use_quantized_matmul and not re_quantize_for_matmul and dtype_dict[weights_dtype]["num_bits"] >= 6:
|
|
group_size = -1
|
|
elif is_linear_type:
|
|
group_size = 2 ** ((3 if (svd_up is not None or using_pre_calculated_svd) else 2) + dtype_dict[weights_dtype]["num_bits"])
|
|
else:
|
|
group_size = 2 ** ((2 if (svd_up is not None or using_pre_calculated_svd) else 1) + dtype_dict[weights_dtype]["num_bits"])
|
|
|
|
if group_size > 0:
|
|
if group_size >= channel_size:
|
|
group_size = channel_size
|
|
num_of_groups = 1
|
|
else:
|
|
num_of_groups = channel_size // group_size
|
|
while num_of_groups * group_size != channel_size: # find something divisible
|
|
num_of_groups -= 1
|
|
if num_of_groups <= 1:
|
|
group_size = channel_size
|
|
num_of_groups = 1
|
|
break
|
|
group_size = channel_size // num_of_groups
|
|
group_size = int(group_size)
|
|
num_of_groups = int(num_of_groups)
|
|
|
|
if num_of_groups > 1:
|
|
if result_shape is None:
|
|
result_shape = weight.shape
|
|
if is_conv_type:
|
|
# output_channel_size, channel_size, X, X
|
|
# output_channel_size, num_of_groups, group_size, X, X
|
|
reduction_axes = 2
|
|
weight = weight.unflatten(1, (num_of_groups, group_size))
|
|
elif is_conv_transpose_type:
|
|
#channel_size, output_channel_size, X, X
|
|
#num_of_groups, group_size, output_channel_size, X, X
|
|
reduction_axes = 1
|
|
weight = weight.unflatten(1, (group_size, num_of_groups))
|
|
else:
|
|
# output_channel_size, channel_size
|
|
# output_channel_size, num_of_groups, group_size
|
|
reduction_axes = -1
|
|
weight = weight.unflatten(-1, (num_of_groups, group_size))
|
|
else:
|
|
group_size = -1
|
|
|
|
cast_scale = True
|
|
transpose_weights = False
|
|
re_quantize_for_matmul = re_quantize_for_matmul or num_of_groups > 1
|
|
if use_quantized_matmul and not re_quantize_for_matmul and not dtype_dict[weights_dtype]["is_packed"]:
|
|
transpose_weights = True
|
|
if not use_tensorwise_fp8_matmul and not dtype_dict[quantized_matmul_dtype]["is_integer"]:
|
|
cast_scale = False
|
|
|
|
weight, scale, zero_point = quantize_weight(weight, reduction_axes, weights_dtype, dtype=(scale_dtype if cast_scale else None), use_stochastic_rounding=(use_stochastic_rounding and not skip_sr))
|
|
|
|
if transpose_weights:
|
|
scale.t_()
|
|
weight.t_()
|
|
weight = prepare_weight_for_matmul(weight)
|
|
|
|
quantized_weight_shape = weight.shape
|
|
if dtype_dict[weights_dtype]["is_packed"]:
|
|
if dtype_dict[weights_dtype]["is_integer"]:
|
|
weight = pack_int(weight, weights_dtype)
|
|
else:
|
|
weight = pack_float(weight, weights_dtype)
|
|
else:
|
|
weight = weight.to(dtype=dtype_dict[weights_dtype]["torch_dtype"])
|
|
|
|
sdnq_dequantizer = SDNQDequantizer(
|
|
result_dtype=torch_dtype,
|
|
result_shape=result_shape,
|
|
original_shape=original_shape,
|
|
original_stride=original_stride,
|
|
quantized_weight_shape=quantized_weight_shape,
|
|
weights_dtype=weights_dtype,
|
|
quantized_matmul_dtype=quantized_matmul_dtype,
|
|
hadamard_group_size=hadamard_group_size,
|
|
group_size=group_size,
|
|
svd_rank=svd_rank,
|
|
svd_steps=svd_steps,
|
|
use_quantized_matmul=use_quantized_matmul,
|
|
re_quantize_for_matmul=re_quantize_for_matmul,
|
|
use_stochastic_rounding=use_stochastic_rounding,
|
|
use_hadamard=bool(use_hadamard or using_pre_rotated_hadamard),
|
|
layer_class_name=layer_class_name,
|
|
)
|
|
|
|
return sdnq_dequantizer, {"weight": weight, "scale": scale, "zero_point": zero_point, "svd_up": svd_up, "svd_down": svd_down}
|
|
|
|
|
|
@devices.inference_context()
|
|
def sdnq_quantize_layer_weight_dynamic(
|
|
weight: torch.FloatTensor,
|
|
layer_class_name: str | None = None,
|
|
weights_dtype: str = "uint4",
|
|
quantized_matmul_dtype: str | None = None,
|
|
group_size: int = 0,
|
|
hadamard_group_size: int = 256,
|
|
svd_rank: int = 32,
|
|
svd_steps: int = 8,
|
|
dynamic_loss_threshold: float | None =None,
|
|
use_svd: bool = False,
|
|
use_hadamard: bool = False,
|
|
use_quantized_matmul: bool = False,
|
|
use_stochastic_rounding: bool = False,
|
|
dequantize_fp32: bool = True,
|
|
hadamard: torch.FloatTensor | None = None,
|
|
param_name: str | None = None,
|
|
torch_dtype: torch.dtype | None = None,
|
|
):
|
|
if torch_dtype is None:
|
|
torch_dtype = weight.dtype
|
|
if dynamic_loss_threshold is None or dynamic_loss_threshold < 0:
|
|
dynamic_loss_threshold = 10 ** -(dtype_dict[weights_dtype]["num_bits"] / 2)
|
|
|
|
weight = weight.detach()
|
|
if weight.dtype != torch.float64:
|
|
weight = weight.to(dtype=torch.float32, copy=False)
|
|
weight_std = weight.std().square_().clamp_(min=1e-8)
|
|
original_weight_fp32 = weight
|
|
|
|
if use_hadamard:
|
|
weight, use_hadamard, hadamard_group_size = apply_hadamard(weight, group_size=hadamard_group_size, hadamard=hadamard, layer_class_name=layer_class_name)
|
|
|
|
if use_svd:
|
|
try:
|
|
weight, svd_up, svd_down = apply_svdquant(weight, rank=svd_rank, niter=svd_steps, dtype=torch_dtype)
|
|
svd_up, svd_down = prepare_svd_for_matmul(svd_up, svd_down, False)
|
|
if use_quantized_matmul:
|
|
svd_up_t, svd_down_t = svd_up.clone().t_(), svd_down.clone().t_()
|
|
svd_up_t, svd_down_t = prepare_svd_for_matmul(svd_up_t, svd_down_t, True)
|
|
else:
|
|
svd_up_t, svd_down_t = None, None
|
|
except Exception:
|
|
svd_up, svd_down = None, None
|
|
svd_up_t, svd_down_t = None, None
|
|
else:
|
|
svd_up, svd_down = None, None
|
|
svd_up_t, svd_down_t = None, None
|
|
|
|
quantization_loss = None
|
|
for i in range(weights_dtype_order.index(weights_dtype), len(weights_dtype_order)):
|
|
current_weights_dtype = weights_dtype_order[i]
|
|
if quantized_matmul_dtype is None and not is_fp8_mm_supported and not dtype_dict[current_weights_dtype]["is_integer"] and dtype_dict[current_weights_dtype]["num_bits"] < 16:
|
|
current_use_quantized_matmul = False
|
|
else:
|
|
current_use_quantized_matmul = use_quantized_matmul
|
|
|
|
sdnq_dequantizer, weight_data = sdnq_quantize_layer_weight(
|
|
weight,
|
|
layer_class_name=layer_class_name,
|
|
weights_dtype=current_weights_dtype,
|
|
quantized_matmul_dtype=quantized_matmul_dtype,
|
|
torch_dtype=torch_dtype,
|
|
hadamard_group_size=hadamard_group_size,
|
|
group_size=group_size,
|
|
svd_rank=svd_rank,
|
|
svd_steps=svd_steps,
|
|
use_svd=False,
|
|
use_hadamard=False,
|
|
using_pre_calculated_svd=use_svd,
|
|
using_pre_rotated_hadamard=use_hadamard,
|
|
use_quantized_matmul=current_use_quantized_matmul,
|
|
use_stochastic_rounding=use_stochastic_rounding,
|
|
dequantize_fp32=dequantize_fp32,
|
|
param_name=param_name,
|
|
)
|
|
|
|
if sdnq_dequantizer.use_quantized_matmul:
|
|
weight_data["svd_up"] = svd_up_t
|
|
weight_data["svd_down"] = svd_down_t
|
|
else:
|
|
weight_data["svd_up"] = svd_up
|
|
weight_data["svd_down"] = svd_down
|
|
|
|
quantization_loss = torch.nn.functional.mse_loss(
|
|
original_weight_fp32,
|
|
sdnq_dequantizer(
|
|
weight_data["weight"],
|
|
weight_data["scale"],
|
|
zero_point=weight_data["zero_point"],
|
|
svd_up=weight_data["svd_up"],
|
|
svd_down=weight_data["svd_down"],
|
|
skip_quantized_matmul=sdnq_dequantizer.use_quantized_matmul,
|
|
dtype=weight.dtype,
|
|
skip_compile=True,
|
|
),
|
|
).div_(weight_std)
|
|
if quantization_loss <= dynamic_loss_threshold:
|
|
del original_weight_fp32
|
|
return sdnq_dequantizer, weight_data
|
|
del original_weight_fp32
|
|
return None
|
|
|
|
|
|
@devices.inference_context()
|
|
def sdnq_quantize_layer(layer, quantization_config: "SDNQConfig", torch_dtype: torch.dtype | None = None, param_name: str = "", quant_kwargs: dict | None = None): # pylint: disable=unused-argument
|
|
if torch_dtype is None:
|
|
torch_dtype = layer.weight.dtype
|
|
if quant_kwargs is None:
|
|
quant_kwargs = get_quant_kwargs(layer, quantization_config, torch_dtype=torch_dtype, param_name=param_name)
|
|
|
|
layer_class_name = layer.__class__.__name__
|
|
if (
|
|
(layer_class_name in embedding_types and not quantization_config.quant_embedding)
|
|
or ((layer_class_name in conv_transpose_types or layer_class_name in conv_types) and not quantization_config.quant_conv)
|
|
):
|
|
quantization_config.modules_to_not_convert.append(param_name)
|
|
return layer, quantization_config
|
|
|
|
return_device = quant_kwargs.pop("return_device")
|
|
quantization_device = quant_kwargs.pop("quantization_device")
|
|
non_blocking = quant_kwargs.pop("non_blocking")
|
|
use_dynamic_quantization = quant_kwargs.pop("use_dynamic_quantization")
|
|
|
|
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, copy=False)
|
|
|
|
if use_dynamic_quantization:
|
|
weight_data = sdnq_quantize_layer_weight_dynamic(layer.weight, **quant_kwargs)
|
|
else:
|
|
weight_data = sdnq_quantize_layer_weight(layer.weight, **quant_kwargs)
|
|
|
|
if weight_data is not None:
|
|
layer.sdnq_dequantizer, weight_data = weight_data
|
|
layer = get_sdnq_wrapper_class(layer, get_forward_func(layer_class_name, layer.sdnq_dequantizer.quantized_matmul_dtype, layer.sdnq_dequantizer.use_quantized_matmul))
|
|
|
|
for key, value in weight_data.items():
|
|
if isinstance(value, (torch.Tensor, torch.nn.Parameter)):
|
|
setattr(layer, key, torch.nn.Parameter(value.to(return_device, non_blocking=non_blocking, copy=False), requires_grad=False))
|
|
setattr(getattr(layer, key), "_is_hf_initialized", True) # noqa: B010
|
|
else:
|
|
setattr(layer, key, value)
|
|
del weight_data
|
|
|
|
if use_dynamic_quantization:
|
|
if layer.sdnq_dequantizer.weights_dtype not in quantization_config.modules_dtype_dict.keys():
|
|
quantization_config.modules_dtype_dict[layer.sdnq_dequantizer.weights_dtype] = [param_name]
|
|
else:
|
|
quantization_config.modules_dtype_dict[layer.sdnq_dequantizer.weights_dtype].append(param_name)
|
|
|
|
if quant_kwargs["use_quantized_matmul"] and not layer.sdnq_dequantizer.use_quantized_matmul:
|
|
quantization_config.modules_to_not_use_matmul.append(param_name)
|
|
else:
|
|
layer.weight = torch.nn.Parameter(layer.weight.to(return_device, dtype=torch_dtype, non_blocking=non_blocking, copy=False), requires_grad=False)
|
|
if use_dynamic_quantization:
|
|
quantization_config.modules_to_not_convert.append(param_name)
|
|
|
|
return layer, quantization_config
|
|
|
|
|
|
@devices.inference_context()
|
|
def apply_sdnq_to_module(model, quantization_config: "SDNQConfig", torch_dtype: torch.dtype | None = None, full_param_name: str = ""): # pylint: disable=unused-argument
|
|
if not list(model.children()):
|
|
return model, quantization_config
|
|
for module_name, module in model.named_children():
|
|
if full_param_name:
|
|
param_name = full_param_name + "." + module_name
|
|
else:
|
|
param_name = module_name
|
|
if hasattr(module, "weight") and module.weight is not None:
|
|
param_name = param_name + ".weight"
|
|
layer_class_name = module.__class__.__name__
|
|
param_in_modules_to_not_convert = check_param_name_in(param_name, quantization_config.modules_to_not_convert)
|
|
if (
|
|
layer_class_name in allowed_types
|
|
and module.weight.dtype in {torch.float64, torch.float32, torch.float16, torch.bfloat16}
|
|
and param_in_modules_to_not_convert is None
|
|
and not (layer_class_name in embedding_types and not quantization_config.quant_embedding)
|
|
and not ((layer_class_name in conv_types or layer_class_name in conv_transpose_types) and not quantization_config.quant_conv)
|
|
):
|
|
module, quantization_config = sdnq_quantize_layer(module, quantization_config, torch_dtype=torch_dtype, param_name=param_name)
|
|
setattr(model, module_name, module)
|
|
elif param_in_modules_to_not_convert is None:
|
|
quantization_config.modules_to_not_convert.append(param_name)
|
|
|
|
module, quantization_config = apply_sdnq_to_module(module, quantization_config, torch_dtype=torch_dtype, full_param_name=param_name)
|
|
setattr(model, module_name, module)
|
|
return model, quantization_config
|
|
|
|
|
|
@devices.inference_context()
|
|
def sdnq_post_load_quant(
|
|
model: torch.nn.Module,
|
|
weights_dtype: str = "int8",
|
|
quantized_matmul_dtype: str | None = None,
|
|
hadamard_group_size: int = 256,
|
|
group_size: int = 0,
|
|
svd_rank: int = 32,
|
|
svd_steps: int = 8,
|
|
dynamic_loss_threshold: float | None = None,
|
|
use_svd: bool = False,
|
|
use_hadamard: bool = False,
|
|
quant_conv: bool = False,
|
|
quant_embedding: bool = False,
|
|
use_quantized_matmul: bool = False,
|
|
use_quantized_matmul_conv: bool = False,
|
|
use_dynamic_quantization: bool = False,
|
|
use_stochastic_rounding: bool = False,
|
|
dequantize_fp32: bool = True,
|
|
non_blocking: bool = False,
|
|
add_skip_keys:bool = True,
|
|
modules_to_not_convert: list[str] | None = None,
|
|
modules_to_not_use_matmul: list[str] | None = None,
|
|
modules_dtype_dict: dict[str, list[str]] | None = None,
|
|
modules_quant_config: dict[str, dict] | None = None,
|
|
quantization_device: torch.device | None = None,
|
|
return_device: torch.device | None = None,
|
|
torch_dtype: torch.dtype | None = None,
|
|
pre_quantized: bool = False,
|
|
):
|
|
if pre_quantized:
|
|
add_skip_keys = False
|
|
use_dynamic_quantization = False
|
|
elif (
|
|
hasattr(model, "quantization_config")
|
|
or (hasattr(model, "config") and hasattr(model.config, "quantization_config"))
|
|
or (hasattr(model, "config") and isinstance(model.config, dict) and "quantization_config" in model.config)
|
|
):
|
|
raise RuntimeError("Quantizing a pre-quantized model is not supported!")
|
|
|
|
quantization_config = SDNQConfig(
|
|
weights_dtype=weights_dtype,
|
|
quantized_matmul_dtype=quantized_matmul_dtype,
|
|
hadamard_group_size=hadamard_group_size,
|
|
group_size=group_size,
|
|
svd_rank=svd_rank,
|
|
svd_steps=svd_steps,
|
|
dynamic_loss_threshold=dynamic_loss_threshold,
|
|
use_svd=use_svd,
|
|
use_hadamard=use_hadamard,
|
|
quant_conv=quant_conv,
|
|
quant_embedding=quant_embedding,
|
|
use_quantized_matmul=use_quantized_matmul,
|
|
use_quantized_matmul_conv=use_quantized_matmul_conv,
|
|
use_dynamic_quantization=use_dynamic_quantization,
|
|
use_stochastic_rounding=use_stochastic_rounding,
|
|
dequantize_fp32=dequantize_fp32,
|
|
non_blocking=non_blocking,
|
|
add_skip_keys=add_skip_keys,
|
|
modules_to_not_convert=modules_to_not_convert,
|
|
modules_to_not_use_matmul=modules_to_not_use_matmul,
|
|
modules_dtype_dict=modules_dtype_dict,
|
|
modules_quant_config=modules_quant_config,
|
|
quantization_device=quantization_device,
|
|
return_device=return_device,
|
|
)
|
|
if add_skip_keys:
|
|
model, quantization_config = add_module_skip_keys(model, quantization_config)
|
|
|
|
model.eval()
|
|
model, quantization_config = apply_sdnq_to_module(model, quantization_config, torch_dtype=torch_dtype)
|
|
|
|
model.quantization_config = quantization_config
|
|
if hasattr(model, "config"):
|
|
try:
|
|
model.config.quantization_config = model.quantization_config
|
|
except Exception:
|
|
pass
|
|
try:
|
|
model.config["quantization_config"] = model.quantization_config.to_dict()
|
|
except Exception:
|
|
pass
|
|
model.quantization_method = QuantizationMethod.SDNQ
|
|
|
|
return model
|
|
|
|
|
|
class SDNQQuantize:
|
|
def __init__(self, hf_quantizer):
|
|
self.hf_quantizer = hf_quantizer
|
|
|
|
def convert(
|
|
self,
|
|
input_dict: dict[str, list[torch.Tensor]],
|
|
model: torch.nn.Module | None = None,
|
|
full_layer_name: str | None = None,
|
|
missing_keys: list[str] | None = None, # pylint: disable=unused-argument
|
|
**kwargs, # pylint: disable=unused-argument
|
|
) -> dict[str, torch.FloatTensor]:
|
|
_module_name, value = tuple(input_dict.items())[0]
|
|
value = value[0]
|
|
self.hf_quantizer.create_quantized_param(model, value, full_layer_name, value.device)
|
|
param, name = get_module_from_name(model, full_layer_name)
|
|
param = getattr(param, name)
|
|
return {full_layer_name: param}
|
|
|
|
@property
|
|
def reverse_op(self):
|
|
raise NotImplementedError
|
|
|
|
|
|
class SDNQQuantizer(DiffusersQuantizer, HfQuantizer):
|
|
r"""
|
|
Diffusers and Transformers Quantizer for SDNQ
|
|
"""
|
|
|
|
requires_parameters_quantization = True
|
|
use_keep_in_fp32_modules = True
|
|
requires_calibration = False
|
|
required_packages = None
|
|
torch_dtype = None
|
|
|
|
def __str__(self):
|
|
return f"SDNQQuantizer(torch_dtype={self.torch_dtype}, requires_parameters_quantization={self.requires_parameters_quantization}, use_keep_in_fp32_modules={self.use_keep_in_fp32_modules}, requires_calibration={self.requires_calibration}, required_packages={self.required_packages})"
|
|
|
|
def check_if_quantized_param(
|
|
self,
|
|
model: torch.nn.Module,
|
|
param_value: torch.Tensor, # pylint: disable=unused-argument
|
|
param_name: str,
|
|
*args, **kwargs, # pylint: disable=unused-argument
|
|
):
|
|
if self.pre_quantized:
|
|
layer, _tensor_name = get_module_from_name(model, param_name)
|
|
if hasattr(layer, "sdnq_dequantizer") and param_name.rsplit(".", maxsplit=1)[-1] in sdnq_keys:
|
|
return True
|
|
elif param_name.endswith(".weight"):
|
|
if check_param_name_in(param_name, self.quantization_config.modules_to_not_convert) is None:
|
|
layer_class_name = get_module_from_name(model, param_name)[0].__class__.__name__
|
|
if layer_class_name in allowed_types:
|
|
if layer_class_name in embedding_types:
|
|
if self.quantization_config.quant_embedding:
|
|
return True
|
|
elif layer_class_name in conv_types or layer_class_name in conv_transpose_types:
|
|
if self.quantization_config.quant_conv:
|
|
return True
|
|
else:
|
|
return True
|
|
self.quantization_config.modules_to_not_convert.append(param_name)
|
|
return False
|
|
|
|
@devices.inference_context()
|
|
def create_quantized_param( # pylint: disable=arguments-differ
|
|
self,
|
|
model: torch.nn.Module,
|
|
param_value: torch.FloatTensor,
|
|
param_name: str,
|
|
target_device: torch.device,
|
|
*args, **kwargs, # pylint: disable=unused-argument
|
|
):
|
|
layer, tensor_name = get_module_from_name(model, param_name)
|
|
|
|
if self.pre_quantized:
|
|
if param_value is not None:
|
|
if tensor_name == "weight":
|
|
return_dtype = param_value.dtype
|
|
elif self.quantization_config.dequantize_fp32 and tensor_name in sdnq_keys:
|
|
if param_value.dtype != torch.float64 and self.torch_dtype != torch.float64:
|
|
return_dtype = torch.float32
|
|
else:
|
|
return_dtype = torch.float64
|
|
else:
|
|
return_dtype = kwargs.get("dtype", param_value.dtype if self.torch_dtype is None else self.torch_dtype)
|
|
|
|
param_value = param_value.detach()
|
|
param_value = param_value.to(target_device, dtype=return_dtype, copy=False)
|
|
|
|
if tensor_name == "weight" and layer.sdnq_dequantizer.use_quantized_matmul and not layer.sdnq_dequantizer.re_quantize_for_matmul:
|
|
param_value = prepare_weight_for_matmul(param_value)
|
|
elif tensor_name == "svd_up":
|
|
param_value, _ = prepare_svd_for_matmul(param_value, None, layer.sdnq_dequantizer.use_quantized_matmul)
|
|
elif tensor_name == "svd_down":
|
|
_, param_value = prepare_svd_for_matmul(None, param_value, layer.sdnq_dequantizer.use_quantized_matmul)
|
|
|
|
param_value = torch.nn.Parameter(param_value, requires_grad=False)
|
|
param_value._is_hf_initialized = True # pylint: disable=protected-access
|
|
setattr(layer, tensor_name, param_value)
|
|
return
|
|
|
|
torch_dtype = kwargs.get("dtype", param_value.dtype if self.torch_dtype is None else self.torch_dtype)
|
|
quant_kwargs = get_quant_kwargs(layer, self.quantization_config, torch_dtype=torch_dtype, param_name=param_name)
|
|
if quant_kwargs["return_device"] is None:
|
|
quant_kwargs["return_device"] = target_device
|
|
if quant_kwargs["quantization_device"] is not None:
|
|
target_device = quant_kwargs["quantization_device"]
|
|
quant_kwargs["quantization_device"] = None
|
|
|
|
param_value = param_value.detach()
|
|
param_value = param_value.to(target_device, non_blocking=self.quantization_config.non_blocking, copy=False).to(dtype=torch.float32 if param_value.dtype != torch.float64 else torch.float64)
|
|
|
|
layer.weight = torch.nn.Parameter(param_value, requires_grad=False)
|
|
layer, self.quantization_config = sdnq_quantize_layer(layer, self.quantization_config, torch_dtype=torch_dtype, param_name=param_name, quant_kwargs=quant_kwargs) # pylint: disable=attribute-defined-outside-init
|
|
|
|
parent_module, tensor_name = get_module_from_name(model, param_name.removesuffix(tensor_name).removesuffix("."))
|
|
setattr(parent_module, tensor_name, layer)
|
|
|
|
def _process_model_before_weight_loading( # pylint: disable=arguments-differ
|
|
self,
|
|
model: torch.nn.Module,
|
|
device_map, # pylint: disable=unused-argument
|
|
keep_in_fp32_modules: list[str] | None = None,
|
|
**kwargs, # pylint: disable=unused-argument
|
|
):
|
|
if self.pre_quantized:
|
|
self.quantization_config.quantization_device = None
|
|
self.quantization_config.return_device = None
|
|
self.quantization_config.non_blocking = False
|
|
self.quantization_config.add_skip_keys = False
|
|
|
|
with init_empty_weights():
|
|
model = sdnq_post_load_quant(model, torch_dtype=self.torch_dtype, pre_quantized=True, **get_quant_args_from_config(self.quantization_config))
|
|
|
|
if self.quantization_config.add_skip_keys:
|
|
if keep_in_fp32_modules is not None:
|
|
self.quantization_config.modules_to_not_convert.extend(keep_in_fp32_modules)
|
|
if hasattr(self, "get_modules_to_not_convert") and hasattr(model, "tie_weights"):
|
|
self.quantization_config.modules_to_not_convert.extend(self.get_modules_to_not_convert(model, add_default_skips=True))
|
|
model, self.quantization_config = add_module_skip_keys(model, self.quantization_config) # pylint: disable=attribute-defined-outside-init
|
|
|
|
|
|
def _process_model_after_weight_loading(self, model, **kwargs): # pylint: disable=unused-argument
|
|
model.quantization_config = self.quantization_config
|
|
model.quantization_method = QuantizationMethod.SDNQ
|
|
if hasattr(model, "config"):
|
|
try:
|
|
model.config.quantization_config = self.quantization_config
|
|
except Exception:
|
|
pass
|
|
try:
|
|
model.config["quantization_config"] = self.quantization_config.to_dict()
|
|
except Exception:
|
|
pass
|
|
|
|
if self.pre_quantized:
|
|
from .loader import post_process_model
|
|
model = post_process_model(model)
|
|
if self.quantization_config.is_training:
|
|
from .training import convert_sdnq_model_to_training
|
|
model = convert_sdnq_model_to_training(
|
|
model,
|
|
dtype=self.torch_dtype,
|
|
quantized_matmul_dtype=self.quantization_config.quantized_matmul_dtype,
|
|
use_grad_ckpt=self.quantization_config.use_grad_ckpt,
|
|
use_quantized_matmul=self.quantization_config.use_quantized_matmul,
|
|
use_stochastic_rounding=self.quantization_config.use_stochastic_rounding,
|
|
dequantize_fp32=self.quantization_config.dequantize_fp32,
|
|
)
|
|
|
|
if shared.opts.diffusers_offload_mode != "none":
|
|
try:
|
|
model = model.to(device=devices.cpu)
|
|
except Exception:
|
|
model = model.to_empty(device=devices.cpu)
|
|
devices.torch_gc(force=True, reason="sdnq")
|
|
return model
|
|
|
|
def get_quantize_ops(self):
|
|
return SDNQQuantize(self)
|
|
|
|
def adjust_max_memory(self, max_memory: dict[str, int | str]) -> dict[str, int | str]:
|
|
max_memory = {key: val * 0.80 for key, val in max_memory.items()}
|
|
return max_memory
|
|
|
|
def adjust_target_dtype(self, target_dtype: torch.dtype) -> torch.dtype: # pylint: disable=unused-argument,arguments-renamed
|
|
return dtype_dict[self.quantization_config.weights_dtype]["target_dtype"]
|
|
|
|
def update_torch_dtype(self, torch_dtype: torch.dtype) -> torch.dtype:
|
|
self.torch_dtype = torch_dtype
|
|
return torch_dtype
|
|
|
|
def get_state_dict_and_metadata(self, state_dict: dict | torch.nn.Module, **kwargs) -> tuple[dict | None, dict]: # pylint: disable=unused-argument, arguments-differ
|
|
# transformers
|
|
if isinstance(state_dict, torch.nn.Module):
|
|
return None, {}
|
|
# diffusers
|
|
return state_dict, {}
|
|
|
|
def get_accelerator_warm_up_factor(self):
|
|
return 32 // dtype_dict[self.quantization_config.weights_dtype]["num_bits"]
|
|
|
|
def _dequantize(self, model):
|
|
return dequantize_sdnq_model(model)
|
|
|
|
def is_serializable(self, *args, **kwargs) -> bool: # pylint: disable=unused-argument, invalid-overridden-method
|
|
return not self.quantization_config.is_training
|
|
|
|
@property
|
|
def supports_safetensors_serialization(self) -> bool:
|
|
return self.is_serializable()
|
|
|
|
@property
|
|
def is_trainable(self):
|
|
return self.quantization_config.is_training
|
|
|
|
@property
|
|
def is_qat_trainable(self) -> bool:
|
|
return self.is_trainable()
|
|
|
|
@property
|
|
def is_compileable(self):
|
|
return True
|
|
|
|
def check_quantized_param(self, *args, **kwargs) -> bool:
|
|
"""
|
|
needed for transformers compatibility, returns self.check_if_quantized_param
|
|
"""
|
|
return self.check_if_quantized_param(*args, **kwargs)
|
|
|
|
def param_needs_quantization(self, model, param_name: str, *args, **kwargs) -> bool:
|
|
"""
|
|
needed for transformers compatibility, returns self.check_if_quantized_param
|
|
"""
|
|
return self.check_if_quantized_param(model, None, param_name, *args, **kwargs)
|
|
|
|
def get_cuda_warm_up_factor(self):
|
|
"""
|
|
needed for transformers compatibility, returns self.get_accelerator_warm_up_factor
|
|
"""
|
|
return self.get_accelerator_warm_up_factor()
|
|
|
|
def update_dtype(self, dtype: torch.dtype) -> torch.dtype:
|
|
"""
|
|
needed for transformers compatibility, returns self.update_torch_dtype
|
|
"""
|
|
return self.update_torch_dtype(dtype)
|
|
|
|
|
|
@dataclass
|
|
class SDNQConfig(QuantizationConfigMixin):
|
|
"""
|
|
This is a wrapper class about all possible attributes and features that you can play with a model that has been
|
|
loaded using `sdnq`.
|
|
|
|
Args:
|
|
weights_dtype (`str`, *optional*, defaults to `"int8"`):
|
|
The target dtype for the weights after quantization.
|
|
See `sdnq.common.accepted_weight_dtypes` for all the supported values.
|
|
These are some of the recommended values to use: ("int8", "int7", "int6", "uint5", "uint4", "uint3", "uint2", "float8_e4m3fn", "float7_e3m3fn", "float6_e3m2fn", "float5_e2m2fn", "float4_e2m1fn", "float3_e1m1fn", "float2_e1m0fn")
|
|
quantized_matmul_dtype (`str`, *optional*, defaults to `None`):
|
|
The target dtype for quantized matmul.
|
|
`None` will use "int8" with integer weight dtypes and "float8_e4m3fn" or "float16" with float weight dtypes.
|
|
Supported values are: ("int8", "float8_e4m3fn", "float16")
|
|
hadamard_group_size (`int`, *optional*, defaults to `0`):
|
|
Used to decide how many elements of a tensor will share the same Hadamard rotation group.
|
|
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.
|
|
group_size = `-1` will disable group sizes and use simple row-wise quantization instead.
|
|
svd_rank (`int`, *optional*, defaults to `32`):
|
|
The rank size used for the SVDQuant algorithm.
|
|
svd_steps (`int`, *optional*, defaults to `8`):
|
|
The number of iterations to use in svd lowrank estimation.
|
|
dynamic_loss_threshold (`float`, *optional*, defaults to `None`):
|
|
The target quantization mse loss threshold to use for dynamic quantization.
|
|
The value `None` or negative values means auto select a threshold based on the weights_dtype.
|
|
use_svd (`bool`, *optional*, defaults to `False`):
|
|
Enabling this option will use SVDQuant algorithm on top of SDNQ quantization.
|
|
use_hadamard (`bool`, *optional*, defaults to `False`):
|
|
Enabling this option will use Hadamard rotation on top of SDNQ quantization.
|
|
use_grad_ckpt (`bool`, *optional*, defaults to `True`):
|
|
This option is only used for training models when `is_training` is enabled or with `sdnq.training.sdnq_training_post_load_quant`.
|
|
Disabling this option will quantize the tensors needed for the backward pass before saving in the forward pass.
|
|
Disable this option only if you are not using gradient checkpointing and enable it when you are using gradient checkpointing.
|
|
quant_conv (`bool`, *optional*, defaults to `False`):
|
|
Enabling this option will quantize the nn.Conv layers in UNet models too.
|
|
Note: nn.Conv layers doesn't support training.
|
|
quant_embedding (`bool`, *optional*, defaults to `False`):
|
|
Enabling this option will quantize the nn.Embedding layers in text models too.
|
|
Note: nn.Embedding layers doesn't support training.
|
|
use_quantized_matmul (`bool`, *optional*, defaults to `False`):
|
|
Enabling this option will use quantized INT8, FP8 or FP16 MatMul on the forward pass and the backward pass instead of BF16 / FP16.
|
|
`False` means no quantized matmul and quantized_matmul_dtype will be ignored.
|
|
use_quantized_matmul_conv (`bool`, *optional*, defaults to `False`):
|
|
Same as use_quantized_matmul but for the nn.Conv layers with UNets like SDXL.
|
|
use_stochastic_rounding (`bool`, *optional*, defaults to `False`):
|
|
Enabling this option will use stochastic rounding on the quantization step.
|
|
Recommended for full finetuning / full rank training.
|
|
use_static_quantization (`bool`, *optional*, defaults to `True`):
|
|
This option is only used for training models with `sdnq.training.sdnq_training_post_load_quant`.
|
|
Disabling this option means model weights will be kept unquantized in their original precision and only the use_quantized_matmul (if enabled) will be used.
|
|
use_dynamic_quantization (`bool`, *optional*, defaults to `False`):
|
|
Enabling this option will dynamically select a per layer quantization type based on the dynamic_loss_threshold.
|
|
weights_dtype will be used as the minimum allowed quantization type when this option is enabled.
|
|
dequantize_fp32 (`bool`, *optional*, defaults to `True`):
|
|
Enabling this option will keep the quant scales in FP32 and compute the de-quant steps in FP32.
|
|
This option is highly recommended to be enabled.
|
|
non_blocking (`bool`, *optional*, defaults to `False`):
|
|
Enabling this option will use non blocking ops when moving layers between the quantization device and the return device.
|
|
add_skip_keys (`bool`, *optional*, defaults to `True`):
|
|
Disabling this option won't add model specific keys to modules_to_not_convert, modules_to_not_use_matmul and modules_dtype_dict.
|
|
quantization_device (`torch.device`, *optional*, defaults to `None`):
|
|
Used to set which device will be used for the quantization calculation on model load.
|
|
return_device (`torch.device`, *optional*, defaults to `None`):
|
|
Used to set which device will the quantized weights be sent back to.
|
|
modules_to_not_convert (`list`, *optional*, default to `None`):
|
|
The list of modules to not quantize. Useful for quantizing models that explicitly require to have some
|
|
modules left in their original precision (e.g. Whisper encoder, Llava encoder, Mixtral gate layers).
|
|
modules_to_not_use_matmul (`list`, *optional*, default to `None`):
|
|
The list of modules to not use quantized matmul. Useful for quantizing models that explicitly require to have some
|
|
modules left in their original activation precision (e.g. AdaLN layers).
|
|
modules_dtype_dict (`dict`, *optional*, default to `None`):
|
|
The dict of dtypes and list of modules. Useful for quantizing some modules with a different dtype.
|
|
modules_quant_config (`dict`, *optional*, default to `None`):
|
|
The dict of modules and a dict of quantization kwargs to use for that module.
|
|
Useful for quantizing some modules with a different quantization config.
|
|
is_training (`bool`, *optional*, defaults to `False`):
|
|
This option is auto enabled when using `sdnq.training.sdnq_training_post_load_quant`.
|
|
Enabling this option with SDNQQuantizer will convert the model to SDNQ Training format after quantization.
|
|
Note: Safetensors serialization is not supported with SDNQ training.
|
|
"""
|
|
|
|
def __init__( # pylint: disable=super-init-not-called
|
|
self,
|
|
weights_dtype: str = "int8",
|
|
quantized_matmul_dtype: str | None = None,
|
|
hadamard_group_size: int = 256,
|
|
group_size: int = 0,
|
|
svd_rank: int = 32,
|
|
svd_steps: int = 8,
|
|
dynamic_loss_threshold: float | None = None,
|
|
use_svd: bool = False,
|
|
use_hadamard: bool = False,
|
|
use_grad_ckpt: bool = True,
|
|
quant_conv: bool = False,
|
|
quant_embedding: bool = False,
|
|
use_quantized_matmul: bool = False,
|
|
use_quantized_matmul_conv: bool = False,
|
|
use_static_quantization: bool = True,
|
|
use_dynamic_quantization: bool = False,
|
|
use_stochastic_rounding: bool = False,
|
|
dequantize_fp32: bool = True,
|
|
non_blocking: bool = False,
|
|
add_skip_keys: bool = True,
|
|
quantization_device: torch.device | None = None,
|
|
return_device: torch.device | None = None,
|
|
modules_to_not_convert: list[str] | None = None,
|
|
modules_to_not_use_matmul: list[str] | None = None,
|
|
modules_dtype_dict: dict[str, list[str]] | None = None,
|
|
modules_quant_config: dict[str, dict] | None = None,
|
|
is_training: bool = False,
|
|
**kwargs, # pylint: disable=unused-argument
|
|
):
|
|
self.weights_dtype = weights_dtype
|
|
self.quantized_matmul_dtype = quantized_matmul_dtype
|
|
self.is_training = is_training
|
|
if self.is_training:
|
|
self.quant_method = QuantizationMethod.SDNQ_TRAINING
|
|
else:
|
|
self.quant_method = QuantizationMethod.SDNQ
|
|
self.hadamard_group_size = hadamard_group_size
|
|
self.group_size = group_size
|
|
self.svd_rank = svd_rank
|
|
self.dynamic_loss_threshold = dynamic_loss_threshold
|
|
self.svd_steps = svd_steps
|
|
self.use_svd = use_svd
|
|
self.use_hadamard = use_hadamard
|
|
self.use_grad_ckpt = use_grad_ckpt
|
|
self.quant_conv = quant_conv
|
|
self.quant_embedding = quant_embedding
|
|
self.use_quantized_matmul = use_quantized_matmul
|
|
self.use_quantized_matmul_conv = use_quantized_matmul_conv
|
|
self.use_static_quantization = use_static_quantization
|
|
self.use_dynamic_quantization = use_dynamic_quantization
|
|
self.use_stochastic_rounding = use_stochastic_rounding
|
|
self.dequantize_fp32 = dequantize_fp32
|
|
self.non_blocking = non_blocking
|
|
self.add_skip_keys = add_skip_keys
|
|
self.quantization_device = quantization_device
|
|
self.return_device = return_device
|
|
self.modules_to_not_convert = modules_to_not_convert
|
|
self.modules_to_not_use_matmul = modules_to_not_use_matmul
|
|
self.modules_dtype_dict = modules_dtype_dict
|
|
self.modules_quant_config = modules_quant_config
|
|
self.is_integer = dtype_dict[self.weights_dtype]["is_integer"]
|
|
self.sdnq_version = sdnq_version
|
|
self.post_init()
|
|
|
|
def post_init(self):
|
|
r"""
|
|
Safety checker that arguments are correct
|
|
"""
|
|
if self.use_quantized_matmul and not check_torch_compile():
|
|
shared.log.warning("SDNQ: Quantized MatMul requires a working Triton install for best performance.")
|
|
if self.weights_dtype not in accepted_weight_dtypes:
|
|
raise ValueError(f"SDNQ only support weight dtypes in {accepted_weight_dtypes} but found {self.weights_dtype}")
|
|
if self.quantized_matmul_dtype is not None and self.quantized_matmul_dtype not in accepted_matmul_dtypes:
|
|
raise ValueError(f"SDNQ only support quantized matmul dtypes in {accepted_matmul_dtypes} but found {self.quantized_matmul_dtype}")
|
|
|
|
if self.modules_to_not_convert is None:
|
|
self.modules_to_not_convert = []
|
|
elif isinstance(self.modules_to_not_convert, str):
|
|
self.modules_to_not_convert = [self.modules_to_not_convert]
|
|
elif isinstance(self.modules_to_not_convert, tuple):
|
|
self.modules_to_not_convert = list(self.modules_to_not_convert)
|
|
elif not isinstance(self.modules_to_not_convert, list):
|
|
raise ValueError(f"modules_to_not_convert must be a list but got {type(self.modules_to_not_convert)}")
|
|
|
|
if self.modules_to_not_use_matmul is None:
|
|
self.modules_to_not_use_matmul = []
|
|
elif isinstance(self.modules_to_not_use_matmul, str):
|
|
self.modules_to_not_use_matmul = [self.modules_to_not_use_matmul]
|
|
elif isinstance(self.modules_to_not_use_matmul, tuple):
|
|
self.modules_to_not_use_matmul = list(self.modules_to_not_use_matmul)
|
|
elif not isinstance(self.modules_to_not_use_matmul, list):
|
|
raise ValueError(f"modules_to_not_use_matmul must be a list but got {type(self.modules_to_not_use_matmul)}")
|
|
|
|
if self.modules_dtype_dict is None:
|
|
self.modules_dtype_dict = {}
|
|
elif not isinstance(self.modules_dtype_dict, dict):
|
|
raise ValueError(f"modules_dtype_dict must be a dict but got {type(self.modules_dtype_dict)}")
|
|
elif len(self.modules_dtype_dict.keys()) > 0:
|
|
self.modules_dtype_dict = self.modules_dtype_dict.copy()
|
|
for key, value in self.modules_dtype_dict.items():
|
|
if isinstance(value, str):
|
|
value = [value]
|
|
self.modules_dtype_dict[key] = value
|
|
elif isinstance(value, tuple):
|
|
value = list(value)
|
|
self.modules_dtype_dict[key] = value
|
|
if not isinstance(key, str) or not isinstance(value, list):
|
|
raise ValueError(f"modules_dtype_dict must be a dictionary of strings and lists but got {type(key)} and {type(value)}")
|
|
|
|
if self.modules_quant_config is None:
|
|
self.modules_quant_config = {}
|
|
|
|
self.modules_to_not_convert = self.modules_to_not_convert.copy()
|
|
self.modules_to_not_use_matmul = self.modules_to_not_use_matmul.copy()
|
|
self.modules_dtype_dict = self.modules_dtype_dict.copy()
|
|
self.modules_quant_config = self.modules_quant_config.copy()
|
|
|
|
# dedupe
|
|
self.modules_to_not_convert = list(set(self.modules_to_not_convert))
|
|
self.modules_to_not_use_matmul = list(set(self.modules_to_not_use_matmul))
|
|
for key, value in self.modules_dtype_dict.items():
|
|
self.modules_dtype_dict[key] = list(set(value))
|
|
|
|
def to_dict(self):
|
|
quantization_config_dict = self.__dict__.copy() # make serializable
|
|
quantization_config_dict["quantization_device"] = str(quantization_config_dict["quantization_device"]) if quantization_config_dict["quantization_device"] is not None else None
|
|
quantization_config_dict["return_device"] = str(quantization_config_dict["return_device"]) if quantization_config_dict["return_device"] is not None else None
|
|
return quantization_config_dict
|
|
|
|
def __str__(self):
|
|
return f"SDNQConfig(weights_dtype={self.weights_dtype} quantization_device={self.quantization_device} return_device={self.return_device} group_size={self.group_size} use_quantized_matmul={self.use_quantized_matmul} quantized_matmul_dtype={self.quantized_matmul_dtype} quant_conv={self.quant_conv} quant_embedding={self.quant_embedding} use_quantized_matmul_conv={self.use_quantized_matmul_conv} use_static_quantization={self.use_static_quantization} use_dynamic_quantization={self.use_dynamic_quantization} dynamic_loss_threshold={self.dynamic_loss_threshold} use_stochastic_rounding={self.use_stochastic_rounding} use_hadamard={self.use_hadamard} hadamard_group_size={self.hadamard_group_size} use_svd={self.use_svd} svd_rank={self.svd_rank} svd_steps={self.svd_steps} dequantize_fp32={self.dequantize_fp32} non_blocking={self.non_blocking} add_skip_keys={self.add_skip_keys} modules_to_not_convert={self.modules_to_not_convert} modules_to_not_use_matmul={self.modules_to_not_use_matmul} modules_dtype_dict={self.modules_dtype_dict} modules_quant_config={self.modules_quant_config} )"
|
|
|
|
|
|
import diffusers.quantizers.auto # noqa: E402,RUF100 # pylint: disable=wrong-import-order
|
|
diffusers.quantizers.auto.AUTO_QUANTIZER_MAPPING["sdnq"] = SDNQQuantizer
|
|
diffusers.quantizers.auto.AUTO_QUANTIZATION_CONFIG_MAPPING["sdnq"] = SDNQConfig
|
|
|
|
diffusers.quantizers.auto.AUTO_QUANTIZER_MAPPING["sdnq_training"] = SDNQQuantizer
|
|
diffusers.quantizers.auto.AUTO_QUANTIZATION_CONFIG_MAPPING["sdnq_training"] = SDNQConfig
|
|
|
|
import transformers.quantizers.auto # noqa: E402,RUF100 # pylint: disable=wrong-import-order
|
|
transformers.quantizers.auto.AUTO_QUANTIZER_MAPPING["sdnq"] = SDNQQuantizer
|
|
transformers.quantizers.auto.AUTO_QUANTIZATION_CONFIG_MAPPING["sdnq"] = SDNQConfig
|
|
|
|
transformers.quantizers.auto.AUTO_QUANTIZER_MAPPING["sdnq_training"] = SDNQQuantizer
|
|
transformers.quantizers.auto.AUTO_QUANTIZATION_CONFIG_MAPPING["sdnq_training"] = SDNQConfig
|
|
|
|
sdnq_quantize_layer_weight_compiled = compile_func(sdnq_quantize_layer_weight)
|