Cleanup SDNQ

This commit is contained in:
Disty0
2026-06-07 16:14:52 +03:00
parent 1f24513507
commit e5249c7a4b
2 changed files with 58 additions and 26 deletions
+8 -10
View File
@@ -74,15 +74,14 @@ def apply_svdquant(weight: torch.FloatTensor, rank: int = 32, niter: int = 8, dt
return weight, svd_up, svd_down
HADAMARD_N2_MATRIX = [[1, 1], [1, -1]]
@devices.inference_context()
def build_hadamard(n: int, dtype: torch.dtype | None = None, device: torch.device | None = None):
if n == 1:
return torch.ones((1, 1), dtype=dtype, device=device)
H = torch.tensor(HADAMARD_N2_MATRIX, dtype=dtype, device=device)
current_size = 2
H = H_N2 = torch.tensor([[1, 1], [1, -1]], dtype=dtype, device=device)
while current_size < n:
H = torch.kron(H, torch.tensor(HADAMARD_N2_MATRIX, dtype=dtype, device=device))
H = torch.kron(H, H_N2)
current_size *= 2
H = H.div_(n**0.5)
H = prepare_weight_for_matmul(H)
@@ -97,13 +96,12 @@ HADAMARD_MATRIX_CACHE = {}
def get_hadamard(n: int, dtype: torch.dtype | None = None, device: torch.device | None = None):
global HADAMARD_MATRIX_CACHE # pylint: disable=global-variable-not-assigned
device = devices.normalize_device(device)
if HADAMARD_MATRIX_CACHE.get(n, None) is None:
HADAMARD_MATRIX_CACHE[n] = {}
if HADAMARD_MATRIX_CACHE[n].get(device, None) is None:
HADAMARD_MATRIX_CACHE[n][device] = {}
if HADAMARD_MATRIX_CACHE[n][device].get(dtype, None) is None:
HADAMARD_MATRIX_CACHE[n][device][dtype] = build_hadamard(n, dtype=dtype, device=device)
return HADAMARD_MATRIX_CACHE[n][device][dtype]
H_key = (n, device, dtype)
H = HADAMARD_MATRIX_CACHE.get(H_key, None)
if H is None:
H = build_hadamard(n, dtype=dtype, device=device)
HADAMARD_MATRIX_CACHE[H_key] = H
return H
@devices.inference_context()
+50 -16
View File
@@ -46,7 +46,27 @@ class QuantizationMethod(str, Enum):
@devices.inference_context()
def sdnq_quantize_layer_weight(weight, layer_class_name=None, weights_dtype="int8", quantized_matmul_dtype=None, torch_dtype=None, group_size=0, hadamard_group_size=128, svd_rank=32, svd_steps=8, use_svd=False, use_hadamard=False, use_quantized_matmul=False, use_stochastic_rounding=False, dequantize_fp32=True, using_pre_calculated_svd=False, using_pre_rotated_hadamard=False, skip_sr=False, param_name=None): # pylint: disable=unused-argument
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 = 128,
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
@@ -54,9 +74,9 @@ def sdnq_quantize_layer_weight(weight, layer_class_name=None, weights_dtype="int
result_shape = None
scale_dtype = None
weight = weight.detach()
original_shape = weight.shape
original_stride = weight.stride()
weight = weight.detach()
if torch_dtype is None:
torch_dtype = weight.dtype
@@ -122,7 +142,7 @@ def sdnq_quantize_layer_weight(weight, layer_class_name=None, weights_dtype="int
scale_dtype = torch_dtype
if use_hadamard:
weight, use_hadamard, hadamard_group_size = apply_hadamard(weight, group_size=hadamard_group_size, layer_class_name=layer_class_name)
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:
@@ -228,20 +248,38 @@ def sdnq_quantize_layer_weight(weight, layer_class_name=None, weights_dtype="int
@devices.inference_context()
def sdnq_quantize_layer_weight_dynamic(weight, layer_class_name=None, weights_dtype="uint4", quantized_matmul_dtype=None, group_size=0, hadamard_group_size=128, svd_rank=32, svd_steps=8, dynamic_loss_threshold=None, use_svd=False, use_hadamard=False, use_quantized_matmul=False, use_stochastic_rounding=False, dequantize_fp32=True, quantization_config=None, torch_dtype=None, param_name=None): # pylint: disable=unused-argument
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 = 128,
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 = weight.detach()
original_weight_fp32 = weight.clone() if use_svd else weight
weight_std = original_weight_fp32.std().square_().clamp_(min=1e-8)
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, layer_class_name=layer_class_name)
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:
@@ -575,10 +613,8 @@ class SDNQQuantizer(DiffusersQuantizer, HfQuantizer):
else:
return_dtype = kwargs.get("dtype", param_value.dtype if self.torch_dtype is None else self.torch_dtype)
if param_value.dtype == return_dtype and devices.same_device(param_value.device, target_device):
param_value = param_value.detach()
else:
param_value = param_value.to(target_device, dtype=return_dtype, copy=False)
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)
@@ -600,10 +636,8 @@ class SDNQQuantizer(DiffusersQuantizer, HfQuantizer):
target_device = quant_kwargs["quantization_device"]
quant_kwargs["quantization_device"] = None
if param_value.dtype in {torch.float32, torch.float64} and devices.same_device(param_value.device, target_device):
param_value = param_value.detach()
else:
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)
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