mirror of
https://github.com/vladmandic/automatic
synced 2026-09-19 17:24:32 +02:00
SDNQ split forward.py into layers and cleanup
This commit is contained in:
Disty0
+35
-36
@@ -14,6 +14,41 @@ from .dequantizer import dequantizer_dict
|
||||
from .forward import get_forward_func
|
||||
|
||||
|
||||
class QuantizationMethod(str, Enum):
|
||||
SDNQ = "sdnq"
|
||||
|
||||
|
||||
def get_scale_asymmetric(weight: torch.FloatTensor, reduction_axes: Union[int, List[int]], weights_dtype: str) -> Tuple[torch.FloatTensor, torch.FloatTensor]:
|
||||
zero_point = torch.amin(weight, dim=reduction_axes, keepdims=True)
|
||||
scale = torch.amax(weight, dim=reduction_axes, keepdims=True).sub_(zero_point).div_(dtype_dict[weights_dtype]["max"] - dtype_dict[weights_dtype]["min"])
|
||||
eps = torch.finfo(scale.dtype).eps # prevent divison by 0
|
||||
scale = torch.where(torch.abs(scale) < eps, eps, scale)
|
||||
if dtype_dict[weights_dtype]["min"] != 0:
|
||||
zero_point.sub_(torch.mul(scale, dtype_dict[weights_dtype]["min"]))
|
||||
return scale, zero_point
|
||||
|
||||
|
||||
def get_scale_symmetric(weight: torch.FloatTensor, reduction_axes: Union[int, List[int]], weights_dtype: str) -> torch.FloatTensor:
|
||||
scale = torch.amax(weight.abs(), dim=reduction_axes, keepdims=True).div_(dtype_dict[weights_dtype]["max"])
|
||||
eps = torch.finfo(scale.dtype).eps # prevent divison by 0
|
||||
scale = torch.where(torch.abs(scale) < eps, eps, scale)
|
||||
return scale
|
||||
|
||||
|
||||
def quantize_weight(weight: torch.FloatTensor, reduction_axes: Union[int, List[int]], weights_dtype: str) -> Tuple[torch.Tensor, torch.FloatTensor, torch.FloatTensor]:
|
||||
if dtype_dict[weights_dtype]["is_unsigned"]:
|
||||
scale, zero_point = get_scale_asymmetric(weight, reduction_axes, weights_dtype)
|
||||
quantized_weight = torch.sub(weight, zero_point).div_(scale)
|
||||
else:
|
||||
scale = get_scale_symmetric(weight, reduction_axes, weights_dtype)
|
||||
quantized_weight = torch.div(weight, scale)
|
||||
zero_point = None
|
||||
if dtype_dict[weights_dtype]["is_integer"]:
|
||||
quantized_weight.round_()
|
||||
quantized_weight = quantized_weight.clamp_(dtype_dict[weights_dtype]["min"], dtype_dict[weights_dtype]["max"]).to(dtype_dict[weights_dtype]["torch_dtype"])
|
||||
return quantized_weight, scale, zero_point
|
||||
|
||||
|
||||
@devices.inference_context()
|
||||
def sdnq_quantize_layer(layer, weights_dtype="int8", torch_dtype=None, group_size=0, quant_conv=False, use_quantized_matmul=False, use_quantized_matmul_conv=False, dequantize_fp32=False, quantization_device=None, return_device=None, param_name=None): # pylint: disable=unused-argument
|
||||
layer_class_name = layer.__class__.__name__
|
||||
@@ -153,7 +188,6 @@ def sdnq_quantize_layer(layer, weights_dtype="int8", torch_dtype=None, group_siz
|
||||
|
||||
layer.forward = get_forward_func(layer_class_name, use_quantized_matmul, dtype_dict[weights_dtype]["is_integer"], use_tensorwise_fp8_matmul)
|
||||
layer.forward = layer.forward.__get__(layer, layer.__class__)
|
||||
#devices.torch_gc(force=False, reason=f"SDNQ param_name: {param_name}")
|
||||
return layer
|
||||
|
||||
|
||||
@@ -195,41 +229,6 @@ def apply_sdnq_to_module(model, weights_dtype="int8", torch_dtype=None, group_si
|
||||
return model
|
||||
|
||||
|
||||
def get_scale_asymmetric(weight: torch.FloatTensor, reduction_axes: Union[int, List[int]], weights_dtype: str) -> Tuple[torch.FloatTensor, torch.FloatTensor]:
|
||||
zero_point = torch.amin(weight, dim=reduction_axes, keepdims=True)
|
||||
scale = torch.amax(weight, dim=reduction_axes, keepdims=True).sub_(zero_point).div_(dtype_dict[weights_dtype]["max"] - dtype_dict[weights_dtype]["min"])
|
||||
eps = torch.finfo(scale.dtype).eps # prevent divison by 0
|
||||
scale = torch.where(torch.abs(scale) < eps, eps, scale)
|
||||
if dtype_dict[weights_dtype]["min"] != 0:
|
||||
zero_point.sub_(torch.mul(scale, dtype_dict[weights_dtype]["min"]))
|
||||
return scale, zero_point
|
||||
|
||||
|
||||
def get_scale_symmetric(weight: torch.FloatTensor, reduction_axes: Union[int, List[int]], weights_dtype: str) -> torch.FloatTensor:
|
||||
scale = torch.amax(weight.abs(), dim=reduction_axes, keepdims=True).div_(dtype_dict[weights_dtype]["max"])
|
||||
eps = torch.finfo(scale.dtype).eps # prevent divison by 0
|
||||
scale = torch.where(torch.abs(scale) < eps, eps, scale)
|
||||
return scale
|
||||
|
||||
|
||||
def quantize_weight(weight: torch.FloatTensor, reduction_axes: Union[int, List[int]], weights_dtype: str) -> Tuple[torch.Tensor, torch.FloatTensor, torch.FloatTensor]:
|
||||
if dtype_dict[weights_dtype]["is_unsigned"]:
|
||||
scale, zero_point = get_scale_asymmetric(weight, reduction_axes, weights_dtype)
|
||||
quantized_weight = torch.sub(weight, zero_point).div_(scale)
|
||||
else:
|
||||
scale = get_scale_symmetric(weight, reduction_axes, weights_dtype)
|
||||
quantized_weight = torch.div(weight, scale)
|
||||
zero_point = None
|
||||
if dtype_dict[weights_dtype]["is_integer"]:
|
||||
quantized_weight.round_()
|
||||
quantized_weight = quantized_weight.clamp_(dtype_dict[weights_dtype]["min"], dtype_dict[weights_dtype]["max"]).to(dtype_dict[weights_dtype]["torch_dtype"])
|
||||
return quantized_weight, scale, zero_point
|
||||
|
||||
|
||||
class QuantizationMethod(str, Enum):
|
||||
SDNQ = "sdnq"
|
||||
|
||||
|
||||
class SDNQQuantizer(DiffusersQuantizer):
|
||||
r"""
|
||||
Diffusers Quantizer for SDNQ
|
||||
|
||||
+11
-2
@@ -1,7 +1,7 @@
|
||||
# pylint: disable=redefined-builtin,no-member,protected-access
|
||||
|
||||
import torch
|
||||
from modules import devices
|
||||
from modules import devices, shared
|
||||
|
||||
torch_version = float(torch.__version__[:3])
|
||||
|
||||
@@ -30,7 +30,9 @@ if hasattr(torch, "float8_e4m3fnuz"):
|
||||
if hasattr(torch, "float8_e5m2fnuz"):
|
||||
dtype_dict["float8_e5m2fnuz"] = {"min": -57344, "max": 57344, "num_bits": 8, "target_dtype": "fp8", "torch_dtype": torch.float8_e5m2fnuz, "storage_dtype": torch.float8_e5m2fnuz, "is_unsigned": False, "is_integer": False}
|
||||
|
||||
use_tensorwise_fp8_matmul = True # Direct tensorwise only exist on H100 hardware, sdnq will use software tensorwise with this setting
|
||||
use_torch_compile = shared.opts.sdnq_dequantize_compile # this setting requires a full restart of the webui to apply
|
||||
use_tensorwise_fp8_matmul = True # row-wise FP8 only exist on H100 hardware, sdnq will use software row-wise with tensorwise hardware with this setting
|
||||
|
||||
quantized_matmul_dtypes = ("int8", "int7", "int6", "int5", "int4", "int3", "int2", "float8_e4m3fn", "float8_e5m2")
|
||||
if devices.backend in {"cpu", "openvino"}:
|
||||
quantized_matmul_dtypes += ("float8_e4m3fnuz", "float8_e5m2fnuz")
|
||||
@@ -39,3 +41,10 @@ linear_types = ("Linear",)
|
||||
conv_types = ("Conv1d", "Conv2d", "Conv3d")
|
||||
conv_transpose_types = ("ConvTranspose1d", "ConvTranspose2d", "ConvTranspose3d")
|
||||
allowed_types = linear_types + conv_types + conv_transpose_types
|
||||
|
||||
if use_torch_compile:
|
||||
try:
|
||||
torch._dynamo.config.cache_size_limit = max(8192, torch._dynamo.config.cache_size_limit)
|
||||
torch._dynamo.config.accumulated_recompile_limit = max(8192, torch._dynamo.config.accumulated_recompile_limit)
|
||||
except Exception as e:
|
||||
shared.log.warning(f"Quantization: type=sdnq Failed to increase the cache size for torch.compile: {e}")
|
||||
|
||||
@@ -3,7 +3,7 @@
|
||||
import torch
|
||||
from modules import shared
|
||||
|
||||
from .common import dtype_dict
|
||||
from .common import dtype_dict, use_torch_compile
|
||||
from .packed_int import pack_int_symetric, unpack_int_symetric, packed_int_function_dict
|
||||
|
||||
|
||||
@@ -173,10 +173,8 @@ dequantizer_dict = {
|
||||
}
|
||||
|
||||
|
||||
if shared.opts.sdnq_dequantize_compile:
|
||||
if use_torch_compile:
|
||||
try:
|
||||
torch._dynamo.config.cache_size_limit = max(8192, torch._dynamo.config.cache_size_limit)
|
||||
torch._dynamo.config.accumulated_recompile_limit = max(8192, torch._dynamo.config.accumulated_recompile_limit)
|
||||
dequantize_asymmetric_compiled = torch.compile(dequantize_asymmetric, fullgraph=True, dynamic=False)
|
||||
dequantize_symmetric_compiled = torch.compile(dequantize_symmetric, fullgraph=True, dynamic=False)
|
||||
dequantize_packed_int_asymmetric_compiled = torch.compile(dequantize_packed_int_asymmetric, fullgraph=True, dynamic=False)
|
||||
|
||||
+14
-374
@@ -1,410 +1,50 @@
|
||||
# pylint: disable=redefined-builtin,no-member,protected-access
|
||||
# pylint: disable=protected-access
|
||||
|
||||
from typing import Callable
|
||||
|
||||
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:
|
||||
from .layers.conv.conv_int8 import quantized_conv_forward_int8_matmul
|
||||
return quantized_conv_forward_int8_matmul
|
||||
else:
|
||||
if use_tensorwise_fp8_matmul:
|
||||
from .layers.conv.conv_fp8_tensorwise import quantized_conv_forward_fp8_matmul_tensorwise
|
||||
return quantized_conv_forward_fp8_matmul_tensorwise
|
||||
else:
|
||||
from .layers.conv.conv_fp8 import quantized_conv_forward_fp8_matmul
|
||||
return quantized_conv_forward_fp8_matmul
|
||||
else:
|
||||
from .layers.conv.forward import quantized_conv_forward
|
||||
return quantized_conv_forward
|
||||
elif layer_class_name in conv_transpose_types:
|
||||
if layer_class_name.endswith("1d"):
|
||||
from .layers.conv.forward import quantized_conv_transpose_1d_forward
|
||||
return quantized_conv_transpose_1d_forward
|
||||
elif layer_class_name.endswith("2d"):
|
||||
from .layers.conv.forward import quantized_conv_transpose_2d_forward
|
||||
return quantized_conv_transpose_2d_forward
|
||||
elif layer_class_name.endswith("3d"):
|
||||
from .layers.conv.forward import quantized_conv_transpose_3d_forward
|
||||
return quantized_conv_transpose_3d_forward
|
||||
else:
|
||||
if use_quantized_matmul:
|
||||
if is_integer:
|
||||
from .layers.linear.linear_int8 import quantized_linear_forward_int8_matmul
|
||||
return quantized_linear_forward_int8_matmul
|
||||
else:
|
||||
if use_tensorwise_fp8_matmul:
|
||||
from .layers.linear.linear_fp8_tensorwise import quantized_linear_forward_fp8_matmul_tensorwise
|
||||
return quantized_linear_forward_fp8_matmul_tensorwise
|
||||
else:
|
||||
from .layers.linear.linear_fp8 import quantized_linear_forward_fp8_matmul
|
||||
return quantized_linear_forward_fp8_matmul
|
||||
else:
|
||||
from .layers.linear.forward import quantized_linear_forward
|
||||
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)
|
||||
torch._dynamo.config.accumulated_recompile_limit = max(8192, torch._dynamo.config.accumulated_recompile_limit)
|
||||
int8_matmul = torch.compile(int8_matmul, fullgraph=True, dynamic=False)
|
||||
fp8_matmul = torch.compile(fp8_matmul, fullgraph=True, dynamic=False)
|
||||
fp8_matmul_tensorwise = torch.compile(fp8_matmul_tensorwise, fullgraph=True, dynamic=False)
|
||||
conv_int8_matmul = torch.compile(conv_int8_matmul, fullgraph=True, dynamic=False)
|
||||
conv_fp8_matmul = torch.compile(conv_fp8_matmul, fullgraph=True, dynamic=False)
|
||||
conv_fp8_matmul_tensorwise = torch.compile(conv_fp8_matmul_tensorwise, fullgraph=True, dynamic=False)
|
||||
except Exception as e:
|
||||
shared.log.warning(f"Quantization: type=sdnq MatMul using torch.compile is not available: {e}")
|
||||
|
||||
@@ -0,0 +1,71 @@
|
||||
# pylint: disable=relative-beyond-top-level,redefined-builtin,protected-access
|
||||
|
||||
from typing import List
|
||||
|
||||
import torch
|
||||
|
||||
from ...common import use_torch_compile # noqa: TID252
|
||||
from ..linear.linear_fp8 import quantize_fp8_matmul_input # noqa: TID252
|
||||
from .conv import get_conv_args, process_conv_input
|
||||
|
||||
|
||||
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 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,
|
||||
)
|
||||
|
||||
|
||||
if use_torch_compile:
|
||||
try:
|
||||
conv_fp8_matmul = torch.compile(conv_fp8_matmul, fullgraph=True, dynamic=False)
|
||||
except Exception:
|
||||
pass
|
||||
@@ -0,0 +1,70 @@
|
||||
# pylint: disable=relative-beyond-top-level,redefined-builtin,protected-access
|
||||
|
||||
from typing import List
|
||||
|
||||
import torch
|
||||
|
||||
from ...common import use_torch_compile # noqa: TID252
|
||||
from ...dequantizer import dequantize_symmetric, dequantize_symmetric_with_bias # noqa: TID252
|
||||
from ..linear.linear_fp8_tensorwise import quantize_fp8_matmul_input_tensorwise # noqa: TID252
|
||||
from .conv import get_conv_args, process_conv_input
|
||||
|
||||
|
||||
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 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,
|
||||
)
|
||||
|
||||
|
||||
if use_torch_compile:
|
||||
try:
|
||||
conv_fp8_matmul_tensorwise = torch.compile(conv_fp8_matmul_tensorwise, fullgraph=True, dynamic=False)
|
||||
except Exception:
|
||||
pass
|
||||
@@ -0,0 +1,76 @@
|
||||
# pylint: disable=relative-beyond-top-level,redefined-builtin,protected-access
|
||||
|
||||
from typing import List
|
||||
|
||||
import torch
|
||||
|
||||
from ...common import use_torch_compile # noqa: TID252
|
||||
from ...packed_int import unpack_int_symetric # noqa: TID252
|
||||
from ...dequantizer import dequantize_symmetric, dequantize_symmetric_with_bias # noqa: TID252
|
||||
from ..linear.linear_int8 import quantize_int8_matmul_input # noqa: TID252
|
||||
from .conv import get_conv_args, process_conv_input
|
||||
|
||||
|
||||
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_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,
|
||||
)
|
||||
|
||||
|
||||
if use_torch_compile:
|
||||
try:
|
||||
conv_int8_matmul = torch.compile(conv_int8_matmul, fullgraph=True, dynamic=False)
|
||||
except Exception:
|
||||
pass
|
||||
@@ -0,0 +1,93 @@
|
||||
# pylint: disable=relative-beyond-top-level,redefined-builtin,protected-access
|
||||
|
||||
from typing import Optional
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
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 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 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)
|
||||
@@ -0,0 +1,7 @@
|
||||
# pylint: disable=relative-beyond-top-level,redefined-builtin,protected-access
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
def quantized_linear_forward(self, input: torch.FloatTensor) -> torch.FloatTensor:
|
||||
return torch.nn.functional.linear(input, self.sdnq_dequantizer(self.weight), self.bias)
|
||||
@@ -0,0 +1,41 @@
|
||||
# pylint: disable=relative-beyond-top-level,redefined-builtin,protected-access
|
||||
|
||||
from typing import Tuple
|
||||
|
||||
import torch
|
||||
|
||||
from ...common import use_torch_compile # noqa: TID252
|
||||
|
||||
|
||||
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 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)
|
||||
|
||||
|
||||
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)
|
||||
|
||||
|
||||
if use_torch_compile:
|
||||
try:
|
||||
fp8_matmul = torch.compile(fp8_matmul, fullgraph=True, dynamic=False)
|
||||
except Exception:
|
||||
pass
|
||||
@@ -0,0 +1,48 @@
|
||||
# pylint: disable=relative-beyond-top-level,redefined-builtin,protected-access
|
||||
|
||||
from typing import Tuple
|
||||
|
||||
import torch
|
||||
|
||||
from ...common import use_torch_compile # noqa: TID252
|
||||
from ...dequantizer import dequantize_symmetric, dequantize_symmetric_with_bias # noqa: TID252
|
||||
|
||||
|
||||
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 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 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)
|
||||
|
||||
|
||||
if use_torch_compile:
|
||||
try:
|
||||
fp8_matmul_tensorwise = torch.compile(fp8_matmul_tensorwise, fullgraph=True, dynamic=False)
|
||||
except Exception:
|
||||
pass
|
||||
@@ -0,0 +1,52 @@
|
||||
# pylint: disable=relative-beyond-top-level,redefined-builtin,protected-access
|
||||
|
||||
from typing import Tuple
|
||||
|
||||
import torch
|
||||
|
||||
from ...common import use_torch_compile # noqa: TID252
|
||||
from ...packed_int import unpack_int_symetric # noqa: TID252
|
||||
from ...dequantizer import dequantize_symmetric, dequantize_symmetric_with_bias # noqa: TID252
|
||||
|
||||
|
||||
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 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 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)
|
||||
|
||||
|
||||
if use_torch_compile:
|
||||
try:
|
||||
int8_matmul = torch.compile(int8_matmul, fullgraph=True, dynamic=False)
|
||||
except Exception:
|
||||
pass
|
||||
@@ -1,6 +1,7 @@
|
||||
# pylint: disable=redefined-builtin,no-member,protected-access
|
||||
|
||||
from typing import Optional
|
||||
|
||||
import torch
|
||||
|
||||
from .common import dtype_dict
|
||||
|
||||
Reference in New Issue
Block a user