From 47dcab3522d3c57ed1eed43a35152f37ffbba6c6 Mon Sep 17 00:00:00 2001 From: Disty0 Date: Fri, 9 Jan 2026 00:34:32 +0300 Subject: [PATCH] update sdnq --- modules/sdnq/common.py | 7 +++++++ modules/sdnq/dequantizer.py | 27 +++++++++++++++++++++++---- modules/sdnq/loader.py | 2 +- modules/sdnq/quantizer.py | 14 ++++++++++---- 4 files changed, 41 insertions(+), 9 deletions(-) diff --git a/modules/sdnq/common.py b/modules/sdnq/common.py index 8660e49d8..b776a22f1 100644 --- a/modules/sdnq/common.py +++ b/modules/sdnq/common.py @@ -299,6 +299,13 @@ module_skip_keys_dict = { ["blocks.0.adaLN_modulation.1.weight", "x_embedder", "t_embedder", "y_embedder", "final_layer"], {} ], + "LTX2VideoTransformer3DModel": [ + [ + "audio_time_embed", "time_embed", "audio_caption_projection", "caption_projection", "proj_in", "audio_proj_in", "proj_out", "audio_proj_out", + "av_cross_attn_audio_scale_shift", "av_cross_attn_audio_v2a_gate", "av_cross_attn_video_a2v_gate", "av_cross_attn_video_scale_shift", + ], + {} + ], "Lumina2Transformer2DModel": [ ["layers.0.norm1.linear.weight", "time_caption_embed", "x_embedder", "norm_out"], {} diff --git a/modules/sdnq/dequantizer.py b/modules/sdnq/dequantizer.py index e058b04d1..076d01475 100644 --- a/modules/sdnq/dequantizer.py +++ b/modules/sdnq/dequantizer.py @@ -230,6 +230,25 @@ def dequantize_sdnq_model(model: torch.nn.Module): # SDNQDequantizer has to be a dataclass for torch.compile @dataclass class SDNQDequantizer: + result_dtype: torch.dtype + result_shape: torch.Size + original_shape: torch.Size + original_stride: List[int] + quantized_weight_shape: torch.Size + weights_dtype: str + quantized_matmul_dtype: str + group_size: int + svd_rank: int + svd_steps: int + use_quantized_matmul: bool + re_quantize_for_matmul: bool + use_stochastic_rounding: bool + layer_class_name: str + is_packed: bool + is_unsigned: bool + is_integer: bool + is_integer_matmul: bool + def __init__( self, result_dtype: torch.dtype, @@ -247,10 +266,6 @@ class SDNQDequantizer: use_stochastic_rounding: bool, layer_class_name: str, ): - self.is_packed = dtype_dict[weights_dtype]["is_packed"] - self.is_unsigned = dtype_dict[weights_dtype]["is_unsigned"] - self.is_integer = dtype_dict[weights_dtype]["is_integer"] - self.is_integer_matmul = dtype_dict[quantized_matmul_dtype]["is_integer"] self.result_dtype = result_dtype self.result_shape = result_shape self.original_shape = original_shape @@ -265,6 +280,10 @@ class SDNQDequantizer: self.re_quantize_for_matmul = re_quantize_for_matmul self.use_stochastic_rounding = use_stochastic_rounding self.layer_class_name = layer_class_name + self.is_packed = dtype_dict[weights_dtype]["is_packed"] + self.is_unsigned = dtype_dict[weights_dtype]["is_unsigned"] + self.is_integer = dtype_dict[weights_dtype]["is_integer"] + self.is_integer_matmul = dtype_dict[quantized_matmul_dtype]["is_integer"] @devices.inference_context() def re_quantize_matmul(self, weight, scale, zero_point, svd_up, svd_down): # pylint: disable=unused-argument diff --git a/modules/sdnq/loader.py b/modules/sdnq/loader.py index 7a9d60c29..f4ea338ae 100644 --- a/modules/sdnq/loader.py +++ b/modules/sdnq/loader.py @@ -25,7 +25,7 @@ def unset_config_on_save(quantization_config: SDNQConfig) -> SDNQConfig: return quantization_config -def save_sdnq_model(model: ModelMixin, model_path: str, max_shard_size: str = "10GB", is_pipeline: bool = False, sdnq_config: SDNQConfig = None) -> None: +def save_sdnq_model(model: ModelMixin, model_path: str, max_shard_size: str = "5GB", is_pipeline: bool = False, sdnq_config: SDNQConfig = None) -> None: if is_pipeline: for module_name in get_module_names(model): module = getattr(model, module_name, None) diff --git a/modules/sdnq/quantizer.py b/modules/sdnq/quantizer.py index 40da1b718..cb341ec88 100644 --- a/modules/sdnq/quantizer.py +++ b/modules/sdnq/quantizer.py @@ -1106,18 +1106,24 @@ class SDNQConfig(QuantizationConfigMixin): self.modules_dtype_dict = self.modules_dtype_dict.copy() def to_dict(self): - dct = self.__dict__.copy() # make serializable - dct["quantization_device"] = str(dct["quantization_device"]) if dct["quantization_device"] is not None else None - dct["return_device"] = str(dct["return_device"]) if dct["return_device"] is not None else None - return dct + 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 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)