mirror of https://github.com/vladmandic/automatic
410 lines
19 KiB
Python
410 lines
19 KiB
Python
# pylint: disable=redefined-builtin,no-member,protected-access
|
|
|
|
from typing import Callable, List, Tuple, Optional
|
|
import torch
|
|
from modules import shared
|
|
|
|
from .common import conv_types, conv_transpose_types
|
|
from .dequantizer import dequantize_symmetric, dequantize_symmetric_with_bias
|
|
from .packed_int import unpack_int_symetric
|
|
|
|
|
|
def get_forward_func(layer_class_name: str, use_quantized_matmul: bool, is_integer: bool, use_tensorwise_fp8_matmul: bool) -> Callable: # pylint: disable=inconsistent-return-statements
|
|
if layer_class_name in conv_types:
|
|
if use_quantized_matmul:
|
|
if is_integer:
|
|
return quantized_conv_forward_int8_matmul
|
|
else:
|
|
if use_tensorwise_fp8_matmul:
|
|
return quantized_conv_forward_fp8_matmul_tensorwise
|
|
else:
|
|
return quantized_conv_forward_fp8_matmul
|
|
else:
|
|
return quantized_conv_forward
|
|
elif layer_class_name in conv_transpose_types:
|
|
if layer_class_name.endswith("1d"):
|
|
return quantized_conv_transpose_1d_forward
|
|
elif layer_class_name.endswith("2d"):
|
|
return quantized_conv_transpose_2d_forward
|
|
elif layer_class_name.endswith("3d"):
|
|
return quantized_conv_transpose_3d_forward
|
|
else:
|
|
if use_quantized_matmul:
|
|
if is_integer:
|
|
return quantized_linear_forward_int8_matmul
|
|
else:
|
|
if use_tensorwise_fp8_matmul:
|
|
return quantized_linear_forward_fp8_matmul_tensorwise
|
|
else:
|
|
return quantized_linear_forward_fp8_matmul
|
|
else:
|
|
return quantized_linear_forward
|
|
|
|
|
|
def quantize_fp8_matmul_input(input: torch.FloatTensor) -> Tuple[torch.Tensor, torch.FloatTensor]:
|
|
input = input.flatten(0,-2).contiguous()
|
|
input_scale = torch.amax(input.abs(), dim=-1, keepdims=True).div_(448)
|
|
input = torch.div(input, input_scale).clamp_(-448, 448).to(dtype=torch.float8_e4m3fn)
|
|
input_scale = input_scale.to(dtype=torch.float32)
|
|
return input, input_scale
|
|
|
|
|
|
def quantize_fp8_matmul_input_tensorwise(input: torch.FloatTensor, scale: torch.FloatTensor) -> Tuple[torch.Tensor, torch.FloatTensor]:
|
|
input = input.flatten(0,-2).contiguous()
|
|
input_scale = torch.amax(input.abs(), dim=-1, keepdims=True).div_(448)
|
|
input = torch.div(input, input_scale).clamp_(-448, 448).to(dtype=torch.float8_e4m3fn)
|
|
scale = torch.mul(input_scale, scale)
|
|
if scale.dtype == torch.float16: # fp16 will overflow
|
|
scale = scale.to(dtype=torch.float32)
|
|
return input, scale
|
|
|
|
|
|
def quantize_int8_matmul_input(input: torch.FloatTensor, scale: torch.FloatTensor) -> Tuple[torch.CharTensor, torch.FloatTensor]:
|
|
input = input.flatten(0,-2).contiguous()
|
|
input_scale = torch.amax(input.abs(), dim=-1, keepdims=True).div_(127)
|
|
input = torch.div(input, input_scale).round_().clamp_(-128, 127).to(dtype=torch.int8)
|
|
scale = torch.mul(input_scale, scale)
|
|
if scale.dtype == torch.float16: # fp16 will overflow
|
|
scale = scale.to(dtype=torch.float32)
|
|
return input, scale
|
|
|
|
|
|
def fp8_matmul(
|
|
input: torch.FloatTensor,
|
|
weight: torch.Tensor,
|
|
bias: torch.FloatTensor,
|
|
scale: torch.FloatTensor,
|
|
) -> torch.FloatTensor:
|
|
return_dtype = input.dtype
|
|
output_shape = list(input.shape)
|
|
output_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)
|
|
|
|
|
|
# sm89 doesn't support row wise scale in Windows
|
|
def fp8_matmul_tensorwise(
|
|
input: torch.FloatTensor,
|
|
weight: torch.Tensor,
|
|
bias: torch.FloatTensor,
|
|
scale: torch.FloatTensor,
|
|
) -> torch.FloatTensor:
|
|
return_dtype = input.dtype
|
|
output_shape = list(input.shape)
|
|
output_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:
|
|
return dequantize_symmetric_with_bias(torch._scaled_mm(input, weight, scale_a=dummy_input_scale, scale_b=dummy_input_scale, bias=None, out_dtype=scale.dtype), scale, bias, return_dtype, 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=scale.dtype), scale, return_dtype, output_shape)
|
|
|
|
|
|
def int8_matmul(
|
|
input: torch.FloatTensor,
|
|
weight: torch.Tensor,
|
|
bias: torch.FloatTensor,
|
|
scale: torch.FloatTensor,
|
|
quantized_weight_shape: torch.Size,
|
|
weights_dtype: str,
|
|
) -> torch.FloatTensor:
|
|
if quantized_weight_shape is not None:
|
|
weight = unpack_int_symetric(weight, quantized_weight_shape, weights_dtype, dtype=torch.int8, transpose=True)
|
|
return_dtype = input.dtype
|
|
output_shape = list(input.shape)
|
|
output_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)
|
|
else:
|
|
return dequantize_symmetric(torch._int_mm(input, weight), scale, return_dtype, output_shape)
|
|
|
|
|
|
def process_conv_input(conv_type, input, reversed_padding_repeated_twice, padding_mode, result_shape, stride, padding, dilation):
|
|
if conv_type == 1:
|
|
batch_size, _, L_in = input.shape
|
|
C_out, _, K_l = result_shape
|
|
L_out = (L_in + 2 * padding[1] - dilation[1] * (K_l - 1) - 1) // stride[1] + 1
|
|
mm_output_shape = (batch_size, L_out, C_out)
|
|
kernel_size = (1, K_l)
|
|
if conv_type == 2:
|
|
batch_size, _, H_in, W_in = input.shape
|
|
C_out, _, K_h, K_w = result_shape
|
|
H_out = (H_in + 2 * padding[0] - dilation[0] * (K_h - 1) - 1) // stride[0] + 1
|
|
W_out = (W_in + 2 * padding[1] - dilation[1] * (K_w - 1) - 1) // stride[1] + 1
|
|
mm_output_shape = (batch_size, H_out, W_out, C_out)
|
|
kernel_size = (K_h, K_w)
|
|
else:
|
|
batch_size, _, D_in, H_in, W_in = input.shape
|
|
C_out, _, K_d, K_h, K_w = result_shape
|
|
D_out = (D_in + 2 * padding[0] - dilation[0] * (K_d - 1) - 1) // stride[0] + 1
|
|
H_out = (H_in + 2 * padding[1] - dilation[1] * (K_h - 1) - 1) // stride[1] + 1
|
|
W_out = (W_in + 2 * padding[2] - dilation[2] * (K_w - 1) - 1) // stride[2] + 1
|
|
mm_output_shape = (batch_size, D_out, H_out, W_out, C_out)
|
|
kernel_size = (K_d, K_h, K_w)
|
|
|
|
if padding_mode != "zeros":
|
|
input = torch.nn.functional.pad(input, reversed_padding_repeated_twice, mode=padding_mode)
|
|
padding = (0,) * (conv_type if conv_type != 1 else 2)
|
|
elif conv_type == 3:
|
|
input = torch.nn.functional.pad(input, reversed_padding_repeated_twice)
|
|
|
|
if conv_type == 1:
|
|
input = input.unsqueeze(2)
|
|
|
|
if conv_type == 3:
|
|
K_D_eff = K_d + (K_d - 1) * (dilation[0] - 1)
|
|
K_H_eff = K_h + (K_h - 1) * (dilation[0] - 1)
|
|
K_W_eff = K_w + (K_w - 1) * (dilation[0] - 1)
|
|
input = input.unfold(2, K_D_eff, stride[0]).unfold(3, K_H_eff, stride[1]).unfold(4, K_W_eff, stride[2])
|
|
if dilation[0] > 1:
|
|
input = input[..., ::dilation[0], :, :]
|
|
if dilation[1] > 1:
|
|
input = input[..., ::dilation[1], :]
|
|
if dilation[2] > 1:
|
|
input = input[..., ::dilation[2]]
|
|
input = input.permute(0, 2, 3, 4, 1, 5, 6, 7).reshape(batch_size, D_out * H_out * W_out, -1)
|
|
else:
|
|
input = torch.nn.functional.unfold(input, kernel_size=kernel_size, padding=padding, stride=stride, dilation=dilation).transpose(1,2)
|
|
return input, mm_output_shape
|
|
|
|
|
|
def conv_fp8_matmul(
|
|
input: torch.FloatTensor,
|
|
weight: torch.Tensor,
|
|
bias: torch.FloatTensor,
|
|
scale: torch.FloatTensor,
|
|
result_shape: torch.Size,
|
|
reversed_padding_repeated_twice: List[int],
|
|
padding_mode: str, conv_type: int,
|
|
groups: int, stride: List[int],
|
|
padding: List[int], dilation: List[int],
|
|
) -> torch.FloatTensor:
|
|
return_dtype = input.dtype
|
|
input, mm_output_shape = process_conv_input(conv_type, input, reversed_padding_repeated_twice, padding_mode, result_shape, stride, padding, dilation)
|
|
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)
|
|
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)
|
|
result = []
|
|
if bias is not None:
|
|
bias = bias.reshape(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))
|
|
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)
|
|
|
|
if conv_type == 1:
|
|
result = result.transpose(1,2)
|
|
elif conv_type == 2:
|
|
result = result.permute(0,3,1,2)
|
|
elif conv_type == 3:
|
|
result = result.permute(0,4,1,2,3)
|
|
return result
|
|
|
|
|
|
def conv_fp8_matmul_tensorwise(
|
|
input: torch.FloatTensor,
|
|
weight: torch.Tensor,
|
|
bias: torch.FloatTensor,
|
|
scale: torch.FloatTensor,
|
|
result_shape: torch.Size,
|
|
reversed_padding_repeated_twice: List[int],
|
|
padding_mode: str, conv_type: int,
|
|
groups: int, stride: List[int],
|
|
padding: List[int], dilation: List[int],
|
|
) -> torch.FloatTensor:
|
|
return_dtype = input.dtype
|
|
input, mm_output_shape = process_conv_input(conv_type, input, reversed_padding_repeated_twice, padding_mode, result_shape, stride, padding, dilation)
|
|
input, scale = quantize_fp8_matmul_input_tensorwise(input, scale)
|
|
dummy_input_scale = torch.ones(1, device=input.device, dtype=torch.float32)
|
|
|
|
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)
|
|
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 = torch.cat(result, dim=-1)
|
|
if bias is not None:
|
|
dequantize_symmetric_with_bias(result, scale, bias, return_dtype, mm_output_shape)
|
|
else:
|
|
dequantize_symmetric(result, scale, return_dtype, mm_output_shape)
|
|
|
|
if conv_type == 1:
|
|
result = result.transpose(1,2)
|
|
elif conv_type == 2:
|
|
result = result.permute(0,3,1,2)
|
|
elif conv_type == 3:
|
|
result = result.permute(0,4,1,2,3)
|
|
return result
|
|
|
|
|
|
def conv_int8_matmul(
|
|
input: torch.FloatTensor,
|
|
weight: torch.CharTensor,
|
|
bias: torch.FloatTensor,
|
|
scale: torch.FloatTensor,
|
|
result_shape: torch.Size,
|
|
quantized_weight_shape: torch.Size,
|
|
weights_dtype: str,
|
|
reversed_padding_repeated_twice: List[int],
|
|
padding_mode: str, conv_type: int,
|
|
groups: int, stride: List[int],
|
|
padding: List[int], dilation: List[int],
|
|
) -> torch.FloatTensor:
|
|
return_dtype = input.dtype
|
|
input, mm_output_shape = process_conv_input(conv_type, input, reversed_padding_repeated_twice, padding_mode, result_shape, stride, padding, dilation)
|
|
input, scale = quantize_int8_matmul_input(input, scale)
|
|
if quantized_weight_shape is not None:
|
|
weight = unpack_int_symetric(weight, quantized_weight_shape, weights_dtype, dtype=torch.int8, transpose=True)
|
|
|
|
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)
|
|
result = []
|
|
for i in range(groups):
|
|
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)
|
|
else:
|
|
result = dequantize_symmetric(result, scale, return_dtype, mm_output_shape)
|
|
|
|
if conv_type == 1:
|
|
result = result.transpose(1,2)
|
|
elif conv_type == 2:
|
|
result = result.permute(0,3,1,2)
|
|
elif conv_type == 3:
|
|
result = result.permute(0,4,1,2,3)
|
|
return result
|
|
|
|
|
|
def quantized_linear_forward_fp8_matmul(self, input: torch.FloatTensor) -> torch.FloatTensor:
|
|
if torch.numel(input) / input.shape[-1] < 32:
|
|
return torch.nn.functional.linear(input, self.sdnq_dequantizer(self.weight, skip_quantized_matmul=True), self.bias)
|
|
return fp8_matmul(input, self.weight, self.bias, self.sdnq_dequantizer.scale)
|
|
|
|
|
|
def quantized_linear_forward_fp8_matmul_tensorwise(self, input: torch.FloatTensor) -> torch.FloatTensor:
|
|
if torch.numel(input) / input.shape[-1] < 32:
|
|
return torch.nn.functional.linear(input, self.sdnq_dequantizer(self.weight, skip_quantized_matmul=True), self.bias)
|
|
return fp8_matmul_tensorwise(input, self.weight, self.bias, self.sdnq_dequantizer.scale)
|
|
|
|
|
|
def quantized_linear_forward_int8_matmul(self, input: torch.FloatTensor) -> torch.FloatTensor:
|
|
if torch.numel(input) / input.shape[-1] < 32:
|
|
return torch.nn.functional.linear(input, self.sdnq_dequantizer(self.weight, skip_quantized_matmul=True), self.bias)
|
|
return int8_matmul(input, self.weight, self.bias, self.sdnq_dequantizer.scale, getattr(self.sdnq_dequantizer, "quantized_weight_shape", None), self.sdnq_dequantizer.weights_dtype)
|
|
|
|
|
|
def quantized_linear_forward(self, input: torch.FloatTensor) -> torch.FloatTensor:
|
|
return torch.nn.functional.linear(input, self.sdnq_dequantizer(self.weight), self.bias)
|
|
|
|
|
|
def get_conv_args(input_ndim: int, stride, padding, dilation):
|
|
if input_ndim == 3:
|
|
conv_type = 1
|
|
elif input_ndim == 4:
|
|
conv_type = 2
|
|
else:
|
|
conv_type = 3
|
|
if isinstance(stride, int):
|
|
stride = (stride,) * conv_type
|
|
if isinstance(padding, int):
|
|
padding = (padding,) * conv_type
|
|
if isinstance(dilation, int):
|
|
dilation = (dilation,) * conv_type
|
|
if conv_type == 1:
|
|
stride = (1, stride[0])
|
|
padding = (0, padding[0])
|
|
dilation = (1, dilation[0])
|
|
return conv_type, stride, padding, dilation
|
|
|
|
|
|
def quantized_conv_forward_fp8_matmul(self, input) -> torch.FloatTensor:
|
|
if torch.numel(input) / input.shape[2] < 32:
|
|
return self._conv_forward(input, self.sdnq_dequantizer(self.weight, skip_quantized_matmul=True), self.bias)
|
|
conv_type, stride, padding, dilation = get_conv_args(input.ndim, self.stride, self.padding, self.dilation)
|
|
return conv_fp8_matmul(
|
|
input, self.weight, self.bias,
|
|
self.sdnq_dequantizer.scale,
|
|
self.sdnq_dequantizer.result_shape,
|
|
self._reversed_padding_repeated_twice,
|
|
self.padding_mode, conv_type,
|
|
self.groups, stride, padding, dilation,
|
|
)
|
|
|
|
|
|
def quantized_conv_forward_fp8_matmul_tensorwise(self, input) -> torch.FloatTensor:
|
|
if torch.numel(input) / input.shape[2] < 32:
|
|
return self._conv_forward(input, self.sdnq_dequantizer(self.weight, skip_quantized_matmul=True), self.bias)
|
|
conv_type, stride, padding, dilation = get_conv_args(input.ndim, self.stride, self.padding, self.dilation)
|
|
return conv_fp8_matmul_tensorwise(
|
|
input, self.weight, self.bias,
|
|
self.sdnq_dequantizer.scale,
|
|
self.sdnq_dequantizer.result_shape,
|
|
self._reversed_padding_repeated_twice,
|
|
self.padding_mode, conv_type,
|
|
self.groups, stride, padding, dilation,
|
|
)
|
|
|
|
|
|
def quantized_conv_forward_int8_matmul(self, input) -> torch.FloatTensor:
|
|
if torch.numel(input) / input.shape[2] < 32:
|
|
return self._conv_forward(input, self.sdnq_dequantizer(self.weight, skip_quantized_matmul=True), self.bias)
|
|
conv_type, stride, padding, dilation = get_conv_args(input.ndim, self.stride, self.padding, self.dilation)
|
|
return conv_int8_matmul(
|
|
input, self.weight, self.bias,
|
|
self.sdnq_dequantizer.scale,
|
|
self.sdnq_dequantizer.result_shape,
|
|
getattr(self.sdnq_dequantizer, "quantized_weight_shape", None),
|
|
self.sdnq_dequantizer.weights_dtype,
|
|
self._reversed_padding_repeated_twice,
|
|
self.padding_mode, conv_type,
|
|
self.groups, stride, padding, dilation,
|
|
)
|
|
|
|
|
|
def quantized_conv_forward(self, input) -> torch.FloatTensor:
|
|
return self._conv_forward(input, self.sdnq_dequantizer(self.weight), self.bias)
|
|
|
|
|
|
def quantized_conv_transpose_1d_forward(self, input: torch.FloatTensor, output_size: Optional[list[int]] = None) -> torch.FloatTensor:
|
|
output_padding = self._output_padding(input, output_size, self.stride, self.padding, self.kernel_size, 1, self.dilation)
|
|
return torch.nn.functional.conv_transpose1d(input, self.sdnq_dequantizer(self.weight), self.bias, self.stride, self.padding, output_padding, self.groups, self.dilation)
|
|
|
|
|
|
def quantized_conv_transpose_2d_forward(self, input: torch.FloatTensor, output_size: Optional[list[int]] = None) -> torch.FloatTensor:
|
|
output_padding = self._output_padding(input, output_size, self.stride, self.padding, self.kernel_size, 2, self.dilation)
|
|
return torch.nn.functional.conv_transpose2d(input, self.sdnq_dequantizer(self.weight), self.bias, self.stride, self.padding, output_padding, self.groups, self.dilation)
|
|
|
|
|
|
def quantized_conv_transpose_3d_forward(self, input: torch.FloatTensor, output_size: Optional[list[int]] = None) -> torch.FloatTensor:
|
|
output_padding = self._output_padding(input, output_size, self.stride, self.padding, self.kernel_size, 3, self.dilation)
|
|
return torch.nn.functional.conv_transpose3d(input, self.sdnq_dequantizer(self.weight), self.bias, self.stride, self.padding, output_padding, self.groups, self.dilation)
|
|
|
|
|
|
if shared.opts.sdnq_dequantize_compile:
|
|
try:
|
|
torch._dynamo.config.cache_size_limit = max(8192, torch._dynamo.config.cache_size_limit)
|
|
int8_matmul = torch.compile(int8_matmul, fullgraph=True)
|
|
fp8_matmul = torch.compile(fp8_matmul, fullgraph=True)
|
|
fp8_matmul_tensorwise = torch.compile(fp8_matmul_tensorwise, fullgraph=True)
|
|
conv_int8_matmul = torch.compile(conv_int8_matmul, fullgraph=True)
|
|
conv_fp8_matmul = torch.compile(conv_fp8_matmul, fullgraph=True)
|
|
conv_fp8_matmul_tensorwise = torch.compile(conv_fp8_matmul_tensorwise, fullgraph=True)
|
|
except Exception as e:
|
|
shared.log.warning(f"Quantization: type=sdnq MatMul using torch.compile is not available: {e}")
|