From 1433dfe3de23e88348a98840bb442b62f2b10709 Mon Sep 17 00:00:00 2001 From: Disty0 Date: Wed, 28 May 2025 03:16:29 +0300 Subject: [PATCH] SDNQ fix high RAM usage with pre mode --- CHANGELOG.md | 1 + modules/model_quant_sdnq.py | 33 +++++++++++++++++++++------------ 2 files changed, 22 insertions(+), 12 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 132f6f67c..1d4ae0edf 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -15,6 +15,7 @@ - Implement better layer hijacks - Fix Conv quant - Fix lora weight change + * Fix high RAM usage with pre mode - **IPEX** - Disabe Dynamic Attention by default on PyTorch 2.7 - Remove GradScaler hijack and use torch.amp.GradScaler instead diff --git a/modules/model_quant_sdnq.py b/modules/model_quant_sdnq.py index 0e4d1bff5..89ac6f60c 100644 --- a/modules/model_quant_sdnq.py +++ b/modules/model_quant_sdnq.py @@ -39,7 +39,7 @@ class QuantizationMethod(str, Enum): SDNQ = "sdnq" -def sdnq_quantize_layer(layer, weights_dtype="int8", torch_dtype=None, group_size=0, quant_conv=False, use_quantized_matmul=False, param_name=None): # pylint: disable=unused-argument +def sdnq_quantize_layer(layer, weights_dtype="int8", torch_dtype=None, group_size=0, quant_conv=False, use_quantized_matmul=False, param_name=None, pre_mode=False): # pylint: disable=unused-argument layer_class_name = layer.__class__.__name__ if layer_class_name in allowed_types: is_conv_type = False @@ -95,8 +95,10 @@ def sdnq_quantize_layer(layer, weights_dtype="int8", torch_dtype=None, group_siz layer.weight.requires_grad = False if shared.opts.diffusers_offload_mode in {"none", "model"}: return_device = devices.device - else: + elif pre_mode: return_device = devices.cpu + else: + return_device = layer.weight.device layer.weight.data = layer.weight.data.to(devices.device, dtype=torch.float32) if dtype_dict[weights_dtype]["is_unsigned"]: @@ -124,8 +126,9 @@ def sdnq_quantize_layer(layer, weights_dtype="int8", torch_dtype=None, group_siz result_shape=result_shape, use_quantized_matmul=use_quantized_matmul, weights_dtype=weights_dtype, - ).to(return_device) + ) layer.weight.data = layer.sdnq_decompressor.pack_weight(layer.weight.data).to(return_device) + layer.sdnq_decompressor = layer.sdnq_decompressor.to(return_device) if is_linear_type: if use_quantized_matmul: @@ -449,11 +452,16 @@ class SDNQQuantizer(DiffusersQuantizer): state_dict: Dict[str, Any], **kwargs, ): + param_value.data = param_value.data.clone() # safetensors is unable to release the cpu memory without this if param_name.endswith(".weight"): split_param_name = param_name.split(".") if param_name not in self.modules_to_not_convert and not any(param in split_param_name for param in self.modules_to_not_convert): - layer, _ = get_module_from_name(model, param_name) - return layer.__class__.__name__ in allowed_types + 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 conv_types or layer_class_name in conv_transpose_types: + return self.quantization_config.quant_conv + else: + return True return False def check_quantized_param(self, *args, **kwargs) -> bool: @@ -465,16 +473,16 @@ class SDNQQuantizer(DiffusersQuantizer): def create_quantized_param( # pylint: disable=arguments-differ self, model, - param_value: "torch.Tensor", + param_value: torch.FloatTensor, param_name: str, - target_device: "torch.device", + target_device: torch.device, # pylint: disable=unused-argument state_dict: Dict[str, Any], # pylint: disable=unused-argument unexpected_keys: List[str], # pylint: disable=unused-argument **kwargs, ): # load the model params to target_device first - layer, tensor_name = get_module_from_name(model, param_name) - layer.weight = torch.nn.Parameter(param_value.to(device=target_device), requires_grad=False) + layer, _ = get_module_from_name(model, param_name) + layer.weight = torch.nn.Parameter(param_value.to(device=devices.device, dtype=torch.float32), requires_grad=False) layer = sdnq_quantize_layer( layer, weights_dtype=self.quantization_config.weights_dtype, @@ -483,16 +491,17 @@ class SDNQQuantizer(DiffusersQuantizer): quant_conv=self.quantization_config.quant_conv, use_quantized_matmul=self.quantization_config.use_quantized_matmul, param_name=param_name, + pre_mode=True, ) def adjust_max_memory(self, max_memory: Dict[str, Union[int, str]]) -> Dict[str, Union[int, str]]: - max_memory = {key: val * 0.70 for key, val in max_memory.items()} + 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 + 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" = None) -> "torch.dtype": + def update_torch_dtype(self, torch_dtype: torch.dtype = None) -> torch.dtype: if torch_dtype is None: torch_dtype = devices.dtype self.torch_dtype = torch_dtype