diff --git a/modules/sdnq/common.py b/modules/sdnq/common.py index 021457ffa..1a0e038f4 100644 --- a/modules/sdnq/common.py +++ b/modules/sdnq/common.py @@ -380,8 +380,9 @@ int_mm_func = None if use_openvino_mm: try: - from .kernels.openvino_mm import openvino_int_mm + from .kernels.openvino_mm import openvino_int_mm, openvino_fp_mm int_mm_func = openvino_int_mm + fp_mm_func = openvino_fp_mm except Exception: use_openvino_mm = False elif use_triton_mm: diff --git a/modules/sdnq/kernels/openvino_mm.py b/modules/sdnq/kernels/openvino_mm.py index 25ae94b0a..5190266cc 100644 --- a/modules/sdnq/kernels/openvino_mm.py +++ b/modules/sdnq/kernels/openvino_mm.py @@ -6,35 +6,36 @@ from openvino.properties import hint as ov_hints core = ov.Core() -NPU_MUL = 32 # NPU uses FP16 x INT8 -> FP16 instead of INT8 x INT8 -> INT32 and FP16 output overflows OV_DEVICE: str = os.environ.get("SDNQ_OPENVINO_DEVICE", "CPU") -OV_COMPILED_CACHE: dict[tuple[str, tuple[int,int] | None, tuple[int,int] | None], tuple[ov.InferRequest, str]] = {} -core.set_property(OV_DEVICE, {ov_hints.execution_mode: ov_hints.ExecutionMode.ACCURACY}) +OV_COMPILED_CACHE: dict[tuple[str, tuple[int,int] | None, str, tuple[int,int] | None], tuple[ov.InferRequest, str]] = {} + +if OV_DEVICE == "NPU": + OV_DEVICE = "HETERO:NPU,CPU" +for ov_device in core.get_available_devices(): + core.set_property(ov_device, {ov_hints.execution_mode: ov_hints.ExecutionMode.ACCURACY}) -def ov_int_mm(A: torch.CharTensor, B: torch.CharTensor, infer_request: ov.InferRequest, out_name: str) -> torch.FloatTensor: +def ov_mm(A: torch.CharTensor, B: torch.CharTensor, infer_request: ov.InferRequest, out_name: str) -> torch.FloatTensor: C = torch.empty((A.shape[0], B.shape[-1]), device="cpu", dtype=torch.float32) infer_request.set_tensor("A", ov.Tensor(A.detach().contiguous().to("cpu").numpy(), shared_memory=True)) infer_request.set_tensor("B", ov.Tensor(B.detach().contiguous().to("cpu").numpy(), shared_memory=True)) infer_request.set_tensor(out_name, ov.Tensor(C.numpy(), shared_memory=True)) infer_request.infer() C = C.to(A.device) - if OV_DEVICE == "NPU": - C.mul_(NPU_MUL**2) return C @torch.library.custom_op("sdnq::openvino_int_mm", mutates_args=()) def openvino_int_mm(Tensor_A: torch.Tensor, Tensor_B: torch.Tensor) -> torch.Tensor: - if OV_DEVICE in {"NPU", "CPU"}: - cache_key = (OV_DEVICE, Tensor_A.shape, Tensor_B.shape) + if "GPU" not in OV_DEVICE: + cache_key = (OV_DEVICE, "int8", Tensor_A.shape, Tensor_B.shape) else: - cache_key = (OV_DEVICE, None, None) + cache_key = (OV_DEVICE, "int8", None, None) infer_request, out_name = OV_COMPILED_CACHE.get(cache_key, (None, None)) if infer_request is not None: - return ov_int_mm(Tensor_A, Tensor_B, infer_request, out_name) + return ov_mm(Tensor_A, Tensor_B, infer_request, out_name) - if OV_DEVICE in {"NPU", "CPU"}: + if "GPU" not in OV_DEVICE: shape_a = ov.Shape(Tensor_A.shape) shape_b = ov.Shape(Tensor_B.shape) else: @@ -42,23 +43,100 @@ def openvino_int_mm(Tensor_A: torch.Tensor, Tensor_B: torch.Tensor) -> torch.Ten shape_b = ov.PartialShape([-1,-1]) input_a = ov_ops.parameter(shape_a, ov.Type.i8, name="A") input_b = ov_ops.parameter(shape_b, ov.Type.i8, name="B") + a = ov_ops.convert(input_a, ov.Type.f32) + b = ov_ops.convert(input_b, ov.Type.f32) + low = ov_ops.constant(-128.0, dtype=ov.Type.f32) high = ov_ops.constant(127.0, dtype=ov.Type.f32) + a = ov_ops.fake_quantize(a, low, high, low, high, 256) + b = ov_ops.fake_quantize(b, low, high, low, high, 256) - a = ov_ops.fake_quantize(ov_ops.convert(input_a, ov.Type.f32), low, high, low, high, 256) - b = ov_ops.fake_quantize(ov_ops.convert(input_b, ov.Type.f32), low, high, low, high, 256) - if OV_DEVICE == "NPU": - a = ov_ops.divide(a, ov_ops.constant(NPU_MUL, dtype=ov.Type.f32)) - b = ov_ops.divide(b, ov_ops.constant(NPU_MUL, dtype=ov.Type.f32)) + # NPU uses FP16 x INT8 -> FP16 instead of INT8 x INT8 -> INT32 and FP16 output overflows + if "NPU" in OV_DEVICE: + fp16_scale = 0.25012213 * Tensor_B.shape[-2] + in_scale = ov_ops.constant(fp16_scale ** 0.5, dtype=ov.Type.f32) + out_scale = ov_ops.constant(fp16_scale, dtype=ov.Type.f32, name="out_scale_const") + a = ov_ops.divide(a, in_scale) + b = ov_ops.divide(b, in_scale) + out = ov_ops.matmul(a, b, False, False) + out = ov_ops.multiply(out, out_scale, name="out_scale") + else: + out = ov_ops.matmul(a, b, False, False) - ov_model = ov.Model([ov_ops.matmul(a, b, False, False)], [input_a, input_b], "ov_int8_mm") + ov_model = ov.Model([out], [input_a, input_b], "ov_int8_mm") + if "NPU" in OV_DEVICE: # NPU can't use FP32 for regular multiplications + for node in ov_model.get_ops(): + if node.get_friendly_name() in {"out_scale", "out_scale_const"}: + node.get_rt_info()["affinity"] = "CPU" + else: + node.get_rt_info()["affinity"] = "NPU" ov_model = core.compile_model(ov_model, OV_DEVICE) infer_request = ov_model.create_infer_request() out_name = ov_model.outputs[0] OV_COMPILED_CACHE[cache_key] = (infer_request, out_name) - return ov_int_mm(Tensor_A, Tensor_B, infer_request, out_name) + return ov_mm(Tensor_A, Tensor_B, infer_request, out_name) @openvino_int_mm.register_fake def openvino_int_mm_fake(A: torch.Tensor, B: torch.Tensor) -> torch.Tensor: return torch.mm(A.to(dtype=torch.float32), B.to(dtype=torch.float32)) + + +@torch.library.custom_op("sdnq::openvino_fp_mm", mutates_args=()) +def openvino_fp_mm(Tensor_A: torch.Tensor, Tensor_B: torch.Tensor) -> torch.Tensor: + mm_dtype = "fp16" if Tensor_B.dtype == torch.float16 else "fp8" + if mm_dtype == "fp8": + Tensor_A = Tensor_A.to(dtype=torch.float16) + Tensor_B = Tensor_B.to(dtype=torch.float16) + if "GPU" not in OV_DEVICE: + cache_key = (OV_DEVICE, mm_dtype, Tensor_A.shape, Tensor_B.shape) + else: + cache_key = (OV_DEVICE, mm_dtype, None, None) + infer_request, out_name = OV_COMPILED_CACHE.get(cache_key, (None, None)) + if infer_request is not None: + return ov_mm(Tensor_A, Tensor_B, infer_request, out_name) + + if "GPU" not in OV_DEVICE: + shape_a = ov.Shape(Tensor_A.shape) + shape_b = ov.Shape(Tensor_B.shape) + else: + shape_a = ov.PartialShape([-1,-1]) + shape_b = ov.PartialShape([-1,-1]) + input_a = ov_ops.parameter(shape_a, ov.Type.f16, name="A") + input_b = ov_ops.parameter(shape_b, ov.Type.f16, name="B") + a = ov_ops.convert(input_a, ov.Type.f32) + b = ov_ops.convert(input_b, ov.Type.f32) + + if mm_dtype == "fp8": + low = ov_ops.constant(-448.0, dtype=ov.Type.f32) + high = ov_ops.constant(448.0, dtype=ov.Type.f32) + a = ov_ops.fake_quantize(a, low, high, low, high, 256) + b = ov_ops.fake_quantize(b, low, high, low, high, 256) + fp16_scale = 4 * Tensor_B.shape[-2] + else: + fp16_scale = 65536 * Tensor_B.shape[-2] + + in_scale = ov_ops.constant(fp16_scale**0.5, dtype=ov.Type.f32) + out_scale = ov_ops.constant(fp16_scale, dtype=ov.Type.f32, name="out_scale_const") + a = ov_ops.convert(ov_ops.divide(a, in_scale), ov.Type.f16) + b = ov_ops.convert(ov_ops.divide(b, in_scale), ov.Type.f16) + out = ov_ops.matmul(a, b, False, False, name="fp_mm") + out = ov_ops.multiply(ov_ops.convert(out, ov.Type.f32), out_scale, name="out_scale") + + ov_model = ov.Model([out], [input_a, input_b], "ov_fp_mm") + if "NPU" in OV_DEVICE: # NPU can't use FP32 for regular multiplications + for node in ov_model.get_ops(): + if node.get_friendly_name() in {"out_scale", "out_scale_const"}: + node.get_rt_info()["affinity"] = "CPU" + else: + node.get_rt_info()["affinity"] = "NPU" + ov_model = core.compile_model(ov_model, OV_DEVICE) + infer_request = ov_model.create_infer_request() + out_name = ov_model.outputs[0] + + OV_COMPILED_CACHE[cache_key] = (infer_request, out_name) + return ov_mm(Tensor_A, Tensor_B, infer_request, out_name) + +@openvino_fp_mm.register_fake +def openvino_fp_mm_fake(A: torch.Tensor, B: torch.Tensor) -> torch.Tensor: + return torch.mm(A.to(dtype=torch.float32), B.to(dtype=torch.float32)) diff --git a/modules/sdnq/layers/linear/linear_fp8_tensorwise.py b/modules/sdnq/layers/linear/linear_fp8_tensorwise.py index 09f349b4d..e730d5668 100644 --- a/modules/sdnq/layers/linear/linear_fp8_tensorwise.py +++ b/modules/sdnq/layers/linear/linear_fp8_tensorwise.py @@ -48,9 +48,9 @@ def fp8_matmul_tensorwise( input, input_scale = quantize_fp_mm_input_tensorwise(input, dtype=scale.dtype) input, weight = check_mats(input, weight, allow_contiguous_mm=False) if bias is not None: - return dequantize_symmetric_with_bias(torch._scaled_mm(input, weight, scale_a=dummy_input_scale, scale_b=dummy_input_scale, bias=None, out_dtype=input_scale.dtype).to(dtype=input_scale.dtype).mul_(input_scale), scale, bias, dtype=return_dtype, result_shape=output_shape) + return dequantize_symmetric_with_bias(torch._scaled_mm(input, weight, scale_a=dummy_input_scale, scale_b=dummy_input_scale, bias=None, out_dtype=input_scale.dtype).mul_(input_scale), scale, bias, dtype=return_dtype, result_shape=output_shape) else: - return dequantize_symmetric(torch._scaled_mm(input, weight, scale_a=dummy_input_scale, scale_b=dummy_input_scale, bias=None, out_dtype=input_scale.dtype).to(dtype=input_scale.dtype).mul_(input_scale), scale, dtype=return_dtype, result_shape=output_shape) + return dequantize_symmetric(torch._scaled_mm(input, weight, scale_a=dummy_input_scale, scale_b=dummy_input_scale, bias=None, out_dtype=input_scale.dtype).mul_(input_scale), scale, dtype=return_dtype, result_shape=output_shape) def quantized_linear_forward_fp8_matmul_tensorwise(self, input: torch.FloatTensor) -> torch.FloatTensor: diff --git a/modules/sdnq/quantizer.py b/modules/sdnq/quantizer.py index aaf4fe911..41f06cda7 100644 --- a/modules/sdnq/quantizer.py +++ b/modules/sdnq/quantizer.py @@ -589,18 +589,6 @@ class SDNQQuantizer(DiffusersQuantizer, HfQuantizer): self.quantization_config.modules_to_not_convert.append(param_name) return False - def check_quantized_param(self, *args, **kwargs) -> bool: - """ - needed for transformers compatibility, returns self.check_if_quantized_param - """ - return self.check_if_quantized_param(*args, **kwargs) - - def param_needs_quantization(self, model, param_name: str, *args, **kwargs) -> bool: - """ - needed for transformers compatibility, returns self.check_if_quantized_param - """ - return self.check_if_quantized_param(model, None, param_name, *args, **kwargs) - @devices.inference_context() def create_quantized_param( # pylint: disable=arguments-differ self, @@ -656,16 +644,6 @@ class SDNQQuantizer(DiffusersQuantizer, HfQuantizer): parent_module, tensor_name = get_module_from_name(model, param_name.removesuffix(tensor_name).removesuffix(".")) setattr(parent_module, tensor_name, layer) - def get_quantize_ops(self): - return SDNQQuantize(self) - - def adjust_max_memory(self, max_memory: dict[str, int | str]) -> dict[str, int | str]: - max_memory = {key: val * 0.80 for key, val in max_memory.items()} - return max_memory - - def adjust_target_dtype(self, target_dtype: torch.dtype) -> torch.dtype: # pylint: disable=unused-argument,arguments-renamed - return dtype_dict[self.quantization_config.weights_dtype]["target_dtype"] - def _process_model_before_weight_loading( # pylint: disable=arguments-differ self, model: torch.nn.Module, @@ -726,6 +704,20 @@ class SDNQQuantizer(DiffusersQuantizer, HfQuantizer): devices.torch_gc(force=True, reason="sdnq") return model + def get_quantize_ops(self): + return SDNQQuantize(self) + + def adjust_max_memory(self, max_memory: dict[str, int | str]) -> dict[str, int | str]: + max_memory = {key: val * 0.80 for key, val in max_memory.items()} + return max_memory + + def adjust_target_dtype(self, target_dtype: torch.dtype) -> torch.dtype: # pylint: disable=unused-argument,arguments-renamed + return dtype_dict[self.quantization_config.weights_dtype]["target_dtype"] + + def update_torch_dtype(self, torch_dtype: torch.dtype) -> torch.dtype: + self.torch_dtype = torch_dtype + return torch_dtype + def get_state_dict_and_metadata(self, state_dict: dict | torch.nn.Module, **kwargs) -> tuple[dict | None, dict]: # pylint: disable=unused-argument, arguments-differ # transformers if isinstance(state_dict, torch.nn.Module): @@ -736,12 +728,6 @@ class SDNQQuantizer(DiffusersQuantizer, HfQuantizer): def get_accelerator_warm_up_factor(self): return 32 // dtype_dict[self.quantization_config.weights_dtype]["num_bits"] - def get_cuda_warm_up_factor(self): - """ - needed for transformers compatibility, returns self.get_accelerator_warm_up_factor - """ - return self.get_accelerator_warm_up_factor() - def _dequantize(self, model): return dequantize_sdnq_model(model) @@ -764,6 +750,30 @@ class SDNQQuantizer(DiffusersQuantizer, HfQuantizer): def is_compileable(self): return True + def check_quantized_param(self, *args, **kwargs) -> bool: + """ + needed for transformers compatibility, returns self.check_if_quantized_param + """ + return self.check_if_quantized_param(*args, **kwargs) + + def param_needs_quantization(self, model, param_name: str, *args, **kwargs) -> bool: + """ + needed for transformers compatibility, returns self.check_if_quantized_param + """ + return self.check_if_quantized_param(model, None, param_name, *args, **kwargs) + + def get_cuda_warm_up_factor(self): + """ + needed for transformers compatibility, returns self.get_accelerator_warm_up_factor + """ + return self.get_accelerator_warm_up_factor() + + def update_dtype(self, dtype: torch.dtype) -> torch.dtype: + """ + needed for transformers compatibility, returns self.update_torch_dtype + """ + return self.update_torch_dtype(dtype) + @dataclass class SDNQConfig(QuantizationConfigMixin):