From 569e9099d78c9e288fe6e536258ecd31f215b1d4 Mon Sep 17 00:00:00 2001 From: Disty0 Date: Wed, 28 May 2025 12:44:07 +0300 Subject: [PATCH] Use torch.amax instead of torch.max --- modules/model_quant_sdnq.py | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/modules/model_quant_sdnq.py b/modules/model_quant_sdnq.py index 26abcd9eb..ac7ce71f6 100644 --- a/modules/model_quant_sdnq.py +++ b/modules/model_quant_sdnq.py @@ -93,13 +93,13 @@ def sdnq_quantize_layer(layer, weights_dtype="int8", torch_dtype=None, group_siz layer.weight.data = layer.weight.reshape(new_shape) layer.weight.requires_grad = False - if pre_mode: + if shared.opts.diffusers_offload_mode in {"none", "model"}: + return_device = devices.device + elif pre_mode: if shared.opts.device_map != "gpu": return_device = devices.cpu else: return_device = devices.device - elif shared.opts.diffusers_offload_mode in {"none", "model"}: - return_device = devices.device else: return_device = layer.weight.device if not pre_mode: @@ -264,8 +264,14 @@ def unpack_int4(packed_tensor: torch.Tensor, shape: torch.Size, dtype: Optional[ return result +def quantize_fp8_matmul_input(input: torch.FloatTensor) -> Tuple[torch.FloatTensor, torch.FloatTensor]: + input_scale = torch.div(input.abs().amax(dim=-1), 448).unsqueeze(-1) + input = torch.div(input, input_scale).clamp_(-448, 448).to(torch.float8_e4m3fn) + return input, input_scale + + def quantize_int8_matmul_input(input: torch.FloatTensor, scale: torch.FloatTensor) -> Tuple[torch.ByteTensor, torch.FloatTensor]: - input_scale = torch.div(input.abs().max(dim=-1).values, 127).unsqueeze(-1) + input_scale = torch.div(input.abs().amax(dim=-1), 127).unsqueeze(-1) input = torch.div(input, input_scale).round_().clamp_(-128, 127).to(torch.int8).flatten(0,-2).contiguous() scale = torch.mul(input_scale, scale).flatten(0,-2).contiguous() if scale.dtype == torch.float16: # fp16 will overflow