mirror of
https://github.com/vladmandic/automatic
synced 2026-09-19 09:14:35 +02:00
Refactor SDNQDequantizer
This commit is contained in:
Disty0
+30
-21
@@ -7,29 +7,38 @@ from modules import shared, devices
|
||||
|
||||
|
||||
dtype_dict = {
|
||||
"int8": {"min": -128, "max": 127, "num_bits": 8, "target_dtype": torch.int8, "torch_dtype": torch.int8, "storage_dtype": torch.int8, "is_unsigned": False, "is_integer": True},
|
||||
"int7": {"min": -64, "max": 63, "num_bits": 7, "target_dtype": "int7", "torch_dtype": torch.int8, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": True},
|
||||
"int6": {"min": -32, "max": 31, "num_bits": 6, "target_dtype": "int6", "torch_dtype": torch.int8, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": True},
|
||||
"int5": {"min": -16, "max": 15, "num_bits": 5, "target_dtype": "int5", "torch_dtype": torch.int8, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": True},
|
||||
"int4": {"min": -8, "max": 7, "num_bits": 4, "target_dtype": "int4", "torch_dtype": torch.int8, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": True},
|
||||
"int3": {"min": -4, "max": 3, "num_bits": 3, "target_dtype": "int3", "torch_dtype": torch.int8, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": True},
|
||||
"int2": {"min": -2, "max": 1, "num_bits": 2, "target_dtype": "int2", "torch_dtype": torch.int8, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": True},
|
||||
"uint8": {"min": 0, "max": 255, "num_bits": 8, "target_dtype": torch.uint8, "torch_dtype": torch.uint8, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": True},
|
||||
"uint7": {"min": 0, "max": 127, "num_bits": 7, "target_dtype": "uint7", "torch_dtype": torch.uint8, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": True},
|
||||
"uint6": {"min": 0, "max": 63, "num_bits": 6, "target_dtype": "uint6", "torch_dtype": torch.uint8, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": True},
|
||||
"uint5": {"min": 0, "max": 31, "num_bits": 5, "target_dtype": "uint5", "torch_dtype": torch.uint8, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": True},
|
||||
"uint4": {"min": 0, "max": 15, "num_bits": 4, "target_dtype": "uint4", "torch_dtype": torch.uint8, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": True},
|
||||
"uint3": {"min": 0, "max": 7, "num_bits": 3, "target_dtype": "uint3", "torch_dtype": torch.uint8, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": True},
|
||||
"uint2": {"min": 0, "max": 3, "num_bits": 2, "target_dtype": "uint2", "torch_dtype": torch.uint8, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": True},
|
||||
"uint1": {"min": 0, "max": 1, "num_bits": 1, "target_dtype": torch.bool, "torch_dtype": torch.bool, "storage_dtype": torch.bool, "is_unsigned": True, "is_integer": True},
|
||||
"float8_e4m3fn": {"min": -448, "max": 448, "num_bits": 8, "target_dtype": torch.float8_e4m3fn, "torch_dtype": torch.float8_e4m3fn, "storage_dtype": torch.float8_e4m3fn, "is_unsigned": False, "is_integer": False},
|
||||
"float8_e5m2": {"min": -57344, "max": 57344, "num_bits": 8, "target_dtype": torch.float8_e5m2, "torch_dtype": torch.float8_e5m2, "storage_dtype": torch.float8_e5m2, "is_unsigned": False, "is_integer": False},
|
||||
"int32": {"min": -2147483648, "max": 2147483647, "num_bits": 32, "sign": 1, "exponent": 0, "mantissa": 31, "target_dtype": torch.int32, "torch_dtype": torch.int32, "storage_dtype": torch.int32, "is_unsigned": False, "is_integer": True, "is_packed": False},
|
||||
"int16": {"min": -32768, "max": 32767, "num_bits": 16, "sign": 1, "exponent": 0, "mantissa": 15, "target_dtype": torch.int16, "torch_dtype": torch.int16, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": True, "is_packed": False},
|
||||
"int8": {"min": -128, "max": 127, "num_bits": 8, "sign": 1, "exponent": 0, "mantissa": 7, "target_dtype": torch.int8, "torch_dtype": torch.int8, "storage_dtype": torch.int8, "is_unsigned": False, "is_integer": True, "is_packed": False},
|
||||
"int7": {"min": -64, "max": 63, "num_bits": 7, "sign": 1, "exponent": 0, "mantissa": 6, "target_dtype": "int7", "torch_dtype": torch.int8, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": True, "is_packed": True},
|
||||
"int6": {"min": -32, "max": 31, "num_bits": 6, "sign": 1, "exponent": 0, "mantissa": 5, "target_dtype": "int6", "torch_dtype": torch.int8, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": True, "is_packed": True},
|
||||
"int5": {"min": -16, "max": 15, "num_bits": 5, "sign": 1, "exponent": 0, "mantissa": 4, "target_dtype": "int5", "torch_dtype": torch.int8, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": True, "is_packed": True},
|
||||
"int4": {"min": -8, "max": 7, "num_bits": 4, "sign": 1, "exponent": 0, "mantissa": 3, "target_dtype": "int4", "torch_dtype": torch.int8, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": True, "is_packed": True},
|
||||
"int3": {"min": -4, "max": 3, "num_bits": 3, "sign": 1, "exponent": 0, "mantissa": 2, "target_dtype": "int3", "torch_dtype": torch.int8, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": True, "is_packed": True},
|
||||
"int2": {"min": -2, "max": 1, "num_bits": 2, "sign": 1, "exponent": 0, "mantissa": 1, "target_dtype": "int2", "torch_dtype": torch.int8, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": True, "is_packed": True},
|
||||
"uint32": {"min": 0, "max": 4294967295, "num_bits": 32, "sign": 0, "exponent": 0, "mantissa": 31, "target_dtype": torch.uint32, "torch_dtype": torch.uint32, "storage_dtype": torch.uint32, "is_unsigned": True, "is_integer": True, "is_packed": False},
|
||||
"uint16": {"min": 0, "max": 65535, "num_bits": 16, "sign": 0, "exponent": 0, "mantissa": 16, "target_dtype": torch.uint16, "torch_dtype": torch.uint16, "storage_dtype": torch.uint16, "is_unsigned": True, "is_integer": True, "is_packed": False},
|
||||
"uint8": {"min": 0, "max": 255, "num_bits": 8, "sign": 0, "exponent": 0, "mantissa": 8, "target_dtype": torch.uint8, "torch_dtype": torch.uint8, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": True, "is_packed": False},
|
||||
"uint7": {"min": 0, "max": 127, "num_bits": 7, "sign": 0, "exponent": 0, "mantissa": 7, "target_dtype": "uint7", "torch_dtype": torch.uint8, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": True, "is_packed": True},
|
||||
"uint6": {"min": 0, "max": 63, "num_bits": 6, "sign": 0, "exponent": 0, "mantissa": 6, "target_dtype": "uint6", "torch_dtype": torch.uint8, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": True, "is_packed": True},
|
||||
"uint5": {"min": 0, "max": 31, "num_bits": 5, "sign": 0, "exponent": 0, "mantissa": 5, "target_dtype": "uint5", "torch_dtype": torch.uint8, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": True, "is_packed": True},
|
||||
"uint4": {"min": 0, "max": 15, "num_bits": 4, "sign": 0, "exponent": 0, "mantissa": 4, "target_dtype": "uint4", "torch_dtype": torch.uint8, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": True, "is_packed": True},
|
||||
"uint3": {"min": 0, "max": 7, "num_bits": 3, "sign": 0, "exponent": 0, "mantissa": 3, "target_dtype": "uint3", "torch_dtype": torch.uint8, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": True, "is_packed": True},
|
||||
"uint2": {"min": 0, "max": 3, "num_bits": 2, "sign": 0, "exponent": 0, "mantissa": 2, "target_dtype": "uint2", "torch_dtype": torch.uint8, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": True, "is_packed": True},
|
||||
"uint1": {"min": 0, "max": 1, "num_bits": 1, "sign": 0, "exponent": 0, "mantissa": 1, "target_dtype": torch.bool, "torch_dtype": torch.bool, "storage_dtype": torch.bool, "is_unsigned": True, "is_integer": True, "is_packed": True},
|
||||
"float32": {"min": -3.40282e+38, "max": 3.40282e+38, "num_bits": 32, "sign": 1, "exponent": 8, "mantissa": 23, "target_dtype": torch.float32, "torch_dtype": torch.float32, "storage_dtype": torch.float32, "is_unsigned": False, "is_integer": False, "is_packed": False},
|
||||
"bfloat16": {"min": -3.38953e+38, "max": 3.38953e+38, "num_bits": 16, "sign": 1, "exponent": 8, "mantissa": 7, "target_dtype": torch.bfloat16, "torch_dtype": torch.bfloat16, "storage_dtype": torch.bfloat16, "is_unsigned": False, "is_integer": False, "is_packed": False},
|
||||
"float16": {"min": -65504, "max": 65504, "num_bits": 16, "sign": 1, "exponent": 5, "mantissa": 10, "target_dtype": torch.float16, "torch_dtype": torch.float16, "storage_dtype": torch.float16, "is_unsigned": False, "is_integer": False, "is_packed": False},
|
||||
"float8_e4m3fn": {"min": -448, "max": 448, "num_bits": 8, "sign": 1, "exponent": 4, "mantissa": 3, "target_dtype": torch.float8_e4m3fn, "torch_dtype": torch.float8_e4m3fn, "storage_dtype": torch.float8_e4m3fn, "is_unsigned": False, "is_integer": False, "is_packed": False},
|
||||
"float8_e5m2": {"min": -57344, "max": 57344, "num_bits": 8, "sign": 1, "exponent": 5, "mantissa": 2, "target_dtype": torch.float8_e5m2, "torch_dtype": torch.float8_e5m2, "storage_dtype": torch.float8_e5m2, "is_unsigned": False, "is_integer": False, "is_packed": False},
|
||||
}
|
||||
|
||||
dtype_dict["fp8"] = dtype_dict["float8_e4m3fn"]
|
||||
dtype_dict["bool"] = dtype_dict["uint1"]
|
||||
if hasattr(torch, "float8_e4m3fnuz"):
|
||||
dtype_dict["float8_e4m3fnuz"] = {"min": -240, "max": 240, "num_bits": 8, "target_dtype": "fp8", "torch_dtype": torch.float8_e4m3fnuz, "storage_dtype": torch.float8_e4m3fnuz, "is_unsigned": False, "is_integer": False}
|
||||
dtype_dict["float8_e4m3fnuz"] = {"min": -240, "max": 240, "num_bits": 8, "sign": 1, "exponent": 4, "mantissa": 3, "target_dtype": "fp8", "torch_dtype": torch.float8_e4m3fnuz, "storage_dtype": torch.float8_e4m3fnuz, "is_unsigned": False, "is_integer": False, "is_packed": False}
|
||||
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}
|
||||
dtype_dict["float8_e5m2fnuz"] = {"min": -57344, "max": 57344, "num_bits": 8, "sign": 1, "exponent": 5, "mantissa": 2, "target_dtype": "fp8", "torch_dtype": torch.float8_e5m2fnuz, "storage_dtype": torch.float8_e5m2fnuz, "is_unsigned": False, "is_integer": False, "is_packed": False}
|
||||
|
||||
linear_types = {"Linear"}
|
||||
conv_types = {"Conv1d", "Conv2d", "Conv3d"}
|
||||
@@ -43,12 +52,12 @@ is_rdna2 = bool(devices.backend == "rocm" and int(getattr(torch.cuda.get_device_
|
||||
|
||||
if os.environ.get("SDNQ_USE_TENSORWISE_FP8_MM", None) is None:
|
||||
# row-wise FP8 only exist on H100 hardware, sdnq will use software row-wise with tensorwise hardware with this setting
|
||||
use_tensorwise_fp8_matmul = bool(devices.backend == "cuda" and torch.cuda.get_device_capability(devices.device) < (9,0))
|
||||
use_tensorwise_fp8_matmul = bool(devices.backend != "cuda" or (devices.backend == "cuda" and torch.cuda.get_device_capability(devices.device) < (9,0)))
|
||||
else:
|
||||
use_tensorwise_fp8_matmul = os.environ.get("SDNQ_USE_TENSORWISE_FP8_MM", "0").lower() not in {"0", "false", "no"}
|
||||
|
||||
if os.environ.get("SDNQ_USE_CONTIGUOUS_MM", None) is None:
|
||||
use_contiguous_mm = bool(is_rdna2 or devices.backend in {"cpu", "ipex", "zluda"})
|
||||
use_contiguous_mm = bool(is_rdna2 or devices.backend in {"ipex", "mps", "cpu", "openvino", "zluda"})
|
||||
else:
|
||||
use_contiguous_mm = bool(os.environ.get("SDNQ_USE_CONTIGUOUS_MM", "0").lower() not in {"0", "false", "no"})
|
||||
|
||||
|
||||
+122
-172
@@ -1,14 +1,16 @@
|
||||
# pylint: disable=redefined-builtin,no-member,protected-access
|
||||
|
||||
from typing import Tuple, Optional
|
||||
from typing import List, Tuple, Optional
|
||||
|
||||
import torch
|
||||
|
||||
from modules import devices
|
||||
from .common import dtype_dict, compile_func, use_contiguous_mm, use_tensorwise_fp8_matmul
|
||||
from .packed_int import pack_int_symetric, unpack_int_symetric, pack_int_asymetric, unpack_int_asymetric
|
||||
from .packed_int import unpack_int_symetric, unpack_int_asymetric
|
||||
|
||||
|
||||
def dequantize_asymmetric(weight: torch.ByteTensor, scale: torch.FloatTensor, zero_point: torch.FloatTensor, dtype: torch.dtype, result_shape: torch.Size, svd_up: Optional[torch.FloatTensor] = None, svd_down: Optional[torch.FloatTensor] = None, skip_quantized_matmul: bool = False) -> torch.FloatTensor:
|
||||
@devices.inference_context()
|
||||
def dequantize_asymmetric(weight: torch.ByteTensor, scale: torch.FloatTensor, zero_point: torch.FloatTensor, svd_up: Optional[torch.FloatTensor] = None, svd_down: Optional[torch.FloatTensor] = None, dtype: Optional[torch.dtype] = None, result_shape: Optional[torch.Size] = None, skip_quantized_matmul: bool = False) -> torch.FloatTensor:
|
||||
result = torch.addcmul(zero_point, weight.to(dtype=scale.dtype), scale)
|
||||
if result_shape is not None:
|
||||
result = result.view(result_shape)
|
||||
@@ -28,7 +30,8 @@ def dequantize_asymmetric(weight: torch.ByteTensor, scale: torch.FloatTensor, ze
|
||||
return result
|
||||
|
||||
|
||||
def dequantize_symmetric(weight: torch.CharTensor, scale: torch.FloatTensor, dtype: torch.dtype, result_shape: torch.Size, svd_up: Optional[torch.FloatTensor] = None, svd_down: Optional[torch.FloatTensor] = None, skip_quantized_matmul: bool = False) -> torch.FloatTensor:
|
||||
@devices.inference_context()
|
||||
def dequantize_symmetric(weight: torch.CharTensor, scale: torch.FloatTensor, svd_up: Optional[torch.FloatTensor] = None, svd_down: Optional[torch.FloatTensor] = None, dtype: Optional[torch.dtype] = None, result_shape: Optional[torch.Size] = None, skip_quantized_matmul: bool = False) -> torch.FloatTensor:
|
||||
result = weight.to(dtype=scale.dtype).mul_(scale)
|
||||
if skip_quantized_matmul:
|
||||
result.t_()
|
||||
@@ -50,32 +53,72 @@ def dequantize_symmetric(weight: torch.CharTensor, scale: torch.FloatTensor, dty
|
||||
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).view(result_shape)
|
||||
@devices.inference_context()
|
||||
def dequantize_symmetric_with_bias(weight: torch.CharTensor, scale: torch.FloatTensor, bias: torch.FloatTensor, dtype: Optional[torch.dtype] = None, result_shape: Optional[torch.Size] = None) -> torch.FloatTensor:
|
||||
result = torch.addcmul(bias, weight.to(dtype=scale.dtype), scale)
|
||||
if result_shape is not None:
|
||||
result = result.view(result_shape)
|
||||
if dtype is not None:
|
||||
result = result.to(dtype=dtype)
|
||||
return result
|
||||
|
||||
|
||||
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, svd_up: Optional[torch.FloatTensor] = None, svd_down: Optional[torch.FloatTensor] = None, skip_quantized_matmul: bool = False) -> torch.FloatTensor:
|
||||
return dequantize_asymmetric(unpack_int_asymetric(weight, shape, weights_dtype), scale, zero_point, dtype, result_shape, svd_up=svd_up, svd_down=svd_down, skip_quantized_matmul=skip_quantized_matmul)
|
||||
@devices.inference_context()
|
||||
def dequantize_packed_int_asymmetric(weight: torch.ByteTensor, scale: torch.FloatTensor, zero_point: torch.FloatTensor, shape: torch.Size, weights_dtype: str, dtype: Optional[torch.dtype] = None, result_shape: Optional[torch.Size] = None, svd_up: Optional[torch.FloatTensor] = None, svd_down: Optional[torch.FloatTensor] = None, skip_quantized_matmul: bool = False) -> torch.FloatTensor:
|
||||
return dequantize_asymmetric(unpack_int_asymetric(weight, shape, weights_dtype), scale, zero_point, dtype=dtype, result_shape=result_shape, svd_up=svd_up, svd_down=svd_down, skip_quantized_matmul=skip_quantized_matmul)
|
||||
|
||||
|
||||
def dequantize_packed_int_symmetric(weight: torch.ByteTensor, scale: torch.FloatTensor, shape: torch.Size, dtype: torch.dtype, result_shape: torch.Size, weights_dtype: str, svd_up: Optional[torch.FloatTensor] = None, svd_down: Optional[torch.FloatTensor] = None, skip_quantized_matmul: bool = False) -> torch.FloatTensor:
|
||||
return dequantize_symmetric(unpack_int_symetric(weight, shape, weights_dtype, dtype=scale.dtype), scale, dtype, result_shape, svd_up=svd_up, svd_down=svd_down, skip_quantized_matmul=skip_quantized_matmul)
|
||||
@devices.inference_context()
|
||||
def dequantize_packed_int_symmetric(weight: torch.ByteTensor, scale: torch.FloatTensor, shape: torch.Size, weights_dtype: str, dtype: Optional[torch.dtype] = None, result_shape: Optional[torch.Size] = None, svd_up: Optional[torch.FloatTensor] = None, svd_down: Optional[torch.FloatTensor] = None, skip_quantized_matmul: bool = False) -> torch.FloatTensor:
|
||||
return dequantize_symmetric(unpack_int_symetric(weight, shape, weights_dtype, dtype=scale.dtype), scale, dtype=dtype, result_shape=result_shape, svd_up=svd_up, svd_down=svd_down, skip_quantized_matmul=skip_quantized_matmul)
|
||||
|
||||
|
||||
@devices.inference_context()
|
||||
def quantize_int8(input: torch.FloatTensor, dim: int = -1) -> Tuple[torch.CharTensor, torch.FloatTensor]:
|
||||
scale = torch.amax(input.abs(), dim=dim, keepdims=True).div_(127)
|
||||
input = torch.div(input, scale).round_().clamp_(-128, 127).to(dtype=torch.int8)
|
||||
return input, scale
|
||||
|
||||
|
||||
@devices.inference_context()
|
||||
def quantize_int8_sr(input: torch.FloatTensor, dim: int = -1) -> Tuple[torch.CharTensor, torch.FloatTensor]:
|
||||
scale = torch.amax(input.abs(), dim=dim, keepdims=True).div_(127)
|
||||
input = torch.normal(0, 0.1, input.shape, device=input.device, dtype=input.dtype
|
||||
).addcdiv_(input, scale).round_().clamp_(-128, 127).to(dtype=torch.int8)
|
||||
return input, scale
|
||||
|
||||
|
||||
@devices.inference_context()
|
||||
def quantize_fp8(input: torch.FloatTensor, dim: int = -1, is_e5: bool = False) -> Tuple[torch.Tensor, torch.FloatTensor]:
|
||||
max_range = 57344 if is_e5 else 448
|
||||
fp8_dtype = torch.float8_e5m2 if is_e5 else torch.float8_e4m3fn
|
||||
if is_e5:
|
||||
max_range = 57344
|
||||
fp8_dtype = torch.float8_e5m2
|
||||
else:
|
||||
max_range = 448
|
||||
fp8_dtype = torch.float8_e4m3fn
|
||||
scale = torch.amax(input.abs(), dim=dim, keepdims=True).div_(max_range)
|
||||
input = torch.div(input, scale).nan_to_num_().clamp_(-max_range, max_range).to(dtype=fp8_dtype)
|
||||
return input, scale
|
||||
|
||||
|
||||
@devices.inference_context()
|
||||
def quantize_fp8_sr(input: torch.FloatTensor, dim: int = -1, is_e5: bool = False) -> Tuple[torch.Tensor, torch.FloatTensor]:
|
||||
if is_e5:
|
||||
max_range = 57344
|
||||
fp8_dtype = torch.float8_e5m2
|
||||
mantissa_difference = 2097152
|
||||
else:
|
||||
max_range = 448
|
||||
fp8_dtype = torch.float8_e4m3fn
|
||||
mantissa_difference = 1048576
|
||||
scale = torch.amax(input.abs(), dim=dim, keepdims=True).div_(max_range)
|
||||
input = torch.div(input, scale).to(dtype=torch.float32).view(dtype=torch.int32)
|
||||
input = input.add_(torch.randint_like(input, low=0, high=mantissa_difference)).bitwise_and_(-mantissa_difference).view(dtype=torch.float32)
|
||||
input = input.nan_to_num_().clamp_(-max_range, max_range).to(dtype=fp8_dtype)
|
||||
return input, scale
|
||||
|
||||
|
||||
@devices.inference_context()
|
||||
def re_quantize_int8(weight: torch.FloatTensor) -> Tuple[torch.CharTensor, torch.FloatTensor]:
|
||||
if weight.ndim > 2: # convs
|
||||
weight = weight.flatten(1,-1)
|
||||
@@ -88,6 +131,7 @@ def re_quantize_int8(weight: torch.FloatTensor) -> Tuple[torch.CharTensor, torch
|
||||
return weight, scale
|
||||
|
||||
|
||||
@devices.inference_context()
|
||||
def re_quantize_fp8(weight: torch.FloatTensor, is_e5: bool = False) -> Tuple[torch.CharTensor, torch.FloatTensor]:
|
||||
if weight.ndim > 2: # convs
|
||||
weight = weight.flatten(1,-1)
|
||||
@@ -98,207 +142,113 @@ def re_quantize_fp8(weight: torch.FloatTensor, is_e5: bool = False) -> Tuple[tor
|
||||
return weight, scale
|
||||
|
||||
|
||||
def re_quantize_matmul_asymmetric(weight: torch.ByteTensor, scale: torch.FloatTensor, zero_point: torch.FloatTensor, result_shape: torch.Size, svd_up: Optional[torch.FloatTensor] = None, svd_down: Optional[torch.FloatTensor] = None) -> Tuple[torch.CharTensor, torch.FloatTensor]:
|
||||
return re_quantize_int8(dequantize_asymmetric(weight, scale, zero_point, scale.dtype, result_shape, svd_up=svd_up, svd_down=svd_down))
|
||||
@devices.inference_context()
|
||||
def re_quantize_matmul_asymmetric(weight: torch.ByteTensor, scale: torch.FloatTensor, zero_point: torch.FloatTensor, result_shape: Optional[torch.Size] = None, svd_up: Optional[torch.FloatTensor] = None, svd_down: Optional[torch.FloatTensor] = None) -> Tuple[torch.CharTensor, torch.FloatTensor]:
|
||||
return re_quantize_int8(dequantize_asymmetric(weight, scale, zero_point, dtype=scale.dtype, result_shape=result_shape, svd_up=svd_up, svd_down=svd_down))
|
||||
|
||||
|
||||
def re_quantize_matmul_symmetric(weight: torch.CharTensor, scale: torch.FloatTensor, result_shape: torch.Size, svd_up: Optional[torch.FloatTensor] = None, svd_down: Optional[torch.FloatTensor] = None) -> Tuple[torch.CharTensor, torch.FloatTensor]:
|
||||
return re_quantize_int8(dequantize_symmetric(weight, scale, scale.dtype, result_shape, svd_up=svd_up, svd_down=svd_down))
|
||||
@devices.inference_context()
|
||||
def re_quantize_matmul_symmetric(weight: torch.CharTensor, scale: torch.FloatTensor, result_shape: Optional[torch.Size] = None, svd_up: Optional[torch.FloatTensor] = None, svd_down: Optional[torch.FloatTensor] = None) -> Tuple[torch.CharTensor, torch.FloatTensor]:
|
||||
return re_quantize_int8(dequantize_symmetric(weight, scale, dtype=scale.dtype, result_shape=result_shape, svd_up=svd_up, svd_down=svd_down))
|
||||
|
||||
|
||||
def re_quantize_matmul_packed_int_asymmetric(weight: torch.ByteTensor, scale: torch.FloatTensor, zero_point: torch.FloatTensor, shape: torch.Size, result_shape: torch.Size, weights_dtype: str, svd_up: Optional[torch.FloatTensor] = None, svd_down: Optional[torch.FloatTensor] = None) -> Tuple[torch.CharTensor, torch.FloatTensor]:
|
||||
return re_quantize_matmul_asymmetric(unpack_int_asymetric(weight, shape, weights_dtype), scale, zero_point, result_shape, svd_up=svd_up, svd_down=svd_down)
|
||||
@devices.inference_context()
|
||||
def re_quantize_matmul_packed_int_asymmetric(weight: torch.ByteTensor, scale: torch.FloatTensor, zero_point: torch.FloatTensor, shape: torch.Size, weights_dtype: str, result_shape: torch.Size, svd_up: Optional[torch.FloatTensor] = None, svd_down: Optional[torch.FloatTensor] = None) -> Tuple[torch.CharTensor, torch.FloatTensor]:
|
||||
return re_quantize_matmul_asymmetric(unpack_int_asymetric(weight, shape, weights_dtype), scale, zero_point, result_shape=result_shape, svd_up=svd_up, svd_down=svd_down)
|
||||
|
||||
|
||||
def re_quantize_matmul_packed_int_symmetric(weight: torch.ByteTensor, scale: torch.FloatTensor, shape: torch.Size, result_shape: torch.Size, weights_dtype: str, svd_up: Optional[torch.FloatTensor] = None, svd_down: Optional[torch.FloatTensor] = None) -> Tuple[torch.CharTensor, torch.FloatTensor]:
|
||||
return re_quantize_matmul_symmetric(unpack_int_symetric(weight, shape, weights_dtype, dtype=scale.dtype), scale, result_shape, svd_up=svd_up, svd_down=svd_down)
|
||||
@devices.inference_context()
|
||||
def re_quantize_matmul_packed_int_symmetric(weight: torch.ByteTensor, scale: torch.FloatTensor, shape: torch.Size, weights_dtype: str, result_shape: Optional[torch.Size] = None, svd_up: Optional[torch.FloatTensor] = None, svd_down: Optional[torch.FloatTensor] = None) -> Tuple[torch.CharTensor, torch.FloatTensor]:
|
||||
return re_quantize_matmul_symmetric(unpack_int_symetric(weight, shape, weights_dtype, dtype=scale.dtype), scale, result_shape=result_shape, svd_down=svd_down)
|
||||
|
||||
|
||||
def dequantize_sdnq_model(model):
|
||||
@devices.inference_context()
|
||||
def dequantize_layer_weight(self: torch.nn.Module, inplace: bool = False):
|
||||
weight = self.sdnq_dequantizer(self.weight, self.scale, self.zero_point, self.svd_up, self.svd_down, skip_quantized_matmul=self.sdnq_dequantizer.use_quantized_matmul)
|
||||
if inplace:
|
||||
self.weight.data = weight
|
||||
self.forward = getattr(torch.nn, self.sdnq_dequantizer.layer_class_name).forward
|
||||
del self.sdnq_dequantizer, self.scale, self.zero_point, self.svd_up, self.svd_down
|
||||
return weight
|
||||
|
||||
|
||||
@devices.inference_context()
|
||||
def dequantize_sdnq_model(model: torch.nn.Module):
|
||||
if hasattr(model, "sdnq_dequantizer"):
|
||||
model.weight = torch.nn.Parameter(model.sdnq_dequantizer(model.weight, model.scale, model.zero_point, model.svd_up, model.svd_down))
|
||||
del model.sdnq_dequantizer, model.scale, model.zero_point, model.svd_up, model.svd_down
|
||||
return model
|
||||
model.weight.data = dequantize_layer_weight(model, inplace=True)
|
||||
has_children = list(model.children())
|
||||
if not has_children:
|
||||
return model
|
||||
for module in model.children():
|
||||
for module_name, module in model.named_children():
|
||||
if hasattr(module, "sdnq_dequantizer"):
|
||||
module.weight = torch.nn.Parameter(module.sdnq_dequantizer(module.weight, module.scale, module.zero_point, module.svd_up, module.svd_down))
|
||||
del module.sdnq_dequantizer, module.scale, module.zero_point, module.svd_up, module.svd_down
|
||||
module.weight.data = dequantize_layer_weight(module, inplace=True)
|
||||
setattr(model, module_name, module)
|
||||
else:
|
||||
module = dequantize_sdnq_model(module)
|
||||
setattr(model, module_name, dequantize_sdnq_model(module))
|
||||
return model
|
||||
|
||||
|
||||
class AsymmetricWeightsDequantizer(torch.nn.Module):
|
||||
class SDNQDequantizer():
|
||||
def __init__(
|
||||
self,
|
||||
result_dtype: torch.dtype,
|
||||
result_shape: torch.Size,
|
||||
original_shape: torch.Size,
|
||||
original_stride: List[int],
|
||||
quantized_weight_shape: torch.Size,
|
||||
weights_dtype: str,
|
||||
group_size: int,
|
||||
svd_rank: int,
|
||||
svd_steps: int,
|
||||
use_quantized_matmul: bool,
|
||||
re_quantize_for_matmul: bool,
|
||||
use_stochastic_rounding: bool,
|
||||
layer_class_name: str,
|
||||
):
|
||||
super().__init__()
|
||||
self.is_packed = False
|
||||
self.is_asym = True
|
||||
self.is_packed = dtype_dict[weights_dtype]["is_packed"]
|
||||
self.is_unsigned = dtype_dict[weights_dtype]["is_unsigned"]
|
||||
self.result_dtype = result_dtype
|
||||
self.result_shape = result_shape
|
||||
self.original_shape = original_shape
|
||||
self.original_stride = original_stride
|
||||
self.quantized_weight_shape = quantized_weight_shape
|
||||
self.weights_dtype = weights_dtype
|
||||
self.group_size = group_size
|
||||
self.svd_rank = svd_rank
|
||||
self.svd_steps = svd_steps
|
||||
self.use_quantized_matmul = use_quantized_matmul
|
||||
self.re_quantize_for_matmul = re_quantize_for_matmul
|
||||
self.use_stochastic_rounding = use_stochastic_rounding
|
||||
self.layer_class_name = layer_class_name
|
||||
|
||||
def pack_weight(self, weight: torch.Tensor) -> torch.Tensor:
|
||||
return weight.to(dtype=dtype_dict[self.weights_dtype]["torch_dtype"])
|
||||
@devices.inference_context()
|
||||
def re_quantize_matmul(self, weight, scale, zero_point, svd_up, svd_down): # pylint: disable=unused-argument
|
||||
if self.is_packed:
|
||||
if self.is_unsigned:
|
||||
return re_quantize_matmul_packed_int_asymmetric_compiled(weight, scale, zero_point, self.quantized_weight_shape, self.weights_dtype, result_shape=self.result_shape, svd_up=svd_up, svd_down=svd_down)
|
||||
else:
|
||||
return re_quantize_matmul_packed_int_symmetric_compiled(weight, scale, self.quantized_weight_shape, self.weights_dtype, result_shape=self.result_shape, svd_up=svd_up, svd_down=svd_down)
|
||||
else:
|
||||
if self.is_unsigned:
|
||||
return re_quantize_matmul_asymmetric_compiled(weight, scale, zero_point, result_shape=self.result_shape, svd_up=svd_up, svd_down=svd_down)
|
||||
else:
|
||||
return re_quantize_matmul_symmetric_compiled(weight, scale, result_shape=self.result_shape, svd_up=svd_up, svd_down=svd_down)
|
||||
|
||||
def re_quantize_matmul(self, weight, scale, zero_point, svd_up, svd_down, **kwargs): # pylint: disable=unused-argument
|
||||
return re_quantize_matmul_asymmetric_compiled(weight, scale, zero_point, self.result_shape, svd_up=svd_up, svd_down=svd_down)
|
||||
|
||||
def forward(self, weight, scale, zero_point, svd_up, svd_down, skip_quantized_matmul=False): # pylint: disable=unused-argument
|
||||
return dequantize_asymmetric_compiled(weight, scale, zero_point, self.result_dtype, self.result_shape, svd_up=svd_up, svd_down=svd_down, skip_quantized_matmul=skip_quantized_matmul)
|
||||
|
||||
|
||||
class SymmetricWeightsDequantizer(torch.nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
result_dtype: torch.dtype,
|
||||
result_shape: torch.Size,
|
||||
original_shape: torch.Size,
|
||||
quantized_weight_shape: torch.Size,
|
||||
weights_dtype: str,
|
||||
group_size: int,
|
||||
svd_rank: int,
|
||||
use_quantized_matmul: bool,
|
||||
re_quantize_for_matmul: bool,
|
||||
):
|
||||
super().__init__()
|
||||
self.is_packed = False
|
||||
self.is_asym = False
|
||||
self.result_dtype = result_dtype
|
||||
self.result_shape = result_shape
|
||||
self.original_shape = original_shape
|
||||
self.quantized_weight_shape = quantized_weight_shape
|
||||
self.weights_dtype = weights_dtype
|
||||
self.group_size = group_size
|
||||
self.svd_rank = svd_rank
|
||||
self.use_quantized_matmul = use_quantized_matmul
|
||||
self.re_quantize_for_matmul = re_quantize_for_matmul
|
||||
|
||||
def pack_weight(self, weight: torch.Tensor) -> torch.Tensor:
|
||||
return weight.to(dtype=dtype_dict[self.weights_dtype]["torch_dtype"])
|
||||
|
||||
def re_quantize_matmul(self, weight, scale, zero_point, svd_up, svd_down, **kwargs): # pylint: disable=unused-argument
|
||||
return re_quantize_matmul_symmetric_compiled(weight, scale, self.result_shape, svd_up=svd_up, svd_down=svd_down)
|
||||
|
||||
def forward(self, weight, scale, zero_point, svd_up, svd_down, skip_quantized_matmul=False): # pylint: disable=unused-argument
|
||||
@devices.inference_context()
|
||||
def __call__(self, weight, scale, zero_point, svd_up, svd_down, skip_quantized_matmul: bool = False, dtype: torch.dtype = None): # pylint: disable=unused-argument
|
||||
skip_quantized_matmul = skip_quantized_matmul and not self.re_quantize_for_matmul
|
||||
return dequantize_symmetric_compiled(weight, scale, self.result_dtype, self.result_shape, svd_up=svd_up, svd_down=svd_down, skip_quantized_matmul=skip_quantized_matmul)
|
||||
|
||||
|
||||
class PackedINTAsymmetricWeightsDequantizer(torch.nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
result_dtype: torch.dtype,
|
||||
result_shape: torch.Size,
|
||||
original_shape: torch.Size,
|
||||
quantized_weight_shape: torch.Size,
|
||||
weights_dtype: str,
|
||||
group_size: int,
|
||||
svd_rank: int,
|
||||
use_quantized_matmul: bool,
|
||||
re_quantize_for_matmul: bool,
|
||||
):
|
||||
super().__init__()
|
||||
self.is_packed = True
|
||||
self.is_asym = True
|
||||
self.result_dtype = result_dtype
|
||||
self.result_shape = result_shape
|
||||
self.original_shape = original_shape
|
||||
self.quantized_weight_shape = quantized_weight_shape
|
||||
self.weights_dtype = weights_dtype
|
||||
self.group_size = group_size
|
||||
self.svd_rank = svd_rank
|
||||
self.use_quantized_matmul = use_quantized_matmul
|
||||
self.re_quantize_for_matmul = re_quantize_for_matmul
|
||||
|
||||
def pack_weight(self, weight: torch.Tensor) -> torch.Tensor:
|
||||
return pack_int_asymetric(weight, self.weights_dtype)
|
||||
|
||||
def re_quantize_matmul(self, weight, scale, zero_point, svd_up, svd_down, **kwargs): # pylint: disable=unused-argument
|
||||
return re_quantize_matmul_packed_int_asymmetric_compiled(weight, scale, zero_point, self.quantized_weight_shape, self.result_shape, self.weights_dtype, svd_up=svd_up, svd_down=svd_down)
|
||||
|
||||
def forward(self, weight, scale, zero_point, svd_up, svd_down, skip_quantized_matmul=False): # pylint: disable=unused-argument
|
||||
return dequantize_packed_int_asymmetric_compiled(weight, scale, zero_point, self.quantized_weight_shape, self.result_dtype, self.result_shape, self.weights_dtype, svd_up=svd_up, svd_down=svd_down, skip_quantized_matmul=skip_quantized_matmul)
|
||||
|
||||
|
||||
class PackedINTSymmetricWeightsDequantizer(torch.nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
result_dtype: torch.dtype,
|
||||
result_shape: torch.Size,
|
||||
original_shape: torch.Size,
|
||||
quantized_weight_shape: torch.Size,
|
||||
weights_dtype: str,
|
||||
group_size: int,
|
||||
svd_rank: int,
|
||||
use_quantized_matmul: bool,
|
||||
re_quantize_for_matmul: bool,
|
||||
):
|
||||
super().__init__()
|
||||
self.is_packed = True
|
||||
self.is_asym = False
|
||||
self.result_dtype = result_dtype
|
||||
self.result_shape = result_shape
|
||||
self.original_shape = original_shape
|
||||
self.quantized_weight_shape = quantized_weight_shape
|
||||
self.weights_dtype = weights_dtype
|
||||
self.group_size = group_size
|
||||
self.svd_rank = svd_rank
|
||||
self.use_quantized_matmul = use_quantized_matmul
|
||||
self.re_quantize_for_matmul = re_quantize_for_matmul
|
||||
|
||||
def pack_weight(self, weight: torch.Tensor) -> torch.Tensor:
|
||||
return pack_int_symetric(weight, self.weights_dtype)
|
||||
|
||||
def re_quantize_matmul(self, weight, scale, zero_point, svd_up, svd_down, **kwargs): # pylint: disable=unused-argument
|
||||
return re_quantize_matmul_packed_int_symmetric_compiled(weight, scale, self.quantized_weight_shape, self.result_shape, self.weights_dtype, svd_up=svd_up, svd_down=svd_down)
|
||||
|
||||
def forward(self, weight, scale, zero_point, svd_up, svd_down, skip_quantized_matmul=False): # pylint: disable=unused-argument
|
||||
skip_quantized_matmul = skip_quantized_matmul and not self.re_quantize_for_matmul
|
||||
return dequantize_packed_int_symmetric_compiled(weight, scale, self.quantized_weight_shape, self.result_dtype, self.result_shape, self.weights_dtype, svd_up=svd_up, svd_down=svd_down, skip_quantized_matmul=skip_quantized_matmul)
|
||||
|
||||
|
||||
dequantizer_dict = {
|
||||
"int8": SymmetricWeightsDequantizer,
|
||||
"int7": PackedINTSymmetricWeightsDequantizer,
|
||||
"int6": PackedINTSymmetricWeightsDequantizer,
|
||||
"int5": PackedINTSymmetricWeightsDequantizer,
|
||||
"int4": PackedINTSymmetricWeightsDequantizer,
|
||||
"int3": PackedINTSymmetricWeightsDequantizer,
|
||||
"int2": PackedINTSymmetricWeightsDequantizer,
|
||||
"uint8": AsymmetricWeightsDequantizer,
|
||||
"uint7": PackedINTAsymmetricWeightsDequantizer,
|
||||
"uint6": PackedINTAsymmetricWeightsDequantizer,
|
||||
"uint5": PackedINTAsymmetricWeightsDequantizer,
|
||||
"uint4": PackedINTAsymmetricWeightsDequantizer,
|
||||
"uint3": PackedINTAsymmetricWeightsDequantizer,
|
||||
"uint2": PackedINTAsymmetricWeightsDequantizer,
|
||||
"uint1": PackedINTAsymmetricWeightsDequantizer,
|
||||
"bool": PackedINTAsymmetricWeightsDequantizer,
|
||||
"float8_e4m3fn": SymmetricWeightsDequantizer,
|
||||
"float8_e4m3fnuz": SymmetricWeightsDequantizer,
|
||||
"float8_e5m2": SymmetricWeightsDequantizer,
|
||||
"float8_e5m2fnuz": SymmetricWeightsDequantizer,
|
||||
}
|
||||
if dtype is None:
|
||||
dtype = self.result_dtype
|
||||
if self.is_packed:
|
||||
if self.is_unsigned:
|
||||
return dequantize_packed_int_asymmetric_compiled(weight, scale, zero_point, self.quantized_weight_shape, self.weights_dtype, dtype=dtype, result_shape=self.result_shape, svd_up=svd_up, svd_down=svd_down, skip_quantized_matmul=skip_quantized_matmul)
|
||||
else:
|
||||
return dequantize_packed_int_symmetric_compiled(weight, scale, self.quantized_weight_shape, self.weights_dtype, dtype=dtype, result_shape=self.result_shape, svd_up=svd_up, svd_down=svd_down, skip_quantized_matmul=skip_quantized_matmul)
|
||||
else:
|
||||
if self.is_unsigned:
|
||||
return dequantize_asymmetric_compiled(weight, scale, zero_point, dtype=dtype, result_shape=self.result_shape, svd_up=svd_up, svd_down=svd_down, skip_quantized_matmul=skip_quantized_matmul)
|
||||
else:
|
||||
return dequantize_symmetric_compiled(weight, scale, dtype=dtype, result_shape=self.result_shape, svd_up=svd_up, svd_down=svd_down, skip_quantized_matmul=skip_quantized_matmul)
|
||||
|
||||
|
||||
dequantize_asymmetric_compiled = compile_func(dequantize_asymmetric)
|
||||
|
||||
@@ -47,9 +47,9 @@ def conv_fp8_matmul_tensorwise(
|
||||
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)
|
||||
dequantize_symmetric_with_bias(result, scale, bias, dtype=return_dtype, result_shape=mm_output_shape)
|
||||
else:
|
||||
dequantize_symmetric(result, scale, return_dtype, mm_output_shape)
|
||||
dequantize_symmetric(result, scale, dtype=return_dtype, result_shape=mm_output_shape)
|
||||
|
||||
if conv_type == 1:
|
||||
result = result.transpose_(1,2)
|
||||
|
||||
@@ -51,9 +51,9 @@ def conv_int8_matmul(
|
||||
result.append(int_mm_func(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)
|
||||
result = dequantize_symmetric_with_bias(result, scale, bias, dtype=return_dtype, result_shape=mm_output_shape)
|
||||
else:
|
||||
result = dequantize_symmetric(result, scale, return_dtype, mm_output_shape)
|
||||
result = dequantize_symmetric(result, scale, dtype=return_dtype, result_shape=mm_output_shape)
|
||||
|
||||
if conv_type == 1:
|
||||
result = result.transpose_(1,2)
|
||||
|
||||
@@ -38,9 +38,9 @@ def fp8_matmul_tensorwise(
|
||||
input, scale = quantize_fp8_matmul_input_tensorwise(input, scale)
|
||||
input, weight = check_mats(input, weight)
|
||||
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)
|
||||
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, 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=scale.dtype), scale, return_dtype, output_shape)
|
||||
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, dtype=return_dtype, result_shape=output_shape)
|
||||
|
||||
|
||||
def quantized_linear_forward_fp8_matmul_tensorwise(self, input: torch.FloatTensor) -> torch.FloatTensor:
|
||||
|
||||
@@ -42,9 +42,9 @@ def int8_matmul(
|
||||
input, scale = quantize_int8_matmul_input(input, scale)
|
||||
input, weight = check_mats(input, weight)
|
||||
if bias is not None:
|
||||
return dequantize_symmetric_with_bias(int_mm_func(input, weight), scale, bias, return_dtype, output_shape)
|
||||
return dequantize_symmetric_with_bias(int_mm_func(input, weight), scale, bias, dtype=return_dtype, result_shape=output_shape)
|
||||
else:
|
||||
return dequantize_symmetric(int_mm_func(input, weight), scale, return_dtype, output_shape)
|
||||
return dequantize_symmetric(int_mm_func(input, weight), scale, dtype=return_dtype, result_shape=output_shape)
|
||||
|
||||
|
||||
def quantized_linear_forward_int8_matmul(self, input: torch.FloatTensor) -> torch.FloatTensor:
|
||||
|
||||
@@ -181,11 +181,11 @@ def apply_options_to_model(model, dtype: torch.dtype = None, dequantize_fp32: bo
|
||||
if use_quantized_matmul and module.sdnq_dequantizer.re_quantize_for_matmul:
|
||||
scale_dtype = module.scale.dtype
|
||||
if module.sdnq_dequantizer.weights_dtype == "int8":
|
||||
module.weight.data, module.scale.data = re_quantize_int8(dequantize_symmetric(module.weight, module.scale, torch.float32, module.sdnq_dequantizer.result_shape))
|
||||
module.weight.data, module.scale.data = re_quantize_int8(dequantize_symmetric(module.weight, module.scale, dtype=torch.float32, result_shape=module.sdnq_dequantizer.result_shape))
|
||||
module.scale.data = module.scale.to(dtype=scale_dtype)
|
||||
else:
|
||||
is_e5 = bool(module.sdnq_dequantizer.weights_dtype == "float8_e5m2")
|
||||
module.weight.data, module.scale.data = re_quantize_fp8(dequantize_symmetric(module.weight, module.scale, torch.float32, module.sdnq_dequantizer.result_shape), is_e5=is_e5)
|
||||
module.weight.data, module.scale.data = re_quantize_fp8(dequantize_symmetric(module.weight, module.scale, dtype=torch.float32, result_shape=module.sdnq_dequantizer.result_shape), is_e5=is_e5)
|
||||
if use_tensorwise_fp8_matmul:
|
||||
module.scale.data = module.scale.to(dtype=scale_dtype)
|
||||
elif not module.sdnq_dequantizer.re_quantize_for_matmul:
|
||||
|
||||
+250
-176
@@ -16,8 +16,9 @@ from accelerate import init_empty_weights
|
||||
from accelerate.utils import set_module_tensor_to_device
|
||||
|
||||
from modules import devices, shared
|
||||
from .common import dtype_dict, common_skip_keys, module_skip_keys_dict, accepted_weights, use_tensorwise_fp8_matmul, allowed_types, conv_types, conv_transpose_types, use_contiguous_mm
|
||||
from .dequantizer import dequantizer_dict, dequantize_sdnq_model
|
||||
from .common import dtype_dict, common_skip_keys, module_skip_keys_dict, accepted_weights, allowed_types, linear_types, conv_types, conv_transpose_types, compile_func, use_tensorwise_fp8_matmul, use_contiguous_mm
|
||||
from .dequantizer import SDNQDequantizer, dequantize_sdnq_model
|
||||
from .packed_int import pack_int_symetric, pack_int_asymetric
|
||||
from .forward import get_forward_func
|
||||
|
||||
|
||||
@@ -25,6 +26,7 @@ class QuantizationMethod(str, Enum):
|
||||
SDNQ = "sdnq"
|
||||
|
||||
|
||||
@devices.inference_context()
|
||||
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"])
|
||||
@@ -33,41 +35,63 @@ def get_scale_asymmetric(weight: torch.FloatTensor, reduction_axes: Union[int, L
|
||||
return scale, zero_point
|
||||
|
||||
|
||||
@devices.inference_context()
|
||||
def get_scale_symmetric(weight: torch.FloatTensor, reduction_axes: Union[int, List[int]], weights_dtype: str) -> torch.FloatTensor:
|
||||
return torch.amax(weight.abs(), dim=reduction_axes, keepdims=True).div_(dtype_dict[weights_dtype]["max"])
|
||||
|
||||
|
||||
def quantize_weight(weight: torch.FloatTensor, reduction_axes: Union[int, List[int]], weights_dtype: str) -> Tuple[torch.Tensor, torch.FloatTensor, torch.FloatTensor]:
|
||||
@devices.inference_context()
|
||||
def quantize_weight(weight: torch.FloatTensor, reduction_axes: Union[int, List[int]], weights_dtype: str, use_stochastic_rounding: bool = False) -> Tuple[torch.Tensor, torch.FloatTensor, torch.FloatTensor]:
|
||||
weight = weight.to(dtype=torch.float32)
|
||||
|
||||
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)
|
||||
quantized_weight = torch.sub(weight, zero_point)
|
||||
scale_inplace = True
|
||||
else:
|
||||
scale = get_scale_symmetric(weight, reduction_axes, weights_dtype)
|
||||
quantized_weight = torch.div(weight, scale)
|
||||
quantized_weight = weight
|
||||
scale_inplace = False
|
||||
zero_point = None
|
||||
if dtype_dict[weights_dtype]["is_integer"]:
|
||||
|
||||
is_integer = dtype_dict[weights_dtype]["is_integer"]
|
||||
if use_stochastic_rounding and is_integer: # this case can be fused with addcdiv_
|
||||
quantized_weight = torch.normal(0, 0.1, weight.shape, device=weight.device, dtype=weight.dtype).addcdiv_(quantized_weight, scale)
|
||||
elif scale_inplace:
|
||||
quantized_weight.div_(scale)
|
||||
else:
|
||||
quantized_weight = torch.div(quantized_weight, scale)
|
||||
|
||||
if is_integer:
|
||||
quantized_weight.round_()
|
||||
else:
|
||||
if use_stochastic_rounding:
|
||||
mantissa_difference = 1 << (23 - dtype_dict[weights_dtype]["mantissa"])
|
||||
quantized_weight = quantized_weight.view(dtype=torch.int32)
|
||||
quantized_weight = torch.randint_like(quantized_weight, low=0, high=mantissa_difference).add_(quantized_weight).bitwise_and_(-mantissa_difference).view(dtype=torch.float32)
|
||||
quantized_weight.nan_to_num_()
|
||||
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 apply_svdquant(weight: torch.FloatTensor, rank: int = 32, niter: int = 8) -> Tuple[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]:
|
||||
reshape_weight = False
|
||||
if weight.ndim > 2: # convs
|
||||
reshape_weight = True
|
||||
weight_shape = weight.shape
|
||||
weight = weight.flatten(1,-1)
|
||||
weight = weight.to(dtype=torch.float32)
|
||||
U, S, svd_down = torch.svd_lowrank(weight, q=rank, niter=niter)
|
||||
svd_up = torch.mul(U, S.unsqueeze(0))
|
||||
svd_down = svd_down.t_()
|
||||
weight = weight.sub_(torch.mm(svd_up, svd_down))
|
||||
weight = weight.sub(torch.mm(svd_up, svd_down))
|
||||
if reshape_weight:
|
||||
weight = weight.unflatten(-1, (*weight_shape[1:],)) # pylint: disable=possibly-used-before-assignment
|
||||
return weight, svd_up, svd_down
|
||||
|
||||
|
||||
@devices.inference_context()
|
||||
def prepare_weight_for_matmul(weight: torch.Tensor) -> torch.Tensor:
|
||||
if use_contiguous_mm:
|
||||
weight = weight.contiguous()
|
||||
@@ -76,6 +100,7 @@ def prepare_weight_for_matmul(weight: torch.Tensor) -> torch.Tensor:
|
||||
return weight
|
||||
|
||||
|
||||
@devices.inference_context()
|
||||
def prepare_svd_for_matmul(svd_up: torch.FloatTensor, svd_down: torch.FloatTensor, use_quantized_matmul: bool) -> Tuple[torch.FloatTensor, torch.FloatTensor]:
|
||||
if svd_up is not None:
|
||||
if use_quantized_matmul:
|
||||
@@ -163,182 +188,220 @@ def add_module_skip_keys(model, modules_to_not_convert: List[str] = None, module
|
||||
|
||||
|
||||
@devices.inference_context()
|
||||
def sdnq_quantize_layer(layer, weights_dtype="int8", torch_dtype=None, group_size=0, svd_rank=32, svd_steps=8, use_svd=False, quant_conv=False, use_quantized_matmul=False, use_quantized_matmul_conv=False, dequantize_fp32=False, non_blocking=False, quantization_device=None, return_device=None, param_name=None): # pylint: disable=unused-argument
|
||||
layer_class_name = layer.__class__.__name__
|
||||
if layer_class_name in allowed_types:
|
||||
num_of_groups = 1
|
||||
is_conv_type = False
|
||||
is_conv_transpose_type = False
|
||||
is_linear_type = False
|
||||
result_shape = None
|
||||
original_shape = layer.weight.shape
|
||||
if torch_dtype is None:
|
||||
torch_dtype = layer.weight.dtype
|
||||
def sdnq_quantize_layer_weight(weight, layer_class_name=None, weights_dtype="int8", torch_dtype=None, group_size=0, svd_rank=32, svd_steps=8, use_svd=False, use_quantized_matmul=False, use_stochastic_rounding=False, dequantize_fp32=False, param_name=None): # pylint: disable=unused-argument
|
||||
num_of_groups = 1
|
||||
is_conv_type = False
|
||||
is_conv_transpose_type = False
|
||||
is_linear_type = False
|
||||
result_shape = None
|
||||
original_shape = weight.shape
|
||||
original_stride = weight.stride()
|
||||
if torch_dtype is None:
|
||||
torch_dtype = weight.dtype
|
||||
|
||||
if layer_class_name in conv_types:
|
||||
if not quant_conv:
|
||||
return layer
|
||||
if dtype_dict[weights_dtype]["num_bits"] < 4:
|
||||
weights_dtype = "uint4"
|
||||
is_conv_type = True
|
||||
reduction_axes = 1
|
||||
output_channel_size, channel_size = layer.weight.shape[:2]
|
||||
group_channel_size = channel_size // layer.groups
|
||||
use_quantized_matmul = False
|
||||
if use_quantized_matmul_conv:
|
||||
use_quantized_matmul = group_channel_size >= 32 and output_channel_size >= 32
|
||||
if use_quantized_matmul and not dtype_dict[weights_dtype]["is_integer"]:
|
||||
use_quantized_matmul = output_channel_size % 16 == 0 and group_channel_size % 16 == 0
|
||||
if use_quantized_matmul and dtype_dict[weights_dtype]["num_bits"] == 8:
|
||||
result_shape = layer.weight.shape
|
||||
layer.weight.data = layer.weight.flatten(1,-1)
|
||||
reduction_axes = -1
|
||||
elif layer_class_name in conv_transpose_types:
|
||||
if not quant_conv:
|
||||
return layer
|
||||
if dtype_dict[weights_dtype]["num_bits"] < 4:
|
||||
weights_dtype = "uint4"
|
||||
is_conv_transpose_type = True
|
||||
reduction_axes = 0
|
||||
channel_size, output_channel_size = layer.weight.shape[:2]
|
||||
use_quantized_matmul = False
|
||||
else:
|
||||
is_linear_type = True
|
||||
if layer_class_name in conv_types:
|
||||
if dtype_dict[weights_dtype]["num_bits"] < 4:
|
||||
weights_dtype = "uint4"
|
||||
is_conv_type = True
|
||||
reduction_axes = 1
|
||||
output_channel_size, channel_size = weight.shape[:2]
|
||||
if use_quantized_matmul:
|
||||
use_quantized_matmul = channel_size >= 32 and output_channel_size >= 32
|
||||
if use_quantized_matmul and not dtype_dict[weights_dtype]["is_integer"]:
|
||||
use_quantized_matmul = output_channel_size % 16 == 0 and channel_size % 16 == 0
|
||||
if use_quantized_matmul and dtype_dict[weights_dtype]["num_bits"] == 8:
|
||||
result_shape = weight.shape
|
||||
weight = weight.flatten(1,-1)
|
||||
reduction_axes = -1
|
||||
try:
|
||||
output_channel_size, channel_size = layer.weight.shape
|
||||
except Exception as e:
|
||||
raise ValueError(f"SDNQ: param_name={param_name} layer_class_name={layer_class_name} layer_weight_shape={layer.weight.shape} weights_dtype={weights_dtype} unsupported") from e
|
||||
elif layer_class_name in conv_transpose_types:
|
||||
if dtype_dict[weights_dtype]["num_bits"] < 4:
|
||||
weights_dtype = "uint4"
|
||||
is_conv_transpose_type = True
|
||||
reduction_axes = 0
|
||||
channel_size, output_channel_size = weight.shape[:2]
|
||||
use_quantized_matmul = False
|
||||
elif layer_class_name in linear_types:
|
||||
is_linear_type = True
|
||||
reduction_axes = -1
|
||||
try:
|
||||
output_channel_size, channel_size = weight.shape
|
||||
except Exception as e:
|
||||
raise ValueError(f"SDNQ: param_name={param_name} layer_class_name={layer_class_name} weight_shape={weight.shape} weights_dtype={weights_dtype} unsupported") from e
|
||||
if use_quantized_matmul:
|
||||
use_quantized_matmul = channel_size >= 32 and output_channel_size >= 32
|
||||
if use_quantized_matmul:
|
||||
use_quantized_matmul = channel_size >= 32 and output_channel_size >= 32
|
||||
if use_quantized_matmul:
|
||||
if dtype_dict[weights_dtype]["is_integer"]:
|
||||
use_quantized_matmul = output_channel_size % 8 == 0 and channel_size % 8 == 0
|
||||
else:
|
||||
use_quantized_matmul = output_channel_size % 16 == 0 and channel_size % 16 == 0
|
||||
|
||||
layer.weight.requires_grad = False
|
||||
if return_device is None:
|
||||
return_device = layer.weight.device
|
||||
if quantization_device is not None:
|
||||
layer.weight.data = layer.weight.to(quantization_device, non_blocking=non_blocking)
|
||||
if layer.weight.dtype != torch.float32:
|
||||
layer.weight.data = layer.weight.to(dtype=torch.float32)
|
||||
|
||||
if use_svd:
|
||||
try:
|
||||
layer.weight.data, svd_up, svd_down = apply_svdquant(layer.weight, rank=svd_rank, niter=svd_steps)
|
||||
if use_quantized_matmul:
|
||||
svd_up = svd_up.t_()
|
||||
svd_down = svd_down.t_()
|
||||
svd_up, svd_down = prepare_svd_for_matmul(svd_up, svd_down, use_quantized_matmul)
|
||||
except Exception:
|
||||
svd_up, svd_down = None, None
|
||||
if dtype_dict[weights_dtype]["is_integer"]:
|
||||
use_quantized_matmul = output_channel_size % 8 == 0 and channel_size % 8 == 0
|
||||
else:
|
||||
use_quantized_matmul = output_channel_size % 16 == 0 and channel_size % 16 == 0
|
||||
else:
|
||||
if weight.ndim > 1:
|
||||
output_channel_size, channel_size = weight.shape[-2:]
|
||||
else:
|
||||
output_channel_size, channel_size = 1, weight.shape[-1]
|
||||
reduction_axes = -1
|
||||
use_quantized_matmul = False
|
||||
|
||||
if use_svd:
|
||||
try:
|
||||
weight, svd_up, svd_down = apply_svdquant(weight, rank=svd_rank, niter=svd_steps)
|
||||
if use_quantized_matmul:
|
||||
svd_up = svd_up.t_()
|
||||
svd_down = svd_down.t_()
|
||||
svd_up, svd_down = prepare_svd_for_matmul(svd_up, svd_down, use_quantized_matmul)
|
||||
except Exception:
|
||||
svd_up, svd_down = None, None
|
||||
else:
|
||||
svd_up, svd_down = None, None
|
||||
|
||||
if group_size == 0:
|
||||
if use_quantized_matmul and dtype_dict[weights_dtype]["num_bits"] >= 6:
|
||||
group_size = -1
|
||||
elif is_linear_type:
|
||||
group_size = 2 ** ((2 if svd_up is None else 3) + dtype_dict[weights_dtype]["num_bits"])
|
||||
if group_size == 0:
|
||||
if use_quantized_matmul and dtype_dict[weights_dtype]["num_bits"] >= 6:
|
||||
group_size = -1
|
||||
elif is_linear_type:
|
||||
group_size = 2 ** ((2 if svd_up is None else 3) + dtype_dict[weights_dtype]["num_bits"])
|
||||
else:
|
||||
group_size = 2 ** ((1 if svd_up is None else 2) + dtype_dict[weights_dtype]["num_bits"])
|
||||
elif use_quantized_matmul and dtype_dict[weights_dtype]["num_bits"] == 8:
|
||||
group_size = -1 # override user value, re-quantizing 8bit into 8bit is pointless
|
||||
elif group_size != -1 and not is_linear_type:
|
||||
group_size = max(group_size // 2, 1)
|
||||
|
||||
if group_size > 0:
|
||||
if group_size >= channel_size:
|
||||
group_size = channel_size
|
||||
num_of_groups = 1
|
||||
else:
|
||||
num_of_groups = channel_size // group_size
|
||||
while num_of_groups * group_size != channel_size: # find something divisible
|
||||
num_of_groups -= 1
|
||||
if num_of_groups <= 1:
|
||||
group_size = channel_size
|
||||
num_of_groups = 1
|
||||
break
|
||||
group_size = channel_size // num_of_groups
|
||||
group_size = int(group_size)
|
||||
num_of_groups = int(num_of_groups)
|
||||
|
||||
if num_of_groups > 1:
|
||||
if result_shape is None:
|
||||
result_shape = weight.shape
|
||||
new_shape = list(result_shape)
|
||||
if is_conv_type:
|
||||
# output_channel_size, channel_size, X, X
|
||||
# output_channel_size, num_of_groups, group_size, X, X
|
||||
new_shape[1] = group_size
|
||||
new_shape.insert(1, num_of_groups)
|
||||
reduction_axes = 2
|
||||
elif is_conv_transpose_type:
|
||||
#channel_size, output_channel_size, X, X
|
||||
#num_of_groups, group_size, output_channel_size, X, X
|
||||
new_shape[0] = group_size
|
||||
new_shape.insert(0, num_of_groups)
|
||||
reduction_axes = 1
|
||||
else:
|
||||
group_size = 2 ** ((1 if svd_up is None else 2) + dtype_dict[weights_dtype]["num_bits"])
|
||||
elif use_quantized_matmul and dtype_dict[weights_dtype]["num_bits"] == 8:
|
||||
group_size = -1 # override user value, re-quantizing 8bit into 8bit is pointless
|
||||
elif group_size != -1 and not is_linear_type:
|
||||
group_size = max(group_size // 2, 1)
|
||||
# output_channel_size, channel_size
|
||||
# output_channel_size, num_of_groups, group_size
|
||||
last_dim_index = weight.ndim
|
||||
new_shape[last_dim_index - 1 : last_dim_index] = (num_of_groups, group_size)
|
||||
weight = weight.reshape(new_shape)
|
||||
else:
|
||||
group_size = -1
|
||||
|
||||
if group_size > 0:
|
||||
if group_size >= channel_size:
|
||||
group_size = channel_size
|
||||
num_of_groups = 1
|
||||
else:
|
||||
num_of_groups = channel_size // group_size
|
||||
while num_of_groups * group_size != channel_size: # find something divisible
|
||||
num_of_groups -= 1
|
||||
if num_of_groups <= 1:
|
||||
group_size = channel_size
|
||||
num_of_groups = 1
|
||||
break
|
||||
group_size = channel_size // num_of_groups
|
||||
group_size = int(group_size)
|
||||
num_of_groups = int(num_of_groups)
|
||||
|
||||
if num_of_groups > 1:
|
||||
if result_shape is None:
|
||||
result_shape = layer.weight.shape
|
||||
new_shape = list(result_shape)
|
||||
if is_conv_type:
|
||||
# output_channel_size, channel_size, X, X
|
||||
# output_channel_size, num_of_groups, group_size, X, X
|
||||
new_shape[1] = group_size
|
||||
new_shape.insert(1, num_of_groups)
|
||||
reduction_axes = 2
|
||||
elif is_conv_transpose_type:
|
||||
#channel_size, output_channel_size, X, X
|
||||
#num_of_groups, group_size, output_channel_size, X, X
|
||||
new_shape[0] = group_size
|
||||
new_shape.insert(0, num_of_groups)
|
||||
reduction_axes = 1
|
||||
elif is_linear_type:
|
||||
# output_channel_size, channel_size
|
||||
# output_channel_size, num_of_groups, group_size
|
||||
last_dim_index = layer.weight.ndim
|
||||
new_shape[last_dim_index - 1 : last_dim_index] = (num_of_groups, group_size)
|
||||
layer.weight.data = layer.weight.reshape(new_shape)
|
||||
|
||||
layer.weight.data, scale, zero_point = quantize_weight(layer.weight, reduction_axes, weights_dtype)
|
||||
if not dequantize_fp32 and not (use_quantized_matmul and not dtype_dict[weights_dtype]["is_integer"] and not use_tensorwise_fp8_matmul):
|
||||
scale = scale.to(dtype=torch_dtype)
|
||||
if zero_point is not None:
|
||||
zero_point = zero_point.to(dtype=torch_dtype)
|
||||
if svd_up is not None:
|
||||
svd_up = svd_up.to(dtype=torch_dtype)
|
||||
svd_down = svd_down.to(dtype=torch_dtype)
|
||||
|
||||
re_quantize_for_matmul = (num_of_groups > 1 or zero_point is not None)
|
||||
if use_quantized_matmul and not re_quantize_for_matmul:
|
||||
scale.t_()
|
||||
layer.weight.t_()
|
||||
layer.weight.data = prepare_weight_for_matmul(layer.weight)
|
||||
if not use_tensorwise_fp8_matmul and not dtype_dict[weights_dtype]["is_integer"]:
|
||||
scale = scale.to(dtype=torch.float32)
|
||||
|
||||
scale = scale.to(return_device, non_blocking=non_blocking)
|
||||
layer.scale = torch.nn.Parameter(scale, requires_grad=False)
|
||||
weight, scale, zero_point = quantize_weight(weight, reduction_axes, weights_dtype)
|
||||
if not dequantize_fp32 and not (use_quantized_matmul and not dtype_dict[weights_dtype]["is_integer"] and not use_tensorwise_fp8_matmul):
|
||||
scale = scale.to(dtype=torch_dtype)
|
||||
if zero_point is not None:
|
||||
zero_point = zero_point.to(return_device, non_blocking=non_blocking)
|
||||
layer.zero_point = torch.nn.Parameter(zero_point, requires_grad=False)
|
||||
else:
|
||||
layer.zero_point = None
|
||||
zero_point = zero_point.to(dtype=torch_dtype)
|
||||
if svd_up is not None:
|
||||
svd_up = svd_up.to(return_device, non_blocking=non_blocking)
|
||||
svd_down = svd_down.to(return_device, non_blocking=non_blocking)
|
||||
layer.svd_up = torch.nn.Parameter(svd_up, requires_grad=False)
|
||||
layer.svd_down = torch.nn.Parameter(svd_down, requires_grad=False)
|
||||
svd_up = svd_up.to(dtype=torch_dtype)
|
||||
svd_down = svd_down.to(dtype=torch_dtype)
|
||||
|
||||
re_quantize_for_matmul = (num_of_groups > 1 or zero_point is not None)
|
||||
if use_quantized_matmul and not re_quantize_for_matmul:
|
||||
scale.t_()
|
||||
weight.t_()
|
||||
weight = prepare_weight_for_matmul(weight)
|
||||
if not use_tensorwise_fp8_matmul and not dtype_dict[weights_dtype]["is_integer"]:
|
||||
scale = scale.to(dtype=torch.float32)
|
||||
|
||||
sdnq_dequantizer = SDNQDequantizer(
|
||||
result_dtype=torch_dtype,
|
||||
result_shape=result_shape,
|
||||
original_shape=original_shape,
|
||||
original_stride=original_stride,
|
||||
quantized_weight_shape=weight.shape,
|
||||
weights_dtype=weights_dtype,
|
||||
group_size=group_size,
|
||||
svd_rank=svd_rank,
|
||||
svd_steps=svd_steps,
|
||||
use_quantized_matmul=use_quantized_matmul,
|
||||
re_quantize_for_matmul=re_quantize_for_matmul,
|
||||
use_stochastic_rounding=use_stochastic_rounding,
|
||||
layer_class_name=layer_class_name,
|
||||
)
|
||||
|
||||
if dtype_dict[weights_dtype]["is_packed"]:
|
||||
if dtype_dict[weights_dtype]["is_unsigned"]:
|
||||
weight = pack_int_asymetric(weight, weights_dtype)
|
||||
else:
|
||||
layer.svd_up, layer.svd_down = None, None
|
||||
weight = pack_int_symetric(weight, weights_dtype)
|
||||
else:
|
||||
weight = weight.to(dtype=dtype_dict[weights_dtype]["torch_dtype"])
|
||||
|
||||
layer.sdnq_dequantizer = dequantizer_dict[weights_dtype](
|
||||
result_dtype=torch_dtype,
|
||||
result_shape=result_shape,
|
||||
original_shape=original_shape,
|
||||
quantized_weight_shape=layer.weight.shape,
|
||||
weights_dtype=weights_dtype,
|
||||
group_size=group_size,
|
||||
svd_rank=svd_rank,
|
||||
use_quantized_matmul=use_quantized_matmul,
|
||||
re_quantize_for_matmul=re_quantize_for_matmul,
|
||||
)
|
||||
layer.weight.data = layer.sdnq_dequantizer.pack_weight(layer.weight).to(return_device, non_blocking=non_blocking)
|
||||
return weight, scale, zero_point, svd_up, svd_down, sdnq_dequantizer, use_quantized_matmul
|
||||
|
||||
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.inference_context()
|
||||
def sdnq_quantize_layer(layer, weights_dtype="int8", torch_dtype=None, group_size=0, svd_rank=32, svd_steps=8, use_svd=False, quant_conv=False, use_quantized_matmul=False, use_quantized_matmul_conv=False, use_stochastic_rounding=False, dequantize_fp32=False, non_blocking=False, quantization_device=None, return_device=None, param_name=None): # pylint: disable=unused-argument
|
||||
layer_class_name = layer.__class__.__name__
|
||||
if layer_class_name in conv_transpose_types or layer_class_name in conv_types:
|
||||
if not quant_conv:
|
||||
return layer
|
||||
use_quantized_matmul = use_quantized_matmul_conv
|
||||
|
||||
layer.weight.requires_grad = False
|
||||
if return_device is None:
|
||||
return_device = layer.weight.device
|
||||
if quantization_device is None:
|
||||
quantization_device = layer.weight.device
|
||||
layer.weight.data = layer.weight.to(quantization_device, non_blocking=non_blocking)
|
||||
|
||||
(
|
||||
layer.weight.data,
|
||||
layer.scale, layer.zero_point,
|
||||
layer.svd_up, layer.svd_down,
|
||||
layer.sdnq_dequantizer,
|
||||
use_quantized_matmul,
|
||||
) = sdnq_quantize_layer_weight(
|
||||
layer.weight,
|
||||
layer_class_name=layer_class_name,
|
||||
weights_dtype=weights_dtype,
|
||||
torch_dtype=torch_dtype,
|
||||
group_size=group_size,
|
||||
svd_rank=svd_rank,
|
||||
svd_steps=svd_steps,
|
||||
use_svd=use_svd,
|
||||
use_quantized_matmul=use_quantized_matmul,
|
||||
use_stochastic_rounding=use_stochastic_rounding,
|
||||
dequantize_fp32=dequantize_fp32,
|
||||
param_name=param_name,
|
||||
)
|
||||
|
||||
layer.scale = torch.nn.Parameter(layer.scale.to(return_device, non_blocking=non_blocking), requires_grad=False)
|
||||
if layer.zero_point is not None:
|
||||
layer.zero_point = torch.nn.Parameter(layer.zero_point.to(return_device, non_blocking=non_blocking), requires_grad=False)
|
||||
if layer.svd_up is not None:
|
||||
layer.svd_up = torch.nn.Parameter(layer.svd_up.to(return_device, non_blocking=non_blocking), requires_grad=False)
|
||||
layer.svd_down = torch.nn.Parameter(layer.svd_down.to(return_device, non_blocking=non_blocking), requires_grad=False)
|
||||
|
||||
layer = layer.to(return_device, non_blocking=non_blocking)
|
||||
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__)
|
||||
return layer
|
||||
|
||||
|
||||
def apply_sdnq_to_module(model, weights_dtype="int8", torch_dtype=None, group_size=0, svd_rank=32, svd_steps=8, use_svd=False, quant_conv=False, use_quantized_matmul=False, use_quantized_matmul_conv=False, dequantize_fp32=False, non_blocking=False, quantization_device=None, return_device=None, modules_to_not_convert: List[str] = None, modules_dtype_dict: Dict[str, List[str]] = None, full_param_name="", op=None): # pylint: disable=unused-argument
|
||||
@devices.inference_context()
|
||||
def apply_sdnq_to_module(model, weights_dtype="int8", torch_dtype=None, group_size=0, svd_rank=32, svd_steps=8, use_svd=False, quant_conv=False, use_quantized_matmul=False, use_quantized_matmul_conv=False, use_stochastic_rounding=False, dequantize_fp32=False, non_blocking=False, quantization_device=None, return_device=None, modules_to_not_convert: List[str] = None, modules_dtype_dict: Dict[str, List[str]] = None, full_param_name="", op=None): # pylint: disable=unused-argument
|
||||
has_children = list(model.children())
|
||||
if not has_children:
|
||||
return model
|
||||
@@ -346,20 +409,20 @@ def apply_sdnq_to_module(model, weights_dtype="int8", torch_dtype=None, group_si
|
||||
modules_to_not_convert = []
|
||||
if modules_dtype_dict is None:
|
||||
modules_dtype_dict = {}
|
||||
for param_name, module in model.named_children():
|
||||
if param_name == "sdnq_dequantizer":
|
||||
continue
|
||||
for module_name, module in model.named_children():
|
||||
if full_param_name:
|
||||
param_name = full_param_name + "." + param_name
|
||||
param_name = full_param_name + "." + module_name
|
||||
else:
|
||||
param_name = module_name
|
||||
if hasattr(module, "weight") and module.weight is not None:
|
||||
param_name = param_name + ".weight"
|
||||
if check_param_name_in(param_name, modules_to_not_convert):
|
||||
continue
|
||||
layer_class_name = module.__class__.__name__
|
||||
if layer_class_name in allowed_types:
|
||||
if layer_class_name in allowed_types and module.weight.dtype in {torch.float32, torch.float16, torch.bfloat16}:
|
||||
if (layer_class_name in conv_types or layer_class_name in conv_transpose_types) and not quant_conv:
|
||||
continue
|
||||
module = sdnq_quantize_layer(
|
||||
setattr(model, module_name, sdnq_quantize_layer(
|
||||
module,
|
||||
weights_dtype=get_minimum_dtype(weights_dtype, param_name, modules_dtype_dict),
|
||||
torch_dtype=torch_dtype,
|
||||
@@ -370,13 +433,14 @@ def apply_sdnq_to_module(model, weights_dtype="int8", torch_dtype=None, group_si
|
||||
quant_conv=quant_conv,
|
||||
use_quantized_matmul=use_quantized_matmul,
|
||||
use_quantized_matmul_conv=use_quantized_matmul_conv,
|
||||
use_stochastic_rounding=use_stochastic_rounding,
|
||||
dequantize_fp32=dequantize_fp32,
|
||||
non_blocking=non_blocking,
|
||||
quantization_device=quantization_device,
|
||||
return_device=return_device,
|
||||
param_name=param_name,
|
||||
)
|
||||
module = apply_sdnq_to_module(
|
||||
))
|
||||
setattr(model, module_name, apply_sdnq_to_module(
|
||||
module,
|
||||
weights_dtype=weights_dtype,
|
||||
torch_dtype=torch_dtype,
|
||||
@@ -387,6 +451,7 @@ def apply_sdnq_to_module(model, weights_dtype="int8", torch_dtype=None, group_si
|
||||
quant_conv=quant_conv,
|
||||
use_quantized_matmul=use_quantized_matmul,
|
||||
use_quantized_matmul_conv=use_quantized_matmul_conv,
|
||||
use_stochastic_rounding=use_stochastic_rounding,
|
||||
dequantize_fp32=dequantize_fp32,
|
||||
non_blocking=non_blocking,
|
||||
quantization_device=quantization_device,
|
||||
@@ -395,10 +460,11 @@ def apply_sdnq_to_module(model, weights_dtype="int8", torch_dtype=None, group_si
|
||||
modules_dtype_dict=modules_dtype_dict,
|
||||
full_param_name=param_name,
|
||||
op=op,
|
||||
)
|
||||
))
|
||||
return model
|
||||
|
||||
|
||||
@devices.inference_context()
|
||||
def sdnq_post_load_quant(
|
||||
model,
|
||||
weights_dtype="int8",
|
||||
@@ -410,6 +476,7 @@ def sdnq_post_load_quant(
|
||||
quant_conv: bool = False,
|
||||
use_quantized_matmul: bool = False,
|
||||
use_quantized_matmul_conv: bool = False,
|
||||
use_stochastic_rounding: bool = False,
|
||||
dequantize_fp32: bool = False,
|
||||
non_blocking: bool = False,
|
||||
add_skip_keys:bool = True,
|
||||
@@ -441,6 +508,7 @@ def sdnq_post_load_quant(
|
||||
quant_conv=quant_conv,
|
||||
use_quantized_matmul=use_quantized_matmul,
|
||||
use_quantized_matmul_conv=use_quantized_matmul_conv,
|
||||
use_stochastic_rounding=use_stochastic_rounding,
|
||||
dequantize_fp32=dequantize_fp32,
|
||||
non_blocking=non_blocking,
|
||||
quantization_device=quantization_device,
|
||||
@@ -595,6 +663,7 @@ class SDNQQuantizer(DiffusersQuantizer, HfQuantizer):
|
||||
quant_conv=self.quantization_config.quant_conv,
|
||||
use_quantized_matmul=self.quantization_config.use_quantized_matmul,
|
||||
use_quantized_matmul_conv=self.quantization_config.use_quantized_matmul_conv,
|
||||
use_stochastic_rounding=self.quantization_config.use_stochastic_rounding,
|
||||
dequantize_fp32=self.quantization_config.dequantize_fp32,
|
||||
non_blocking=self.quantization_config.non_blocking,
|
||||
quantization_device=None,
|
||||
@@ -633,6 +702,7 @@ class SDNQQuantizer(DiffusersQuantizer, HfQuantizer):
|
||||
quantization_config_dict.pop("return_device", None)
|
||||
quantization_config_dict.pop("non_blocking", None)
|
||||
quantization_config_dict.pop("add_skip_keys", None)
|
||||
quantization_config_dict.pop("use_stochastic_rounding", None)
|
||||
with init_empty_weights():
|
||||
model = sdnq_post_load_quant(model, add_skip_keys=False, **quantization_config_dict)
|
||||
|
||||
@@ -754,6 +824,7 @@ class SDNQConfig(QuantizationConfigMixin):
|
||||
quant_conv: bool = False,
|
||||
use_quantized_matmul: bool = False,
|
||||
use_quantized_matmul_conv: bool = False,
|
||||
use_stochastic_rounding: bool = False,
|
||||
dequantize_fp32: bool = False,
|
||||
non_blocking: bool = False,
|
||||
add_skip_keys: bool = True,
|
||||
@@ -772,6 +843,7 @@ class SDNQConfig(QuantizationConfigMixin):
|
||||
self.quant_conv = quant_conv
|
||||
self.use_quantized_matmul = use_quantized_matmul
|
||||
self.use_quantized_matmul_conv = use_quantized_matmul_conv
|
||||
self.use_stochastic_rounding = use_stochastic_rounding
|
||||
self.dequantize_fp32 = dequantize_fp32
|
||||
self.non_blocking = non_blocking
|
||||
self.add_skip_keys = add_skip_keys
|
||||
@@ -831,3 +903,5 @@ diffusers.quantizers.auto.AUTO_QUANTIZATION_CONFIG_MAPPING["sdnq"] = SDNQConfig
|
||||
import transformers.quantizers.auto # noqa: E402,RUF100 # pylint: disable=wrong-import-order
|
||||
transformers.quantizers.auto.AUTO_QUANTIZER_MAPPING["sdnq"] = SDNQQuantizer
|
||||
transformers.quantizers.auto.AUTO_QUANTIZATION_CONFIG_MAPPING["sdnq"] = SDNQConfig
|
||||
|
||||
sdnq_quantize_layer_weight_compiled = compile_func(sdnq_quantize_layer_weight)
|
||||
|
||||
Reference in New Issue
Block a user