Public/automatic
1
0
Fork 0
mirror of https://github.com/vladmandic/automatic synced 2026-09-19 17:24:32 +02:00
Code Issues Packages Projects Releases Wiki Activity

SDNQ use device_map = gpu

Browse Source
This commit is contained in:
Disty0
2025-05-27 19:46:30 +03:00
parent 73999ac710
commit 3618e39cff
2 changed files with 43 additions and 41 deletions
Download Patch File Download Diff File Expand all files Collapse all files
modules/model_quant.py
+15 -11
View File
@@ -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:
modules/model_quant_sdnq.py
+28 -30
View File
@@ -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