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

SDNQ add dtype casting to loader

Browse Source
This commit is contained in:
Disty0
2025-10-06 17:44:52 +03:00
parent 5c042c5fb8
commit c931bf9efa
2 changed files with 45 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files
modules/sdnq/dequantizer.py
+7 -5
View File
@@ -54,9 +54,11 @@ def quantize_int8(input: torch.FloatTensor, dim: int = -1) -> Tuple[torch.CharTe
return input, scale
def quantize_fp8(input: torch.FloatTensor, dim: int = -1) -> Tuple[torch.Tensor, torch.FloatTensor]:
scale = torch.amax(input.abs(), dim=dim, keepdims=True).div_(448)
input = torch.div(input, scale).nan_to_num_().clamp_(-448, 448).to(dtype=torch.float8_e4m3fn)
def quantize_fp8(input: torch.FloatTensor, dim: int = -1, is_e5: bool = False) -> Tuple[torch.Tensor, torch.FloatTensor]:
max_range = 57344 if is_e5 else 448
fp8_dtype = torch.float8_e5m2 if is_e5 else torch.float8_e4m3fn
scale = torch.amax(input.abs(), dim=dim, keepdims=True).div_(max_range)
input = torch.div(input, scale).nan_to_num_().clamp_(-max_range, max_range).to(dtype=fp8_dtype)
return input, scale
@@ -67,10 +69,10 @@ def re_quantize_int8(weight: torch.FloatTensor) -> Tuple[torch.CharTensor, torch
return weight, scale
def re_quantize_fp8(weight: torch.FloatTensor) -> Tuple[torch.CharTensor, torch.FloatTensor]:
def re_quantize_fp8(weight: torch.FloatTensor, is_e5: bool = False) -> Tuple[torch.CharTensor, torch.FloatTensor]:
if weight.ndim > 2: # convs
weight = weight.flatten(1,-1)
weight, scale = quantize_fp8(weight.t(), dim=0)
weight, scale = quantize_fp8(weight.t(), dim=0, is_e5=is_e5)
if not use_tensorwise_fp8_matmul:
scale = scale.to(dtype=torch.float32)
return weight, scale
modules/sdnq/loader.py
+38 -20
View File
@@ -6,7 +6,7 @@ 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 .dequantizer import dequantize_symmetric_compiled, re_quantize_int8, re_quantize_fp8
from .dequantizer import dequantize_symmetric, re_quantize_int8, re_quantize_fp8
def save_sdnq_model(model: ModelMixin, sdnq_config: SDNQConfig, model_path: str, max_shard_size: str = "10GB") -> None:
@@ -14,7 +14,7 @@ def save_sdnq_model(model: ModelMixin, sdnq_config: SDNQConfig, model_path: str,
sdnq_config.to_json_file(os.path.join(model_path, "quantization_config.json"))
def load_sdnq_model(model_cls: ModelMixin, model_path: str, file_name: str = "diffusion_pytorch_model.safetensors", use_quantized_matmul: bool = False) -> ModelMixin:
def load_sdnq_model(model_cls: ModelMixin, model_path: str, file_name: str = "diffusion_pytorch_model.safetensors", dtype: torch.dtype = None, dequantize_fp32: bool = None, use_quantized_matmul: bool = False) -> ModelMixin:
with torch.device("meta"):
with open(os.path.join(model_path, "quantization_config.json"), "r", encoding="utf-8") as f:
quantization_config = json.load(f)
@@ -33,37 +33,55 @@ def load_sdnq_model(model_cls: ModelMixin, model_path: str, file_name: str = "di
state_dict[k] = f.get_tensor(k)
model.load_state_dict(state_dict, assign=True)
del state_dict
if use_quantized_matmul and not quantization_config["use_quantized_matmul"]:
model = enable_quantized_mamtul(model)
model = apply_options_to_model(model, dtype=dtype, dequantize_fp32=dequantize_fp32, use_quantized_matmul=use_quantized_matmul)
return model
def enable_quantized_mamtul(model):
def apply_options_to_model(model, dtype: torch.dtype = None, dequantize_fp32: bool = None, use_quantized_matmul: bool = False):
has_children = list(model.children())
if not has_children:
return model
for module in model.children():
if hasattr(module, "sdnq_dequantizer"):
if not module.sdnq_dequantizer.use_quantized_matmul:
if module.sdnq_dequantizer.weights_dtype in {"int8", "float8_e4m3fn"}:
if module.sdnq_dequantizer.re_quantize_for_matmul:
return_dtype = module.scale.dtype
if dtype is not None:
module.sdnq_dequantizer.result_dtype = dtype
current_scale_dtype = module.svd_up.dtype if module.svd_up is not None else module.scale.dtype
scale_dtype = torch.float32 if dequantize_fp32 is None and current_scale_dtype == torch.float32 else torch.float32 if dequantize_fp32 else module.sdnq_dequantizer.result_dtype
upcast_scale = bool(use_quantized_matmul and use_tensorwise_fp8_matmul and module.sdnq_dequantizer.weights_dtype in {"float8_e4m3fn", "float8_e5m2"})
if upcast_scale:
module.scale.data = module.scale.to(dtype=torch.float32)
else:
module.scale.data = module.scale.to(dtype=scale_dtype)
if module.zero_point is not None:
module.zero_point.data = module.zero_point.to(dtype=scale_dtype)
if module.svd_up is not None:
module.svd_up.data = module.svd_up.to(dtype=scale_dtype)
module.svd_down.data = module.svd_down.to(dtype=scale_dtype)
if use_quantized_matmul != module.sdnq_dequantizer.use_quantized_matmul:
if module.sdnq_dequantizer.weights_dtype in {"int8", "float8_e4m3fn", "float8_e5m2"}:
if use_quantized_matmul and module.sdnq_dequantizer.re_quantize_for_matmul:
scale_dtype = module.scale.dtype
if module.sdnq_dequantizer.weights_dtype == "int8":
module.weight.data, module.scale.data = re_quantize_int8(dequantize_symmetric_compiled(module.weight, module.scale, torch.float32, module.sdnq_dequantizer.result_shape))
module.scale.data = module.scale.to(dtype=return_dtype)
module.weight.data, module.scale.data = re_quantize_int8(dequantize_symmetric(module.weight, module.scale, torch.float32, module.sdnq_dequantizer.result_shape))
module.scale.data = module.scale.to(dtype=scale_dtype)
else:
module.weight.data, module.scale.data = re_quantize_fp8(dequantize_symmetric_compiled(module.weight, module.scale, torch.float32, module.sdnq_dequantizer.result_shape))
is_e5 = bool(module.sdnq_dequantizer.weights_dtype == "float8_e5m2")
module.weight.data, module.scale.data = re_quantize_fp8(dequantize_symmetric(module.weight, module.scale, torch.float32, module.sdnq_dequantizer.result_shape), is_e5=is_e5)
if use_tensorwise_fp8_matmul:
module.scale.data = module.scale.to(dtype=return_dtype)
else:
module.scale.data = module.scale.to(dtype=scale_dtype)
elif not module.sdnq_dequantizer.re_quantize_for_matmul:
module.weight.data, module.scale.data = module.weight.t_(), module.scale.t_()
if use_contiguous_mm:
module.weight.data = module.weight.contiguous()
elif module.weight.is_contiguous():
module.weight.data = module.weight.t_().contiguous().t_()
if use_quantized_matmul:
if use_contiguous_mm:
module.weight.data = module.weight.contiguous()
elif module.weight.is_contiguous():
module.weight.data = module.weight.t_().contiguous().t_()
if module.svd_up is not None:
module.svd_up.data = module.svd_up.t_()
module.svd_down.data = module.svd_down.t_()
module.sdnq_dequantizer.use_quantized_matmul = True
module = enable_quantized_mamtul(module)
module.sdnq_dequantizer.use_quantized_matmul = use_quantized_matmul
module = apply_options_to_model(module, dtype=dtype, dequantize_fp32=dequantize_fp32, use_quantized_matmul=use_quantized_matmul)
return model