From 26f947300067f8f116a9668b1f268a86d62faeb6 Mon Sep 17 00:00:00 2001 From: Alexander Eichhorn Date: Wed, 14 Oct 2026 00:16:32 +0200 Subject: [PATCH] perf(quantization): move int8 convrot weights to the device before widening A CPU-resident Int8ConvrotLinear streamed its weight with one .to(device=, dtype=), which converts on the host and transfers the widened tensor; partial loads of int8 convrot models now copy int8 and cast on the device, bit-identical output. --- invokeai/backend/quantization/int8_convrot.py | 10 ++-- .../backend/quantization/test_int8_convrot.py | 60 +++++++++++++++++++ 2 files changed, 66 insertions(+), 4 deletions(-) diff --git a/invokeai/backend/quantization/int8_convrot.py b/invokeai/backend/quantization/int8_convrot.py index f6d986a1319..226e89dc6fc 100644 --- a/invokeai/backend/quantization/int8_convrot.py +++ b/invokeai/backend/quantization/int8_convrot.py @@ -266,7 +266,9 @@ def __init__( def _dequantized_weight(self, device: torch.device, dtype: torch.dtype) -> torch.Tensor: # `.to(device)` is a no-op in the intended fully-resident regime; under partial load # the model cache may leave these buffers on the CPU, in which case this call streams - # the int8 weight to the compute device per forward (see the class docstring). + # the int8 weight to the compute device per forward (see the class docstring). The move + # and the cast stay two steps: a single `.to(device=, dtype=)` from the CPU converts on + # the CPU first and then copies the widened weight (~5x slower per streamed bf16 layer on an RTX 4090). # # Dequant + derotation run directly in the compute dtype: int8 values are exact in # bf16/fp16, the scale multiply adds ~0.2% relative rounding vs the ~0.4-0.8% int8 @@ -275,7 +277,7 @@ def _dequantized_weight(self, device: torch.device, dtype: torch.dtype) -> torch # fc1 in bf16) instead of tripling through an fp32 intermediate (~1.5 GiB). An fp32 # compute dtype still gets the exact fp32 path for free. scale = as_column_scale(self.weight, self.weight_scale).to(device=device, dtype=dtype) - w = self.weight.to(device=device, dtype=dtype) * scale + w = self.weight.to(device=device).to(dtype=dtype) * scale if self.convrot: w = ( w.view(self.out_features, self.in_features // self.group_size, self.group_size) @@ -290,8 +292,8 @@ def dequant_transient_bytes(self, compute_dtype: torch.dtype) -> int: alongside the product it is multiplied into, and that product is then alive alongside the derotation matmul's output. Each pair is freed before the next allocates. - Under partial load the int8 weight is additionally streamed to the device per call, adding up to - half a weight again; the resident-model regime this scheme targets pays nothing for that. Verified + Under partial load the int8 weight is additionally streamed to the device per call. That copy is at + most half a weight and is freed by the cast, before the product allocates, so the peak stays at two. Verified against the real `z_image_turbo_int8_convrot` checkpoint: exact for most shapes, and 2 MiB under on the two that make cuBLAS take a workspace (10240x3840 and 3840x10240), against a 3 GiB reservation floor. diff --git a/tests/backend/quantization/test_int8_convrot.py b/tests/backend/quantization/test_int8_convrot.py index 421d5934fa2..0f23d630a01 100644 --- a/tests/backend/quantization/test_int8_convrot.py +++ b/tests/backend/quantization/test_int8_convrot.py @@ -10,6 +10,8 @@ import pytest import torch +from torch.utils._python_dispatch import TorchDispatchMode +from torch.utils._pytree import tree_flatten from invokeai.backend.model_manager.taxonomy import ModelFormat from invokeai.backend.quantization.dequantizing_linear import peak_dequant_transient_bytes, requires_sidecar_patching @@ -101,6 +103,64 @@ def test_int8_convrot_linear_state_dict_contract() -> None: assert lin.state_dict()["weight"].dtype == torch.int8 +@pytest.mark.slow +@pytest.mark.skipif(not torch.cuda.is_available(), reason="measures CUDA allocations") +def test_a_layer_left_in_ram_streams_its_int8_codes_within_the_transient_estimate() -> None: + """Partial loading leaves layers in RAM and their forward streams the int8 codes to the compute device. The + streamed copy must decode to exactly what the resident layer computes, and must be freed before the decode + peaks: the node's working-memory reservation is sized from `dequant_transient_bytes`, which has no room for it. + + Measured in requested bytes, as in the nvfp4 counterpart. The resident decode runs first, which also mints the + shared Hadamard and the cuBLAS workspace a process's first matmul allocates; neither is part of the estimate.""" + device, dtype = torch.device("cuda"), torch.bfloat16 + torch.manual_seed(5) + q, scale = _quantize_reference(torch.randn(1024, 4 * CONVROT_GROUP_SIZE), convrot=True) + lin = Int8ConvrotLinear(q, scale, convrot=True) + resident = Int8ConvrotLinear(q.to(device), scale.to(device), convrot=True)._dequantized_weight(device, dtype) + torch.cuda.synchronize() + torch.cuda.reset_peak_memory_stats() + baseline = torch.cuda.memory_stats()["requested_bytes.all.current"] + + streamed = lin._dequantized_weight(device, dtype) + torch.cuda.synchronize() + + measured = torch.cuda.memory_stats()["requested_bytes.all.peak"] - baseline + assert torch.equal(streamed, resident) + if measured == 0: + pytest.skip("this allocator backend (e.g. cudaMallocAsync) does not track requested bytes") + # The estimate leaves out the per-row scale in the compute dtype, a few KB beside a weight-sized tensor. + assert measured <= lin.dequant_transient_bytes(dtype) + lin.out_features * dtype.itemsize + + +class _HostToDeviceCopies(TorchDispatchMode): + """Records every tensor an op produced off the CPU from CPU inputs, i.e. what crossed to the device.""" + + def __init__(self) -> None: + super().__init__() + self.copies: list[torch.Tensor] = [] + + def __torch_dispatch__(self, func, types, args=(), kwargs=None): + out = func(*args, **(kwargs or {})) + inputs, _ = tree_flatten((args, kwargs)) + if any(isinstance(t, torch.Tensor) and t.device.type == "cpu" for t in inputs): + outputs, _ = tree_flatten(out) + self.copies += [t for t in outputs if isinstance(t, torch.Tensor) and t.device.type != "cpu"] + return out + + +def test_a_layer_left_in_ram_moves_its_int8_codes_before_widening_them() -> None: + """A single `.to(device=, dtype=)` from the CPU widens on the host and then copies twice the bytes: ~5x slower + per streamed layer under partial loading (#9773). Pinned through what crosses rather than through timing; the + meta device stands in for the accelerator, so this runs without one.""" + q, scale = _quantize_reference(torch.randn(64, CONVROT_GROUP_SIZE), convrot=True) + lin = Int8ConvrotLinear(q, scale, convrot=True) + + with _HostToDeviceCopies() as moved: + lin._dequantized_weight(torch.device("meta"), torch.bfloat16) + + assert [t.dtype for t in moved.copies if t.shape == q.shape] == [torch.int8] + + class TestTheDerotationMatrixIsSharedNotPerLayer: """It is byte-identical across every layer of a given group size, and a model has hundreds of them: Krea-2's transformer is 264 Linears, so a per-module 256x256 fp32 buffer is 69 MB of