mirror of
https://github.com/vladmandic/automatic
synced 2026-09-19 17:24:32 +02:00
Unify quant options
This commit is contained in:
Disty0
+103
-109
@@ -37,128 +37,122 @@ def get_quant(name):
|
||||
return 'none'
|
||||
|
||||
|
||||
def create_bnb_config(kwargs = None, allow_bnb: bool = True, module: str = 'Model', modules_to_not_convert: list = []):
|
||||
def create_bnb_config(kwargs = None, allow: bool = True, module: str = 'Model', modules_to_not_convert: list = []):
|
||||
from modules import shared, devices
|
||||
if len(shared.opts.bnb_quantization) > 0 and allow_bnb:
|
||||
if 'Model' in shared.opts.bnb_quantization or (module is not None and module in shared.opts.bnb_quantization) or module == 'any':
|
||||
load_bnb()
|
||||
if bnb is None:
|
||||
return kwargs
|
||||
bnb_config = diffusers.BitsAndBytesConfig(
|
||||
load_in_8bit=shared.opts.bnb_quantization_type in ['fp8'],
|
||||
load_in_4bit=shared.opts.bnb_quantization_type in ['nf4', 'fp4'],
|
||||
bnb_4bit_quant_storage=shared.opts.bnb_quantization_storage,
|
||||
bnb_4bit_quant_type=shared.opts.bnb_quantization_type,
|
||||
bnb_4bit_compute_dtype=devices.dtype,
|
||||
llm_int8_skip_modules=modules_to_not_convert,
|
||||
)
|
||||
log.debug(f'Quantization: module={module} type=bnb dtype={shared.opts.bnb_quantization_type} storage={shared.opts.bnb_quantization_storage}')
|
||||
if kwargs is None:
|
||||
return bnb_config
|
||||
else:
|
||||
kwargs['quantization_config'] = bnb_config
|
||||
return kwargs
|
||||
if allow and (module == 'any' or module in shared.opts.bnb_quantization):
|
||||
load_bnb()
|
||||
if bnb is None:
|
||||
return kwargs
|
||||
bnb_config = diffusers.BitsAndBytesConfig(
|
||||
load_in_8bit=shared.opts.bnb_quantization_type in ['fp8'],
|
||||
load_in_4bit=shared.opts.bnb_quantization_type in ['nf4', 'fp4'],
|
||||
bnb_4bit_quant_storage=shared.opts.bnb_quantization_storage,
|
||||
bnb_4bit_quant_type=shared.opts.bnb_quantization_type,
|
||||
bnb_4bit_compute_dtype=devices.dtype,
|
||||
llm_int8_skip_modules=modules_to_not_convert,
|
||||
)
|
||||
log.debug(f'Quantization: module={module} type=bnb dtype={shared.opts.bnb_quantization_type} storage={shared.opts.bnb_quantization_storage}')
|
||||
if kwargs is None:
|
||||
return bnb_config
|
||||
else:
|
||||
kwargs['quantization_config'] = bnb_config
|
||||
return kwargs
|
||||
return kwargs
|
||||
|
||||
|
||||
def create_ao_config(kwargs = None, allow_ao: bool = True, module: str = 'Model', modules_to_not_convert: list = []):
|
||||
def create_ao_config(kwargs = None, allow: bool = True, module: str = 'Model', modules_to_not_convert: list = []):
|
||||
from modules import shared
|
||||
if len(shared.opts.torchao_quantization) > 0 and (shared.opts.torchao_quantization_mode in {'pre', 'auto'}) and allow_ao:
|
||||
if 'Model' in shared.opts.torchao_quantization or (module is not None and module in shared.opts.torchao_quantization) or module == 'any':
|
||||
torchao = load_torchao()
|
||||
if torchao is None:
|
||||
return kwargs
|
||||
if module in {'TE', 'LLM'}:
|
||||
ao_config = transformers.TorchAoConfig(quant_type=shared.opts.torchao_quantization_type, modules_to_not_convert=modules_to_not_convert)
|
||||
else:
|
||||
ao_config = diffusers.TorchAoConfig(shared.opts.torchao_quantization_type, modules_to_not_convert=modules_to_not_convert)
|
||||
log.debug(f'Quantization: module={module} type=torchao dtype={shared.opts.torchao_quantization_type}')
|
||||
if kwargs is None:
|
||||
return ao_config
|
||||
else:
|
||||
kwargs['quantization_config'] = ao_config
|
||||
return kwargs
|
||||
if allow and (shared.opts.torchao_quantization_mode in {'pre', 'auto'}) and (module == 'any' or module in shared.opts.torchao_quantization):
|
||||
torchao = load_torchao()
|
||||
if torchao is None:
|
||||
return kwargs
|
||||
if module in {'TE', 'LLM'}:
|
||||
ao_config = transformers.TorchAoConfig(quant_type=shared.opts.torchao_quantization_type, modules_to_not_convert=modules_to_not_convert)
|
||||
else:
|
||||
ao_config = diffusers.TorchAoConfig(shared.opts.torchao_quantization_type, modules_to_not_convert=modules_to_not_convert)
|
||||
log.debug(f'Quantization: module={module} type=torchao dtype={shared.opts.torchao_quantization_type}')
|
||||
if kwargs is None:
|
||||
return ao_config
|
||||
else:
|
||||
kwargs['quantization_config'] = ao_config
|
||||
return kwargs
|
||||
return kwargs
|
||||
|
||||
|
||||
def create_quanto_config(kwargs = None, allow_quanto: bool = True, module: str = 'Model', modules_to_not_convert: list = []):
|
||||
def create_quanto_config(kwargs = None, allow: bool = True, module: str = 'Model', modules_to_not_convert: list = []):
|
||||
from modules import shared
|
||||
if len(shared.opts.quanto_quantization) > 0 and allow_quanto:
|
||||
if 'Model' in shared.opts.quanto_quantization or (module is not None and module in shared.opts.quanto_quantization) or module == 'any':
|
||||
load_quanto(silent=True)
|
||||
if optimum_quanto is None:
|
||||
return kwargs
|
||||
if module in {'TE', 'LLM'}:
|
||||
quanto_config = transformers.QuantoConfig(weights=shared.opts.quanto_quantization_type, modules_to_not_convert=modules_to_not_convert)
|
||||
quanto_config.weights_dtype = quanto_config.weights
|
||||
else:
|
||||
quanto_config = diffusers.QuantoConfig(weights_dtype=shared.opts.quanto_quantization_type, modules_to_not_convert=modules_to_not_convert)
|
||||
quanto_config.activations = None # patch so it works with transformers
|
||||
quanto_config.weights = quanto_config.weights_dtype
|
||||
log.debug(f'Quantization: module={module} type=quanto dtype={shared.opts.quanto_quantization_type}')
|
||||
if kwargs is None:
|
||||
return quanto_config
|
||||
else:
|
||||
kwargs['quantization_config'] = quanto_config
|
||||
return kwargs
|
||||
if allow and (module == 'any' or module in shared.opts.quanto_quantization):
|
||||
load_quanto(silent=True)
|
||||
if optimum_quanto is None:
|
||||
return kwargs
|
||||
if module in {'TE', 'LLM'}:
|
||||
quanto_config = transformers.QuantoConfig(weights=shared.opts.quanto_quantization_type, modules_to_not_convert=modules_to_not_convert)
|
||||
quanto_config.weights_dtype = quanto_config.weights
|
||||
else:
|
||||
quanto_config = diffusers.QuantoConfig(weights_dtype=shared.opts.quanto_quantization_type, modules_to_not_convert=modules_to_not_convert)
|
||||
quanto_config.activations = None # patch so it works with transformers
|
||||
quanto_config.weights = quanto_config.weights_dtype
|
||||
log.debug(f'Quantization: module={module} type=quanto dtype={shared.opts.quanto_quantization_type}')
|
||||
if kwargs is None:
|
||||
return quanto_config
|
||||
else:
|
||||
kwargs['quantization_config'] = quanto_config
|
||||
return kwargs
|
||||
return kwargs
|
||||
|
||||
|
||||
def create_sdnq_config(kwargs = None, allow_sdnq: bool = True, module: str = 'Model', weights_dtype: str = None, modules_to_not_convert: list = []):
|
||||
def create_sdnq_config(kwargs = None, allow: bool = True, module: str = 'Model', weights_dtype: str = None, modules_to_not_convert: list = []):
|
||||
from modules import devices, shared
|
||||
if len(shared.opts.sdnq_quantize_weights) > 0 and (shared.opts.sdnq_quantize_mode in {'pre', 'auto'}) and allow_sdnq:
|
||||
if 'Model' in shared.opts.sdnq_quantize_weights or (module is not None and module in shared.opts.sdnq_quantize_weights) or module == 'any':
|
||||
from modules.sdnq import SDNQQuantizer, SDNQConfig
|
||||
diffusers.quantizers.auto.AUTO_QUANTIZER_MAPPING["sdnq"] = SDNQQuantizer
|
||||
transformers.quantizers.auto.AUTO_QUANTIZER_MAPPING["sdnq"] = SDNQQuantizer
|
||||
diffusers.quantizers.auto.AUTO_QUANTIZATION_CONFIG_MAPPING["sdnq"] = SDNQConfig
|
||||
transformers.quantizers.auto.AUTO_QUANTIZATION_CONFIG_MAPPING["sdnq"] = SDNQConfig
|
||||
if allow and (shared.opts.sdnq_quantize_mode in {'pre', 'auto'}) and (module == 'any' or module in shared.opts.sdnq_quantize_weights):
|
||||
from modules.sdnq import SDNQQuantizer, SDNQConfig
|
||||
diffusers.quantizers.auto.AUTO_QUANTIZER_MAPPING["sdnq"] = SDNQQuantizer
|
||||
transformers.quantizers.auto.AUTO_QUANTIZER_MAPPING["sdnq"] = SDNQQuantizer
|
||||
diffusers.quantizers.auto.AUTO_QUANTIZATION_CONFIG_MAPPING["sdnq"] = SDNQConfig
|
||||
transformers.quantizers.auto.AUTO_QUANTIZATION_CONFIG_MAPPING["sdnq"] = SDNQConfig
|
||||
|
||||
if weights_dtype is None:
|
||||
if module in {"TE", "LLM"} and shared.opts.sdnq_quantize_weights_mode_te not in {"same as model", "default"}:
|
||||
weights_dtype = shared.opts.sdnq_quantize_weights_mode_te
|
||||
else:
|
||||
weights_dtype = shared.opts.sdnq_quantize_weights_mode
|
||||
if weights_dtype is None or weights_dtype == 'none':
|
||||
return kwargs
|
||||
|
||||
if shared.opts.device_map == "gpu":
|
||||
quantization_device = devices.device
|
||||
return_device = devices.device
|
||||
elif shared.opts.diffusers_offload_mode in {"none", "model"}:
|
||||
quantization_device = devices.device if shared.opts.sdnq_quantize_with_gpu else devices.cpu
|
||||
return_device = devices.device
|
||||
elif shared.opts.sdnq_quantize_with_gpu:
|
||||
quantization_device = devices.device
|
||||
return_device = devices.cpu
|
||||
if weights_dtype is None:
|
||||
if module in {"TE", "LLM"} and shared.opts.sdnq_quantize_weights_mode_te not in {"same as model", "default"}:
|
||||
weights_dtype = shared.opts.sdnq_quantize_weights_mode_te
|
||||
else:
|
||||
quantization_device = None
|
||||
return_device = None
|
||||
weights_dtype = shared.opts.sdnq_quantize_weights_mode
|
||||
if weights_dtype is None or weights_dtype == 'none':
|
||||
return kwargs
|
||||
|
||||
sdnq_config = SDNQConfig(
|
||||
weights_dtype=weights_dtype,
|
||||
group_size=shared.opts.sdnq_quantize_weights_group_size,
|
||||
quant_conv=shared.opts.sdnq_quantize_conv_layers,
|
||||
use_quantized_matmul=shared.opts.sdnq_use_quantized_matmul,
|
||||
use_quantized_matmul_conv=shared.opts.sdnq_use_quantized_matmul_conv,
|
||||
dequantize_fp32=shared.opts.sdnq_dequantize_fp32,
|
||||
quantization_device=quantization_device,
|
||||
return_device=return_device,
|
||||
modules_to_not_convert=modules_to_not_convert,
|
||||
)
|
||||
log.debug(f'Quantization: module="{module}" type=sdnq dtype={weights_dtype} matmul={shared.opts.sdnq_use_quantized_matmul} group_size={shared.opts.sdnq_quantize_weights_group_size} quant_conv={shared.opts.sdnq_quantize_conv_layers} matmul_conv={shared.opts.sdnq_use_quantized_matmul_conv} dequantize_fp32={shared.opts.sdnq_dequantize_fp32} quantize_with_gpu={shared.opts.sdnq_quantize_with_gpu} quantization_device={quantization_device} return_device={return_device}')
|
||||
if kwargs is None:
|
||||
return sdnq_config
|
||||
else:
|
||||
kwargs['quantization_config'] = sdnq_config
|
||||
return kwargs
|
||||
if shared.opts.device_map == "gpu":
|
||||
quantization_device = devices.device
|
||||
return_device = devices.device
|
||||
elif shared.opts.diffusers_offload_mode in {"none", "model"}:
|
||||
quantization_device = devices.device if shared.opts.sdnq_quantize_with_gpu else devices.cpu
|
||||
return_device = devices.device
|
||||
elif shared.opts.sdnq_quantize_with_gpu:
|
||||
quantization_device = devices.device
|
||||
return_device = devices.cpu
|
||||
else:
|
||||
quantization_device = None
|
||||
return_device = None
|
||||
|
||||
sdnq_config = SDNQConfig(
|
||||
weights_dtype=weights_dtype,
|
||||
group_size=shared.opts.sdnq_quantize_weights_group_size,
|
||||
quant_conv=shared.opts.sdnq_quantize_conv_layers,
|
||||
use_quantized_matmul=shared.opts.sdnq_use_quantized_matmul,
|
||||
use_quantized_matmul_conv=shared.opts.sdnq_use_quantized_matmul_conv,
|
||||
dequantize_fp32=shared.opts.sdnq_dequantize_fp32,
|
||||
quantization_device=quantization_device,
|
||||
return_device=return_device,
|
||||
modules_to_not_convert=modules_to_not_convert,
|
||||
)
|
||||
log.debug(f'Quantization: module="{module}" type=sdnq dtype={weights_dtype} matmul={shared.opts.sdnq_use_quantized_matmul} group_size={shared.opts.sdnq_quantize_weights_group_size} quant_conv={shared.opts.sdnq_quantize_conv_layers} matmul_conv={shared.opts.sdnq_use_quantized_matmul_conv} dequantize_fp32={shared.opts.sdnq_dequantize_fp32} quantize_with_gpu={shared.opts.sdnq_quantize_with_gpu} quantization_device={quantization_device} return_device={return_device}')
|
||||
if kwargs is None:
|
||||
return sdnq_config
|
||||
else:
|
||||
kwargs['quantization_config'] = sdnq_config
|
||||
return kwargs
|
||||
return kwargs
|
||||
|
||||
|
||||
def check_quant(module: str = ''):
|
||||
from modules import shared
|
||||
if 'Model' in shared.opts.bnb_quantization or 'Model' in shared.opts.torchao_quantization or 'Model' in shared.opts.quanto_quantization or 'Model' in shared.opts.sdnq_quantize_weights:
|
||||
return True
|
||||
if module in shared.opts.bnb_quantization or module in shared.opts.torchao_quantization or module in shared.opts.quanto_quantization or module in shared.opts.sdnq_quantize_weights:
|
||||
return True
|
||||
return False
|
||||
@@ -166,7 +160,7 @@ def check_quant(module: str = ''):
|
||||
|
||||
def check_nunchaku(module: str = ''):
|
||||
from modules import shared
|
||||
if 'Model' not in shared.opts.nunchaku_quantization and module not in shared.opts.nunchaku_quantization:
|
||||
if module not in shared.opts.nunchaku_quantization:
|
||||
return False
|
||||
from modules import mit_nunchaku
|
||||
mit_nunchaku.install_nunchaku()
|
||||
@@ -178,22 +172,22 @@ def check_nunchaku(module: str = ''):
|
||||
def create_config(kwargs = None, allow: bool = True, module: str = 'Model', modules_to_not_convert = []):
|
||||
if kwargs is None:
|
||||
kwargs = {}
|
||||
kwargs = create_sdnq_config(kwargs, allow_sdnq=allow, module=module, modules_to_not_convert=modules_to_not_convert)
|
||||
kwargs = create_sdnq_config(kwargs, allow=allow, module=module, modules_to_not_convert=modules_to_not_convert)
|
||||
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, modules_to_not_convert=modules_to_not_convert)
|
||||
kwargs = create_bnb_config(kwargs, allow=allow, module=module, modules_to_not_convert=modules_to_not_convert)
|
||||
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_quanto_config(kwargs, allow_quanto=allow, module=module, modules_to_not_convert=modules_to_not_convert)
|
||||
kwargs = create_quanto_config(kwargs, allow=allow, module=module, modules_to_not_convert=modules_to_not_convert)
|
||||
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_ao_config(kwargs, allow_ao=allow, module=module, modules_to_not_convert=modules_to_not_convert)
|
||||
kwargs = create_ao_config(kwargs, allow=allow, module=module, modules_to_not_convert=modules_to_not_convert)
|
||||
if kwargs is not None and 'quantization_config' in kwargs:
|
||||
if debug:
|
||||
log.trace(f'Quantization: type=torchao config={kwargs.get("quantization_config", None)}')
|
||||
@@ -309,13 +303,13 @@ def apply_layerwise(sd_model, quiet:bool=False):
|
||||
m.enable_layerwise_casting(compute_dtype=devices.dtype, storage_dtype=storage_dtype, non_blocking=non_blocking)
|
||||
m.quantization_method = 'LayerWise'
|
||||
log.quiet(quiet, f'Quantization: type=layerwise module={module} cls={cls} storage={storage_dtype} compute={devices.dtype} blocking={not non_blocking}')
|
||||
if module.startswith('transformer') and ('Model' in shared.opts.layerwise_quantization or 'Transformer' in shared.opts.layerwise_quantization):
|
||||
if module.startswith('transformer') and ('Model' in shared.opts.layerwise_quantization):
|
||||
m = getattr(sd_model, module)
|
||||
if hasattr(m, 'enable_layerwise_casting'):
|
||||
m.enable_layerwise_casting(compute_dtype=devices.dtype, storage_dtype=storage_dtype, non_blocking=non_blocking)
|
||||
m.quantization_method = 'LayerWise'
|
||||
log.quiet(quiet, f'Quantization: type=layerwise module={module} cls={cls} storage={storage_dtype} compute={devices.dtype} blocking={not non_blocking}')
|
||||
if module.startswith('text_encoder') and ('Model' in shared.opts.layerwise_quantization or 'TE' in shared.opts.layerwise_quantization) and ('clip' not in cls.lower()):
|
||||
if module.startswith('text_encoder') and ('TE' in shared.opts.layerwise_quantization) and ('clip' not in cls.lower()):
|
||||
m = getattr(sd_model, module)
|
||||
if hasattr(m, 'enable_layerwise_casting'):
|
||||
m.enable_layerwise_casting(compute_dtype=devices.dtype, storage_dtype=storage_dtype, non_blocking=non_blocking)
|
||||
@@ -575,7 +569,7 @@ 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 devices.backend == "ipex" and os.environ.get('UR_L0_ENABLE_RELAXED_ALLOCATION_LIMITS', '0') != '1' and module in {'TE', 'LLM'}:
|
||||
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'
|
||||
|
||||
Reference in New Issue
Block a user