From 35b62a908f4c290840c25d738c8b22709862e129 Mon Sep 17 00:00:00 2001 From: dxqb <183307934+dxqb@users.noreply.github.com> Date: Sat, 1 Aug 2026 00:13:26 +0200 Subject: [PATCH] Chunk weight quantization to bound fp32 VRAM spikes Quantizing a W8A8 weight upcast the whole tensor to fp32, and the non-in-place `mul` made a second full-size copy, so one large layer transiently needed several times its own size in VRAM. On small GPUs that was enough to OOM during quantization. `quantize_int8_tensorwise_chunked` / `quantize_fp8_tensorwise_chunked` quantize the weight in row-blocks of `_QUANTIZE_CHUNK_ELEMENTS`, bounding the fp32 transient to one block; `LinearW8A8.quantize` uses them. The tensorwise scale now comes from `torch.aminmax` instead of `abs().max()`, dropping another full-tensor copy. The shared `quantize_int8` / `quantize_fp8` cores take an explicit fp32 copy and then scale in place, so an already-fp32 caller tensor is never clobbered. Output is unchanged, bit for bit. Co-Authored-By: Claude Opus 5 (1M context) --- modules/module/quantized/LinearW8A8.py | 8 ++-- modules/util/quantization_util.py | 59 ++++++++++++++++++-------- 2 files changed, 45 insertions(+), 22 deletions(-) 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)