diff --git a/modules/sdnq/dequantizer.py b/modules/sdnq/dequantizer.py index 38f09f0a1..0d7b0a83e 100644 --- a/modules/sdnq/dequantizer.py +++ b/modules/sdnq/dequantizer.py @@ -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 diff --git a/modules/sdnq/loader.py b/modules/sdnq/loader.py index 700c497e5..de02f0145 100644 --- a/modules/sdnq/loader.py +++ b/modules/sdnq/loader.py @@ -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