update sdnq

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"],
         {}

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

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)

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)