diff --git a/modules/sdnq/__init__.py b/modules/sdnq/__init__.py index 2654a06e7..62a82e244 100644 --- a/modules/sdnq/__init__.py +++ b/modules/sdnq/__init__.py @@ -159,12 +159,12 @@ def sdnq_quantize_layer(layer, weights_dtype="int8", torch_dtype=None, group_siz zero_point = zero_point.to(torch_dtype) if use_quantized_matmul: - scale = scale.transpose(0,1) - layer.weight.data = layer.weight.transpose(0,1) + scale.transpose_(0,1) + layer.weight.transpose_(0,1) if not dtype_dict[weights_dtype]["is_integer"]: stride = layer.weight.stride() if stride[0] > stride[1] and stride[1] == 1: - layer.weight.data = layer.weight.t().contiguous().t() + layer.weight.data = layer.weight.t_().contiguous().t_() if not use_tensorwise_fp8_matmul: scale = scale.to(torch.float32) diff --git a/modules/sdnq/dequantizer.py b/modules/sdnq/dequantizer.py index 76e830f59..1461b923a 100644 --- a/modules/sdnq/dequantizer.py +++ b/modules/sdnq/dequantizer.py @@ -9,21 +9,21 @@ from .packed_int import pack_int_symetric, unpack_int_symetric, pack_int_asymetr def dequantize_asymmetric(weight: torch.ByteTensor, scale: torch.FloatTensor, zero_point: torch.FloatTensor, dtype: torch.dtype, result_shape: torch.Size) -> torch.FloatTensor: result = torch.addcmul(zero_point, weight.to(dtype=scale.dtype), scale).to(dtype=dtype) if result_shape is not None: - result = result.reshape(result_shape) + result = result.view(result_shape) return result def dequantize_symmetric(weight: torch.CharTensor, scale: torch.FloatTensor, dtype: torch.dtype, result_shape: torch.Size, skip_quantized_matmul: bool = False) -> torch.FloatTensor: result = weight.to(dtype=scale.dtype).mul_(scale).to(dtype=dtype) if skip_quantized_matmul: - result = result.t() + result.t_() if result_shape is not None: - result = result.reshape(result_shape) + result = result.view(result_shape) return result def dequantize_symmetric_with_bias(weight: torch.CharTensor, scale: torch.FloatTensor, bias: torch.FloatTensor, dtype: torch.dtype, result_shape: torch.Size) -> torch.FloatTensor: - return torch.addcmul(bias, weight.to(dtype=scale.dtype), scale).to(dtype=dtype).reshape(result_shape) + return torch.addcmul(bias, weight.to(dtype=scale.dtype), scale).to(dtype=dtype).view(result_shape) def dequantize_packed_int_asymmetric(weight: torch.ByteTensor, scale: torch.FloatTensor, zero_point: torch.FloatTensor, shape: torch.Size, dtype: torch.dtype, result_shape: torch.Size, weights_dtype: str) -> torch.FloatTensor: diff --git a/modules/sdnq/layers/conv/conv_fp8.py b/modules/sdnq/layers/conv/conv_fp8.py index f93d39519..11bf9f424 100644 --- a/modules/sdnq/layers/conv/conv_fp8.py +++ b/modules/sdnq/layers/conv/conv_fp8.py @@ -25,24 +25,24 @@ def conv_fp8_matmul( input, input_scale = quantize_fp8_matmul_input(input) if groups == 1: - result = torch._scaled_mm(input, weight, scale_a=input_scale, scale_b=scale, bias=bias, out_dtype=return_dtype).reshape(mm_output_shape) + result = torch._scaled_mm(input, weight, scale_a=input_scale, scale_b=scale, bias=bias, out_dtype=return_dtype).view(mm_output_shape) else: - scale = scale.reshape(groups, 1, scale.shape[1] // groups) - input_scale = input_scale.reshape(groups, input_scale.shape[0] // groups, 1) - weight = weight.reshape(weight.shape[0], groups, weight.shape[1] // groups).transpose(0,1) - input = input.reshape(input.shape[0], groups, input.shape[1] // groups).transpose(0,1) + scale = scale.view(groups, 1, scale.shape[1] // groups) + input_scale = input_scale.view(groups, input_scale.shape[0] // groups, 1) + weight = weight.view(weight.shape[0], groups, weight.shape[1] // groups) + input = input.view(input.shape[0], groups, input.shape[1] // groups) result = [] if bias is not None: - bias = bias.reshape(groups, bias.shape[0] // groups) + bias = bias.view(groups, bias.shape[0] // groups) for i in range(groups): - result.append(torch._scaled_mm(input[i], weight[i], scale_a=input_scale[i], scale_b=scale[i], bias=bias[i], out_dtype=return_dtype)) + result.append(torch._scaled_mm(input[:, i], weight[:, i], scale_a=input_scale[i], scale_b=scale[i], bias=bias[i], out_dtype=return_dtype)) else: for i in range(groups): - result.append(torch._scaled_mm(input[i], weight[i], scale_a=input_scale[i], scale_b=scale[i], bias=None, out_dtype=return_dtype)) - result = torch.cat(result, dim=-1).reshape(mm_output_shape) + result.append(torch._scaled_mm(input[:, i], weight[:, i], scale_a=input_scale[i], scale_b=scale[i], bias=None, out_dtype=return_dtype)) + result = torch.cat(result, dim=-1).view(mm_output_shape) if conv_type == 1: - result = result.transpose(1,2) + result = result.transpose_(1,2) elif conv_type == 2: result = result.permute(0,3,1,2) elif conv_type == 3: diff --git a/modules/sdnq/layers/conv/conv_fp8_tensorwise.py b/modules/sdnq/layers/conv/conv_fp8_tensorwise.py index 46b53a2d6..a21699855 100644 --- a/modules/sdnq/layers/conv/conv_fp8_tensorwise.py +++ b/modules/sdnq/layers/conv/conv_fp8_tensorwise.py @@ -29,11 +29,11 @@ def conv_fp8_matmul_tensorwise( if groups == 1: result = torch._scaled_mm(input, weight, scale_a=dummy_input_scale, scale_b=dummy_input_scale, bias=None, out_dtype=scale.dtype) else: - weight = weight.reshape(weight.shape[0], groups, weight.shape[1] // groups).transpose(0,1) - input = input.reshape(input.shape[0], groups, input.shape[1] // groups).transpose(0,1) + weight = weight.view(weight.shape[0], groups, weight.shape[1] // groups) + input = input.view(input.shape[0], groups, input.shape[1] // groups) result = [] for i in range(groups): - result.append(torch._scaled_mm(input[i], weight[i], scale_a=dummy_input_scale, scale_b=dummy_input_scale, bias=None, out_dtype=scale.dtype)) + result.append(torch._scaled_mm(input[:, i], weight[:, i], scale_a=dummy_input_scale, scale_b=dummy_input_scale, bias=None, out_dtype=scale.dtype)) result = torch.cat(result, dim=-1) if bias is not None: dequantize_symmetric_with_bias(result, scale, bias, return_dtype, mm_output_shape) @@ -41,7 +41,7 @@ def conv_fp8_matmul_tensorwise( dequantize_symmetric(result, scale, return_dtype, mm_output_shape) if conv_type == 1: - result = result.transpose(1,2) + result = result.transpose_(1,2) elif conv_type == 2: result = result.permute(0,3,1,2) elif conv_type == 3: diff --git a/modules/sdnq/layers/conv/conv_int8.py b/modules/sdnq/layers/conv/conv_int8.py index 02b553d15..625063dfa 100644 --- a/modules/sdnq/layers/conv/conv_int8.py +++ b/modules/sdnq/layers/conv/conv_int8.py @@ -33,11 +33,11 @@ def conv_int8_matmul( if groups == 1: result = torch._int_mm(input, weight) else: - weight = weight.reshape(weight.shape[0], groups, weight.shape[1] // groups).transpose(0,1) - input = input.reshape(input.shape[0], groups, input.shape[1] // groups).transpose(0,1) + weight = weight.view(weight.shape[0], groups, weight.shape[1] // groups) + input = input.view(input.shape[0], groups, input.shape[1] // groups) result = [] for i in range(groups): - result.append(torch._int_mm(input[i], weight[i])) + result.append(torch._int_mm(input[:, i], weight[:, i])) result = torch.cat(result, dim=-1) if bias is not None: result = dequantize_symmetric_with_bias(result, scale, bias, return_dtype, mm_output_shape) @@ -45,7 +45,7 @@ def conv_int8_matmul( result = dequantize_symmetric(result, scale, return_dtype, mm_output_shape) if conv_type == 1: - result = result.transpose(1,2) + result = result.transpose_(1,2) elif conv_type == 2: result = result.permute(0,3,1,2) elif conv_type == 3: diff --git a/modules/sdnq/layers/linear/linear_fp8.py b/modules/sdnq/layers/linear/linear_fp8.py index f6da5c2ef..ba261eb50 100644 --- a/modules/sdnq/layers/linear/linear_fp8.py +++ b/modules/sdnq/layers/linear/linear_fp8.py @@ -21,10 +21,9 @@ def fp8_matmul( scale: torch.FloatTensor, ) -> torch.FloatTensor: return_dtype = input.dtype - output_shape = list(input.shape) - output_shape[-1] = weight.shape[-1] + output_shape = (*input.shape[:-1], weight.shape[-1]) input, input_scale = quantize_fp8_matmul_input(input) - return torch._scaled_mm(input, weight, scale_a=input_scale, scale_b=scale, bias=bias, out_dtype=return_dtype).reshape(output_shape) + return torch._scaled_mm(input, weight, scale_a=input_scale, scale_b=scale, bias=bias, out_dtype=return_dtype).view(output_shape) def quantized_linear_forward_fp8_matmul(self, input: torch.FloatTensor) -> torch.FloatTensor: diff --git a/modules/sdnq/layers/linear/linear_fp8_tensorwise.py b/modules/sdnq/layers/linear/linear_fp8_tensorwise.py index 3fdecaf57..f07154aab 100644 --- a/modules/sdnq/layers/linear/linear_fp8_tensorwise.py +++ b/modules/sdnq/layers/linear/linear_fp8_tensorwise.py @@ -25,8 +25,7 @@ def fp8_matmul_tensorwise( scale: torch.FloatTensor, ) -> torch.FloatTensor: return_dtype = input.dtype - output_shape = list(input.shape) - output_shape[-1] = weight.shape[-1] + output_shape = (*input.shape[:-1], weight.shape[-1]) dummy_input_scale = torch.ones(1, device=input.device, dtype=torch.float32) input, scale = quantize_fp8_matmul_input_tensorwise(input, scale) if bias is not None: diff --git a/modules/sdnq/layers/linear/linear_int8.py b/modules/sdnq/layers/linear/linear_int8.py index 3c0184056..6d94052a2 100644 --- a/modules/sdnq/layers/linear/linear_int8.py +++ b/modules/sdnq/layers/linear/linear_int8.py @@ -30,8 +30,7 @@ def int8_matmul( if quantized_weight_shape is not None: weight = unpack_int_symetric(weight, quantized_weight_shape, weights_dtype, dtype=torch.int8) return_dtype = input.dtype - output_shape = list(input.shape) - output_shape[-1] = weight.shape[-1] + output_shape = (*input.shape[:-1], weight.shape[-1]) input, scale = quantize_int8_matmul_input(input, scale) if bias is not None: return dequantize_symmetric_with_bias(torch._int_mm(input, weight), scale, bias, return_dtype, output_shape) diff --git a/modules/sdnq/packed_int.py b/modules/sdnq/packed_int.py index d946ef4ba..bb0ea1ec8 100644 --- a/modules/sdnq/packed_int.py +++ b/modules/sdnq/packed_int.py @@ -26,7 +26,7 @@ def unpack_int_asymetric(packed_tensor: torch.ByteTensor, shape: torch.Size, wei def pack_uint7(tensor: torch.ByteTensor) -> torch.ByteTensor: - packed_tensor = tensor.contiguous().reshape(-1, 8) + packed_tensor = tensor.contiguous().view(-1, 8) packed_tensor = torch.stack( ( torch.bitwise_or(packed_tensor[:, 0], torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 7], 1), 128)), @@ -43,7 +43,7 @@ def pack_uint7(tensor: torch.ByteTensor) -> torch.ByteTensor: def pack_uint6(tensor: torch.ByteTensor) -> torch.ByteTensor: - packed_tensor = tensor.contiguous().reshape(-1, 4) + packed_tensor = tensor.contiguous().view(-1, 4) packed_tensor = torch.stack( ( torch.bitwise_or(packed_tensor[:, 0], torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 3], 2), 192)), @@ -56,7 +56,7 @@ def pack_uint6(tensor: torch.ByteTensor) -> torch.ByteTensor: def pack_uint5(tensor: torch.ByteTensor) -> torch.ByteTensor: - packed_tensor = tensor.contiguous().reshape(-1, 8) + packed_tensor = tensor.contiguous().view(-1, 8) packed_tensor = torch.stack( ( torch.bitwise_or(packed_tensor[:, 0], torch.bitwise_left_shift(packed_tensor[:, 5], 5)), @@ -83,13 +83,13 @@ def pack_uint5(tensor: torch.ByteTensor) -> torch.ByteTensor: def pack_uint4(tensor: torch.ByteTensor) -> torch.ByteTensor: - packed_tensor = tensor.contiguous().reshape(-1, 2) + packed_tensor = tensor.contiguous().view(-1, 2) packed_tensor = torch.bitwise_or(packed_tensor[:, 0], torch.bitwise_left_shift(packed_tensor[:, 1], 4)) return packed_tensor def pack_uint3(tensor: torch.ByteTensor) -> torch.ByteTensor: - packed_tensor = tensor.contiguous().reshape(-1, 8) + packed_tensor = tensor.contiguous().view(-1, 8) packed_tensor = torch.stack( ( torch.bitwise_or( @@ -114,7 +114,7 @@ def pack_uint3(tensor: torch.ByteTensor) -> torch.ByteTensor: def pack_uint2(tensor: torch.ByteTensor) -> torch.ByteTensor: - packed_tensor = tensor.contiguous().reshape(-1, 4) + packed_tensor = tensor.contiguous().view(-1, 4) packed_tensor = torch.bitwise_or( torch.bitwise_or(packed_tensor[:, 0], torch.bitwise_left_shift(packed_tensor[:, 1], 2)), torch.bitwise_or(torch.bitwise_left_shift(packed_tensor[:, 2], 4), torch.bitwise_left_shift(packed_tensor[:, 3], 6)), @@ -153,7 +153,7 @@ def unpack_uint7(packed_tensor: torch.ByteTensor, shape: torch.Size) -> torch.By ) ), dim=-1 - ).reshape(shape) + ).view(shape) return result @@ -172,7 +172,7 @@ def unpack_uint6(packed_tensor: torch.ByteTensor, shape: torch.Size) -> torch.By ) ), dim=-1 - ).reshape(shape) + ).view(shape) return result @@ -201,12 +201,12 @@ def unpack_uint5(packed_tensor: torch.ByteTensor, shape: torch.Size) -> torch.By ), ), dim=-1 - ).reshape(shape) + ).view(shape) return result def unpack_uint4(packed_tensor: torch.ByteTensor, shape: torch.Size) -> torch.ByteTensor: - result = torch.stack((torch.bitwise_and(packed_tensor, 15), torch.bitwise_right_shift(packed_tensor, 4)), dim=-1).reshape(shape) + result = torch.stack((torch.bitwise_and(packed_tensor, 15), torch.bitwise_right_shift(packed_tensor, 4)), dim=-1).view(shape) return result @@ -229,7 +229,7 @@ def unpack_uint3(packed_tensor: torch.ByteTensor, shape: torch.Size) -> torch.By ), ), dim=-1 - ).reshape(shape) + ).view(shape) return result @@ -242,7 +242,7 @@ def unpack_uint2(packed_tensor: torch.ByteTensor, shape: torch.Size) -> torch.By torch.bitwise_right_shift(packed_tensor, 6), ), dim=-1 - ).reshape(shape) + ).view(shape) return result