perf(quantization): stream int8 convrot weights as int8 under partial loading

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.

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