diff --git a/modules/module/quantized/LinearW8A8.py b/modules/module/quantized/LinearW8A8.py index babbd1a80..7a9170ac9 100644 --- a/modules/module/quantized/LinearW8A8.py +++ b/modules/module/quantized/LinearW8A8.py @@ -5,9 +5,9 @@ from modules.util.quantization_util import ( dequantize, quantize_fp8_axiswise, - quantize_fp8_tensorwise, + quantize_fp8_tensorwise_chunked, quantize_int8_axiswise, - quantize_int8_tensorwise, + quantize_int8_tensorwise_chunked, ) import torch @@ -107,9 +107,9 @@ def quantize(self, device: torch.device | None = None): if device is not None: weight = weight.to(device=device) if self._dtype == torch.int8: - weight, scale = quantize_int8_tensorwise(weight) + weight, scale = quantize_int8_tensorwise_chunked(weight) else: - weight, scale = quantize_fp8_tensorwise(weight) + weight, scale = quantize_fp8_tensorwise_chunked(weight) if device is not None: weight = weight.to(device=orig_device) diff --git a/modules/util/quantization_util.py b/modules/util/quantization_util.py index 3570f670f..e416100f7 100644 --- a/modules/util/quantization_util.py +++ b/modules/util/quantization_util.py @@ -25,19 +25,33 @@ LinearNf4 = None def quantize_int8(x: Tensor, scale: float | Tensor) -> Tensor: - q = x.float().mul(1.0 / scale).round_().clamp_(-128.0, 127.0).to(torch.int8) - return q + xf = x.to(torch.float32, copy=True) + return xf.mul_(1.0 / scale).round_().clamp_(-128.0, 127.0).to(torch.int8) -def quantize_int8_tensorwise_get_scale(x: Tensor) -> float: - abs_max = x.abs().max() - scale = (abs_max.float() / 127.0).clamp(min=1e-30) - return scale +def quantize_int8_tensorwise_get_scale(x: Tensor) -> Tensor: + # max|x| == max(max x, -min x): one pass over x, no full-tensor abs() copy + min_val, max_val = torch.aminmax(x) + abs_max = torch.maximum(max_val, min_val.neg()) + return (abs_max.float() / 127.0).clamp(min=1e-30) -def quantize_int8_tensorwise(x: Tensor) -> tuple[Tensor, float]: +def quantize_int8_tensorwise(x: Tensor) -> tuple[Tensor, Tensor]: scale = quantize_int8_tensorwise_get_scale(x) q = quantize_int8(x, scale) return q, scale +# Quantizing a whole weight at once allocates a full-size fp32 intermediary, which spikes VRAM on small +# GPUs. The chunked variants work in row-blocks, bounding that transient to one block. Load-time only, +# so the Python loop costs nothing. +_QUANTIZE_CHUNK_ELEMENTS = 16 * 1024 * 1024 + +def quantize_int8_tensorwise_chunked(x: Tensor) -> tuple[Tensor, Tensor]: + scale = quantize_int8_tensorwise_get_scale(x) + q = torch.empty_like(x, dtype=torch.int8) + rows = max(1, _QUANTIZE_CHUNK_ELEMENTS // x[0].numel()) + for i in range(0, x.shape[0], rows): + q[i:i + rows] = quantize_int8(x[i:i + rows], scale) + return q, scale + def quantize_int8_axiswise_get_scale(x: Tensor, dim: int) -> Tensor: abs_max = x.abs().amax(dim=dim, keepdim=True) scale = (abs_max.float() / 127.0).clamp(min=1e-30) @@ -49,24 +63,33 @@ def quantize_int8_axiswise(x: Tensor, dim: int) -> tuple[Tensor, Tensor]: return q, scale def quantize_fp8(x: Tensor, scale: float | Tensor) -> Tensor: - q = x.float().mul(1.0 / scale).clamp_(-448.0, 448.0).to(torch.float8_e4m3fn) - return q + xf = x.to(torch.float32, copy=True) + return xf.mul_(1.0 / scale).clamp_(-448.0, 448.0).to(torch.float8_e4m3fn) -def quantize_fp8_tensorwise_get_scale(x: Tensor) -> float: - abs_max = x.abs().max() - scale = (abs_max.float() / 448.0).clamp(min=1e-30) - return scale +def quantize_fp8_tensorwise_get_scale(x: Tensor) -> Tensor: + # max|x| == max(max x, -min x): one pass over x, no full-tensor abs() copy + min_val, max_val = torch.aminmax(x) + abs_max = torch.maximum(max_val, min_val.neg()) + return (abs_max.float() / 448.0).clamp(min=1e-30) + +def quantize_fp8_tensorwise(x: Tensor) -> tuple[Tensor, Tensor]: + scale = quantize_fp8_tensorwise_get_scale(x) + q = quantize_fp8(x, scale) + return q, scale + +def quantize_fp8_tensorwise_chunked(x: Tensor) -> tuple[Tensor, Tensor]: + scale = quantize_fp8_tensorwise_get_scale(x) + q = torch.empty_like(x, dtype=torch.float8_e4m3fn) + rows = max(1, _QUANTIZE_CHUNK_ELEMENTS // x[0].numel()) + for i in range(0, x.shape[0], rows): + q[i:i + rows] = quantize_fp8(x[i:i + rows], scale) + return q, scale def quantize_fp8_axiswise_get_scale(x: Tensor, dim: int) -> Tensor: abs_max = x.abs().amax(dim=dim, keepdim=True) scale = (abs_max.float() / 448.0).clamp(min=1e-30) return scale -def quantize_fp8_tensorwise(x: Tensor) -> tuple[Tensor, float]: - scale = quantize_fp8_tensorwise_get_scale(x) - q = quantize_fp8(x, scale) - return q, scale - def quantize_fp8_axiswise(x: Tensor, dim: int) -> tuple[Tensor, Tensor]: scale = quantize_fp8_axiswise_get_scale(x, dim) q = quantize_fp8(x, scale)