From 3618e39cff8025020d8e727a8835e5de07a8a042 Mon Sep 17 00:00:00 2001 From: Disty0 Date: Tue, 27 May 2025 19:46:30 +0300 Subject: [PATCH] SDNQ use device_map = gpu --- modules/model_quant.py | 26 ++++++++++------- modules/model_quant_sdnq.py | 58 ++++++++++++++++++------------------- 2 files changed, 43 insertions(+), 41 deletions(-) diff --git a/modules/model_quant.py b/modules/model_quant.py index 05c9585b3..5a5dff44a 100644 --- a/modules/model_quant.py +++ b/modules/model_quant.py @@ -151,25 +151,25 @@ def check_nunchaku(module: str = ''): def create_config(kwargs = None, allow: bool = True, module: str = 'Model'): if kwargs is None: kwargs = {} + kwargs = create_sdnq_config(kwargs, allow_sdnq=allow, module=module) + if kwargs is not None and 'quantization_config' in kwargs: + if debug: + log.trace(f'Quantization: type=sdnq config={kwargs.get("quantization_config", None)}') + return kwargs kwargs = create_bnb_config(kwargs, allow_bnb=allow, module=module) if kwargs is not None and 'quantization_config' in kwargs: if debug: log.trace(f'Quantization: type=bnb config={kwargs.get("quantization_config", None)}') return kwargs - kwargs = create_ao_config(kwargs, allow_ao=allow, module=module) - if kwargs is not None and 'quantization_config' in kwargs: - if debug: - log.trace(f'Quantization: type=torchao config={kwargs.get("quantization_config", None)}') - return kwargs kwargs = create_quanto_config(kwargs, allow_quanto=allow, module=module) if kwargs is not None and 'quantization_config' in kwargs: if debug: log.trace(f'Quantization: type=quanto config={kwargs.get("quantization_config", None)}') return kwargs - kwargs = create_sdnq_config(kwargs, allow_sdnq=allow, module=module) + kwargs = create_ao_config(kwargs, allow_ao=allow, module=module) if kwargs is not None and 'quantization_config' in kwargs: if debug: - log.trace(f'Quantization: type=sdnq config={kwargs.get("quantization_config", None)}') + log.trace(f'Quantization: type=torchao config={kwargs.get("quantization_config", None)}') return kwargs return kwargs @@ -349,6 +349,8 @@ def sdnq_quantize_model(model, op=None, sd_model=None, do_gc=True): quant_last_model_name = None quant_last_model_device = None model.to(devices.device) + elif shared.opts.diffusers_offload_mode == "model": + model = model.to(devices.cpu) if do_gc: devices.torch_gc(force=True) return model @@ -521,12 +523,14 @@ def get_dit_args(load_config:dict={}, module:str=None, device_map:bool=False, al # if 'variant' in config: # del config['variant'] if device_map: - if shared.opts.device_map == 'cpu': - config['device_map'] = 'cpu' - if shared.opts.device_map == 'gpu': - config['device_map'] = devices.device if devices.backend == "ipex" and os.environ.get('UR_L0_ENABLE_RELAXED_ALLOCATION_LIMITS', '0') != '1' and module in {'TE', 'LLM'}: config['device_map'] = 'cpu' # alchemist gpus hits the 4GB allocation limit with transformers, UR_L0_ENABLE_RELAXED_ALLOCATION_LIMITS emulates above 4GB allocations + elif shared.opts.device_map == 'cpu': + config['device_map'] = 'cpu' + elif shared.opts.device_map == 'gpu': + config['device_map'] = devices.device + elif 'Model' in shared.opts.sdnq_quantize_weights or (module is not None and module in shared.opts.sdnq_quantize_weights) or module == 'any': + config['device_map'] = devices.device if allow_quant: quant_args = create_config(module=module) else: diff --git a/modules/model_quant_sdnq.py b/modules/model_quant_sdnq.py index 9a1be07ad..6d37774a0 100644 --- a/modules/model_quant_sdnq.py +++ b/modules/model_quant_sdnq.py @@ -92,10 +92,11 @@ def sdnq_quantize_layer(layer, weights_dtype="int8", torch_dtype=None, group_siz new_shape[last_dim_index - 1 : last_dim_index] = (int(num_of_groups), int(group_size)) layer.weight.data = layer.weight.reshape(new_shape) - if shared.opts.diffusers_offload_mode != "none": - return_device = layer.weight.device - else: + layer.weight.requires_grad = False + if shared.opts.diffusers_offload_mode in {"none", "model"}: return_device = devices.device + else: + return_device = devices.cpu layer.weight.data = layer.weight.data.to(devices.device, dtype=torch.float32) if dtype_dict[weights_dtype]["is_unsigned"]: @@ -103,7 +104,7 @@ def sdnq_quantize_layer(layer, weights_dtype="int8", torch_dtype=None, group_siz else: scale = get_int_scale_symmetric(layer.weight, reduction_axes, weights_dtype) zero_point = None - compressed_weight = quantize_weight(layer.weight, scale, zero_point, weights_dtype) + layer.weight.data = quantize_weight(layer.weight, scale, zero_point, weights_dtype) if not shared.opts.sdnq_decompress_fp32: scale = scale.to(torch_dtype) @@ -113,24 +114,17 @@ def sdnq_quantize_layer(layer, weights_dtype="int8", torch_dtype=None, group_siz if use_quantized_matmul: scale = scale.squeeze(-1) if dtype_dict[weights_dtype]["num_bits"] == 8: - compressed_weight = compressed_weight.transpose(0,1) + layer.weight.data = layer.weight.data.transpose(0,1) - decompressor = decompressor_dict[weights_dtype]( + layer.sdnq_decompressor = decompressor_dict[weights_dtype]( scale=scale, zero_point=zero_point, - compressed_weight_shape=compressed_weight.shape, + compressed_weight_shape=layer.weight.data.shape, result_dtype=torch_dtype, result_shape=result_shape, use_quantized_matmul=use_quantized_matmul, - ) - - compressed_weight = decompressor.pack_weight(compressed_weight).to(return_device) - decompressor = decompressor.to(return_device) - - layer.weight.requires_grad = False - layer.weight.data = compressed_weight - layer.sdnq_decompressor = decompressor - del compressed_weight, scale, zero_point + ).to(return_device) + layer.weight.data = layer.sdnq_decompressor.pack_weight(layer.weight.data).to(return_device) if is_linear_type: if use_quantized_matmul: @@ -457,7 +451,12 @@ class SDNQQuantizer(DiffusersQuantizer): state_dict: Dict[str, Any], **kwargs, ): - return param_name.endswith(".weight") + 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 + return False def check_quantized_param(self, *args, **kwargs) -> bool: """ @@ -477,19 +476,16 @@ class SDNQQuantizer(DiffusersQuantizer): ): # load the model params to target_device first layer, tensor_name = get_module_from_name(model, param_name) - layer._parameters[tensor_name] = torch.nn.Parameter(param_value).to(device=target_device) # pylint: disable=protected-access - - 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 = sdnq_quantize_layer( - layer, - weights_dtype=self.quantization_config.weights_dtype, - torch_dtype=self.torch_dtype, - group_size=self.quantization_config.group_size, - quant_conv=self.quantization_config.quant_conv, - use_quantized_matmul=self.quantization_config.use_quantized_matmul, - param_name=param_name, - ) + layer.weight = torch.nn.Parameter(param_value.to(device=target_device), requires_grad=False) + layer = sdnq_quantize_layer( + layer, + weights_dtype=self.quantization_config.weights_dtype, + torch_dtype=self.torch_dtype, + group_size=self.quantization_config.group_size, + quant_conv=self.quantization_config.quant_conv, + use_quantized_matmul=self.quantization_config.use_quantized_matmul, + param_name=param_name, + ) 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()} @@ -519,6 +515,8 @@ class SDNQQuantizer(DiffusersQuantizer): self.modules_to_not_convert.extend(keep_in_fp32_modules) def _process_model_after_weight_loading(self, model, **kwargs): + if shared.opts.diffusers_offload_mode == "model": + model = model.to(devices.cpu) devices.torch_gc(force=True) return model