mirror of
https://github.com/vladmandic/automatic
synced 2026-09-19 01:04:32 +02:00
SDNQ add sdnq_post_load_quant and update Qwen keys
This commit is contained in:
Disty0
+6
-37
@@ -479,7 +479,7 @@ def apply_layerwise(sd_model, quiet:bool=False):
|
||||
def sdnq_quantize_model(model, op=None, sd_model=None, do_gc: bool = True, weights_dtype: str = None, modules_to_not_convert: list = None, modules_dtype_dict: dict = None):
|
||||
global quant_last_model_name, quant_last_model_device # pylint: disable=global-statement
|
||||
from modules import devices, shared, timer
|
||||
from modules.sdnq import apply_sdnq_to_module
|
||||
from modules.sdnq import sdnq_post_load_quant
|
||||
|
||||
if weights_dtype is None:
|
||||
if op is not None and ("text_encoder" in op or op in {"TE", "LLM"}) and shared.opts.sdnq_quantize_weights_mode_te not in {"Same as model", "default"}:
|
||||
@@ -497,13 +497,6 @@ def sdnq_quantize_model(model, op=None, sd_model=None, do_gc: bool = True, weigh
|
||||
if modules_dtype_dict is None:
|
||||
modules_dtype_dict = {}
|
||||
|
||||
if getattr(model, "_keep_in_fp32_modules", None) is not None:
|
||||
modules_to_not_convert.extend(model._keep_in_fp32_modules) # pylint: disable=protected-access
|
||||
if getattr(model, "_skip_layerwise_casting_patterns", None) is not None:
|
||||
modules_to_not_convert.extend(model._skip_layerwise_casting_patterns) # pylint: disable=protected-access
|
||||
if model.__class__.__name__ == "ChromaTransformer2DModel":
|
||||
modules_to_not_convert.append("distilled_guidance_layer")
|
||||
|
||||
sdnq_modules_to_not_convert = [m.strip() for m in re.split(';|,| ', shared.opts.sdnq_modules_to_not_convert) if len(m.strip()) > 1]
|
||||
if len(sdnq_modules_to_not_convert) > 0:
|
||||
modules_to_not_convert.extend(sdnq_modules_to_not_convert)
|
||||
@@ -524,14 +517,9 @@ def sdnq_quantize_model(model, op=None, sd_model=None, do_gc: bool = True, weigh
|
||||
except Exception as e:
|
||||
log.warning(f'Quantization: SDNQ failed to parse sdnq_modules_dtype_dict: {e}')
|
||||
|
||||
model.eval()
|
||||
backup_embeddings = None
|
||||
if hasattr(model, "get_input_embeddings"):
|
||||
backup_embeddings = copy.deepcopy(model.get_input_embeddings())
|
||||
|
||||
t0 = time.time()
|
||||
|
||||
model = apply_sdnq_to_module(
|
||||
model = sdnq_post_load_quant(
|
||||
model,
|
||||
weights_dtype=weights_dtype,
|
||||
torch_dtype=devices.dtype,
|
||||
@@ -550,29 +538,8 @@ def sdnq_quantize_model(model, op=None, sd_model=None, do_gc: bool = True, weigh
|
||||
op=op,
|
||||
)
|
||||
|
||||
from modules.sdnq import SDNQConfig
|
||||
model.quantization_config = SDNQConfig(
|
||||
weights_dtype=weights_dtype,
|
||||
group_size=shared.opts.sdnq_quantize_weights_group_size,
|
||||
svd_rank=shared.opts.sdnq_svd_rank,
|
||||
use_svd=shared.opts.sdnq_use_svd,
|
||||
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,
|
||||
non_blocking=shared.opts.diffusers_offload_nonblocking,
|
||||
quantization_device=quantization_device,
|
||||
return_device=return_device,
|
||||
modules_to_not_convert=modules_to_not_convert,
|
||||
modules_dtype_dict=modules_dtype_dict.copy(),
|
||||
)
|
||||
|
||||
t1 = time.time()
|
||||
timer.load.add('sdnq', t1 - t0)
|
||||
model.quantization_method = 'SDNQ'
|
||||
|
||||
if hasattr(model, "set_input_embeddings") and backup_embeddings is not None:
|
||||
model.set_input_embeddings(backup_embeddings)
|
||||
|
||||
if op is not None and shared.opts.sdnq_quantize_shuffle_weights:
|
||||
if quant_last_model_name is not None:
|
||||
@@ -628,12 +595,14 @@ def optimum_quanto_model(model, op=None, sd_model=None, weights=None, activation
|
||||
from modules import devices, shared
|
||||
quanto = load_quanto('Quantize model: type=Optimum Quanto')
|
||||
global quant_last_model_name, quant_last_model_device # pylint: disable=global-statement
|
||||
if sd_model is not None and ("Flux" in sd_model.__class__.__name__ or "Chroma" in sd_model.__class__.__name__): # LayerNorm is not supported
|
||||
if model.__class__.__name__ in {"FluxTransformer2DModel", "ChromaTransformer2DModel"}: # LayerNorm is not supported
|
||||
exclude_list = ["transformer_blocks.*.norm1.norm", "transformer_blocks.*.norm2", "transformer_blocks.*.norm1_context.norm", "transformer_blocks.*.norm2_context", "single_transformer_blocks.*.norm.norm", "norm_out.norm"]
|
||||
if "Chroma" in sd_model.__class__.__name__:
|
||||
if model.__class__.__name__ == "ChromaTransformer2DModel":
|
||||
# we ignore the distilled guidance layer because it degrades quality too much
|
||||
# see: https://github.com/huggingface/diffusers/pull/11698#issuecomment-2969717180 for more details
|
||||
exclude_list.append("distilled_guidance_layer.*")
|
||||
elif model.__class__.__name__ == "QwenImageTransformer2DModel":
|
||||
exclude_list = ["transformer_blocks.0.img_mod.1.weight", "time_text_embed", "img_in", "txt_in", "proj_out", "norm_out", "pos_embed"]
|
||||
else:
|
||||
exclude_list = None
|
||||
weights = getattr(quanto, weights) if weights is not None else getattr(quanto, shared.opts.optimum_quanto_weights_type)
|
||||
|
||||
@@ -1,12 +1,14 @@
|
||||
from .quantizer import SDNQConfig, SDNQQuantizer, apply_sdnq_to_module, sdnq_quantize_layer
|
||||
from .quantizer import QuantizationMethod, SDNQConfig, SDNQQuantizer, sdnq_post_load_quant, apply_sdnq_to_module, sdnq_quantize_layer
|
||||
from .loader import save_sdnq_model, load_sdnq_model
|
||||
|
||||
|
||||
__all__ = [
|
||||
"QuantizationMethod",
|
||||
"SDNQConfig",
|
||||
"SDNQQuantizer",
|
||||
"apply_sdnq_to_module",
|
||||
"load_sdnq_model",
|
||||
"save_sdnq_model",
|
||||
"sdnq_post_load_quant",
|
||||
"sdnq_quantize_layer",
|
||||
]
|
||||
|
||||
@@ -5,7 +5,7 @@ from safetensors import safe_open
|
||||
from diffusers.models.modeling_utils import ModelMixin
|
||||
|
||||
from .common import use_tensorwise_fp8_matmul, use_contiguous_mm
|
||||
from .quantizer import SDNQConfig, apply_sdnq_to_module
|
||||
from .quantizer import SDNQConfig, sdnq_post_load_quant
|
||||
from .dequantizer import dequantize_symmetric, re_quantize_int8, re_quantize_fp8
|
||||
|
||||
|
||||
@@ -66,7 +66,7 @@ def load_sdnq_model(model_path: str, model_cls: ModelMixin = None, file_name: st
|
||||
model = model_cls._from_config(config) # pylint: disable=protected-access
|
||||
else:
|
||||
raise ValueError(f"Dont know how to load model for {model_cls}")
|
||||
model = apply_sdnq_to_module(model, **quantization_config)
|
||||
model = sdnq_post_load_quant(model, **quantization_config)
|
||||
|
||||
state_dict = {}
|
||||
if file_name:
|
||||
@@ -81,7 +81,6 @@ def load_sdnq_model(model_path: str, model_cls: ModelMixin = None, file_name: st
|
||||
for k in f.keys():
|
||||
state_dict[k] = f.get_tensor(k)
|
||||
model.load_state_dict(state_dict, assign=True)
|
||||
model.quantization_method = "sdnq"
|
||||
del state_dict
|
||||
if dtype is not None or dequantize_fp32 is not None or use_quantized_matmul is not None:
|
||||
model = apply_options_to_model(model, dtype=dtype, dequantize_fp32=dequantize_fp32, use_quantized_matmul=use_quantized_matmul)
|
||||
|
||||
@@ -75,7 +75,7 @@ def sdnq_quantize_layer(layer, weights_dtype="int8", torch_dtype=None, group_siz
|
||||
result_shape = None
|
||||
original_shape = layer.weight.shape
|
||||
if torch_dtype is None:
|
||||
torch_dtype = devices.dtype
|
||||
torch_dtype = layer.weight.dtype
|
||||
|
||||
if layer_class_name in conv_types:
|
||||
if not quant_conv:
|
||||
@@ -321,6 +321,66 @@ def apply_sdnq_to_module(model, weights_dtype="int8", torch_dtype=None, group_si
|
||||
return model
|
||||
|
||||
|
||||
def sdnq_post_load_quant(model, weights_dtype="int8", torch_dtype=None, group_size=0, svd_rank=32, use_svd=False, quant_conv=False, use_quantized_matmul=False, use_quantized_matmul_conv=False, dequantize_fp32=False, non_blocking=False, quantization_device=None, return_device=None, modules_to_not_convert: List[str] = None, modules_dtype_dict: Dict[str, List[str]] = None, op=None): # pylint: disable=unused-argument
|
||||
model.eval()
|
||||
if modules_to_not_convert is None:
|
||||
modules_to_not_convert = []
|
||||
if modules_dtype_dict is None:
|
||||
modules_dtype_dict = {}
|
||||
|
||||
if getattr(model, "_keep_in_fp32_modules", None) is not None:
|
||||
modules_to_not_convert.extend(model._keep_in_fp32_modules) # pylint: disable=protected-access
|
||||
if getattr(model, "_skip_layerwise_casting_patterns", None) is not None:
|
||||
modules_to_not_convert.extend(model._skip_layerwise_casting_patterns) # pylint: disable=protected-access
|
||||
if model.__class__.__name__ == "ChromaTransformer2DModel":
|
||||
modules_to_not_convert.append("distilled_guidance_layer")
|
||||
elif model.__class__.__name__ == "QwenImageTransformer2DModel":
|
||||
modules_to_not_convert.extend(["transformer_blocks.0.img_mod.1.weight", "time_text_embed", "img_in", "txt_in", "proj_out", "norm_out", "pos_embed"])
|
||||
if "minimum_6bit" not in modules_dtype_dict.keys():
|
||||
modules_dtype_dict["minimum_6bit"] = ["img_mod"]
|
||||
else:
|
||||
modules_dtype_dict["minimum_6bit"].append("img_mod")
|
||||
|
||||
model = apply_sdnq_to_module(
|
||||
model,
|
||||
weights_dtype=weights_dtype,
|
||||
torch_dtype=torch_dtype,
|
||||
group_size=group_size,
|
||||
svd_rank=svd_rank,
|
||||
use_svd=use_svd,
|
||||
quant_conv=quant_conv,
|
||||
use_quantized_matmul=use_quantized_matmul,
|
||||
use_quantized_matmul_conv=use_quantized_matmul_conv,
|
||||
dequantize_fp32=dequantize_fp32,
|
||||
non_blocking=non_blocking,
|
||||
quantization_device=quantization_device,
|
||||
return_device=return_device,
|
||||
modules_to_not_convert=modules_to_not_convert,
|
||||
modules_dtype_dict=modules_dtype_dict.copy(),
|
||||
op=op,
|
||||
)
|
||||
|
||||
model.quantization_config = SDNQConfig(
|
||||
weights_dtype=weights_dtype,
|
||||
group_size=group_size,
|
||||
svd_rank=svd_rank,
|
||||
use_svd=use_svd,
|
||||
quant_conv=quant_conv,
|
||||
use_quantized_matmul=use_quantized_matmul,
|
||||
use_quantized_matmul_conv=use_quantized_matmul_conv,
|
||||
dequantize_fp32=dequantize_fp32,
|
||||
non_blocking=non_blocking,
|
||||
quantization_device=quantization_device,
|
||||
return_device=return_device,
|
||||
modules_to_not_convert=modules_to_not_convert,
|
||||
modules_dtype_dict=modules_dtype_dict.copy(),
|
||||
)
|
||||
|
||||
model.quantization_method = QuantizationMethod.SDNQ
|
||||
|
||||
return model
|
||||
|
||||
|
||||
class SDNQQuantizer(DiffusersQuantizer, HfQuantizer):
|
||||
r"""
|
||||
Diffusers Quantizer for SDNQ
|
||||
@@ -384,6 +444,8 @@ class SDNQQuantizer(DiffusersQuantizer, HfQuantizer):
|
||||
*args, **kwargs, # pylint: disable=unused-argument
|
||||
):
|
||||
weights_dtype = self.quantization_config.weights_dtype
|
||||
torch_dtype = param_value.dtype if self.torch_dtype is None else self.torch_dtype
|
||||
|
||||
if len(self.quantization_config.modules_dtype_dict.keys()) > 0:
|
||||
split_param_name = param_name.split(".")
|
||||
for key, value in self.quantization_config.modules_dtype_dict.items():
|
||||
@@ -424,7 +486,7 @@ class SDNQQuantizer(DiffusersQuantizer, HfQuantizer):
|
||||
layer = sdnq_quantize_layer(
|
||||
layer,
|
||||
weights_dtype=weights_dtype,
|
||||
torch_dtype=self.torch_dtype,
|
||||
torch_dtype=torch_dtype,
|
||||
group_size=self.quantization_config.group_size,
|
||||
svd_rank=self.quantization_config.svd_rank,
|
||||
use_svd=self.quantization_config.use_svd,
|
||||
@@ -446,8 +508,6 @@ class SDNQQuantizer(DiffusersQuantizer, HfQuantizer):
|
||||
return dtype_dict[self.quantization_config.weights_dtype]["target_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
|
||||
return torch_dtype
|
||||
|
||||
|
||||
@@ -49,7 +49,8 @@ def load_qwen(checkpoint_info, diffusers_load_config={}):
|
||||
subfolder=transformer_subfolder,
|
||||
cls_name=diffusers.QwenImageTransformer2DModel,
|
||||
load_config=diffusers_load_config,
|
||||
modules_dtype_dict={"minimum_8bit": ["pos_embed", "time_text_embed", "img_in", "txt_in", "norm_out", "img_mod", "transformer_blocks.0.img_mod.1.weight"]},
|
||||
modules_dtype_dict={"minimum_6bit": ["img_mod"]},
|
||||
modules_to_not_convert=["transformer_blocks.0.img_mod.1.weight", "time_text_embed", "img_in", "txt_in", "proj_out", "norm_out", "pos_embed"],
|
||||
)
|
||||
|
||||
repo_te = 'Qwen/Qwen-Image'
|
||||
|
||||
Reference in New Issue
Block a user