Fix QQQ workspace lifetime and reduce packing memory

gptqmodel/looper/qqq_processor.py

@@ -162,7 +162,10 @@ def process(
         # logger.info(f"Quantizing module START: {name}, {gptq[name].shape()}")
         ## Need to return the quantized_weight for offloading
         q = qqq[module.name]
-        wq, q_scales, q_zeros, q_g_idx, duration, avg_loss, damp_percent, q_scales_extra, nsamples = q.quantize()
+        try:
+            wq, q_scales, q_zeros, q_g_idx, duration, avg_loss, damp_percent, q_scales_extra, nsamples = q.quantize()
+        finally:
+            q.free()
 
         q_scales = q_scales.to(CPU)
         q_zeros = q_zeros.to(CPU)

gptqmodel/nn_modules/qlinear/qqq.py

@@ -287,13 +287,8 @@ def list_buffers(self) -> List:
             buf.append(self.reduce_buffer)
         return buf
 
-    #def pack(self, linear: nn.Module, scales: t.Tensor, zeros: t.Tensor, g_idx: t.Tensor = None):
     def pack(self, linear: torch.nn.Module, scales: torch.Tensor, s_extra=None):
-        """Pack a fake-quantized linear layer into this actual Marlin representation.
-        @linear: fake-quantized `torch.nn.Linear` layer to convert (must be of type `torch.float16`)
-        @scales: corresponding quantization scales of shape `(infeatures, groups)`
-        @s_extra: corresponding quantization scales of shape `(1, outfeatures)`
-        """
+        """Pack a fake-quantized linear layer into the Marlin representation."""
         if self.group_size != self.in_features:
             assert s_extra is not None, "s_extra is needed"
         if linear.weight.dtype != torch.float16:
@@ -302,87 +297,112 @@ def pack(self, linear: torch.nn.Module, scales: torch.Tensor, s_extra=None):
                 If you can ensure your GEMM results don't overflow torch.float16, it will still function correctly.
                 Otherwise, it will yield incorrect results."""
             )
-        s = scales.t()
-        w = linear.weight.data.t()
-        if self.group_size != self.in_features:
-            w = w.reshape((-1, self.group_size, self.out_features))
-            w = w.permute(1, 0, 2)
-            w = w.reshape((self.group_size, -1))
-            s = s.reshape((1, -1))
-        w = torch.round(w / s).int()
-        if self.group_size != self.in_features:
-            w += (self.maxq + 1) // 2
-            w = torch.clamp(w, 0, self.maxq)
-        else:
-            w = torch.clamp(w, -self.maxq, self.maxq)
+
+        raw_scales = scales.t()
         if self.group_size != self.in_features:
             s_extra = s_extra.reshape(1, -1).to(dtype=torch.float32)
-            s = (s.reshape(-1, self.out_features) / s_extra).to(dtype=torch.float16)
-
-            w = w.reshape((self.group_size, -1, self.out_features))
-            w = w.permute(1, 0, 2)
-            w = w.reshape((self.in_features, self.out_features)).contiguous()
-            s = s.reshape((-1, len(self._scale_perm)))[:, self._scale_perm]
-            s_extra = s_extra.reshape((-1, len(self._scale_perm_single)))[
-                      :, self._scale_perm_single
-                      ]
-            s_extra = s_extra.reshape((-1, self.out_features)).contiguous()
+            packed_s_group = (raw_scales / s_extra).to(dtype=torch.float16)
+            packed_s_group = packed_s_group.reshape(
+                (-1, len(self._scale_perm))
+            )[:, self._scale_perm].reshape((-1, self.out_features)).contiguous()
+            packed_s_channel = s_extra.reshape(
+                (-1, len(self._scale_perm_single))
+            )[:, self._scale_perm_single].reshape((-1, self.out_features)).contiguous()
         else:
-            # NOTE(zhangying): div 2 ** (8 - self.bits)) to deal with right_shift in unpacking
-            s = (
-                (s / (2 ** (8 - self.bits)))
+            packed_s_group = None
+            packed_s_channel = (
+                (raw_scales / (2 ** (8 - self.bits)))
                 .reshape((-1, len(self._scale_perm_single)))[:, self._scale_perm_single]
                 .to(dtype=torch.float32)
+                .reshape((-1, self.out_features))
+                .contiguous()
             )
-        s = s.reshape((-1, self.out_features)).contiguous()
-        w = w.reshape(
-            (
-                self.in_features // self.tile,
-                self.tile,
-                self.out_features // self.tile,
-                self.tile,
+
+        input_chunk_size = min(1024, self.in_features)
+        output_chunk_size = min(256, self.out_features)
+        input_chunk_size -= input_chunk_size % self.tile
+        output_chunk_size -= output_chunk_size % 64
+        if input_chunk_size == 0 or output_chunk_size == 0:
+            raise ValueError(
+                f"QQQ pack requires dimensions divisible by {self.tile} and 64, "
+                f"got in_features={self.in_features}, out_features={self.out_features}"
             )
+
+        packed_weight = torch.empty(
+            (self.in_features // self.tile, self.out_features * 2),
+            dtype=torch.int32,
+            device=CPU,
         )
-        w = w.permute((0, 2, 1, 3))
-        w = w.reshape((self.in_features // self.tile, self.out_features * self.tile))
-        res = w
-        res = res.reshape((-1, self._perm.numel()))[:, self._perm].reshape(res.shape)
-        q = np.zeros((res.shape[0], res.shape[1] // 8), dtype=np.uint32)
-        res = res.cpu().numpy().astype(np.uint32)
-        if self.group_size != self.in_features:
-            for i in range(8):
-                q |= res[:, i::8] << 4 * i
-        else:
-            for i in range(8):
-                q |= (res[:, i::8] & 0xF) << 4 * i
-        q = torch.from_numpy(q.astype(np.int32)).to(CPU)
+        weight = linear.weight.data
+        perm = self._perm.to(device=weight.device)
+        for input_start in range(0, self.in_features, input_chunk_size):
+            input_end = min(input_start + input_chunk_size, self.in_features)
+            input_size = input_end - input_start
+            if input_size % self.tile != 0:
+                raise ValueError("QQQ pack input chunk is not tile aligned")
+
+            input_indices = torch.arange(
+                input_start,
+                input_end,
+                device=weight.device,
+                dtype=torch.long,
+            )
+            if self.group_size != self.in_features:
+                group_indices = torch.div(input_indices, self.group_size, rounding_mode="floor")
+                scale_chunk = raw_scales.index_select(0, group_indices)
+            else:
+                scale_chunk = raw_scales.expand(input_size, -1)
+
+            for output_start in range(0, self.out_features, output_chunk_size):
+                output_end = min(output_start + output_chunk_size, self.out_features)
+                output_size = output_end - output_start
+                if output_size % 64 != 0:
+                    raise ValueError("QQQ pack output chunk is not permutation aligned")
+
+                weight_chunk = weight[output_start:output_end, input_start:input_end].transpose(0, 1)
+                scale_chunk_view = scale_chunk[:, output_start:output_end]
+                codes = torch.round(weight_chunk / scale_chunk_view).to(dtype=torch.int32)
+                if self.group_size != self.in_features:
+                    codes.add_((self.maxq + 1) // 2).clamp_(0, self.maxq)
+                else:
+                    codes.clamp_(-self.maxq, self.maxq)
+
+                transformed = codes.reshape(
+                    input_size // self.tile,
+                    self.tile,
+                    output_size // self.tile,
+                    self.tile,
+                ).permute(0, 2, 1, 3).reshape(
+                    input_size // self.tile,
+                    output_size * self.tile,
+                )
+                transformed = transformed.reshape(-1, perm.numel()).index_select(1, perm).reshape(
+                    input_size // self.tile,
+                    output_size * self.tile,
+                )
+                packed_chunk = torch.zeros(
+                    (input_size // self.tile, output_size * 2),
+                    dtype=torch.int32,
+                    device=weight.device,
+                )
+                for lane in range(8):
+                    packed_chunk.bitwise_or_(
+                        (transformed[:, lane::8] & 0xF) << (4 * lane)
+                    )
 
-        #self.B[:, :] = q.to(self.B.device)
-        self.register_buffer("B", q)
+                packed_weight[
+                    input_start // self.tile:input_end // self.tile,
+                    output_start * 2:output_end * 2,
+                ].copy_(packed_chunk.to(device=CPU))
 
+        self.register_buffer("B", packed_weight)
         if self.group_size != self.in_features:
-            #self.s_group[:, :] = s.to(self.s_group.device)
-            self.register_buffer("s_group", s.to(CPU))
-
-            #self.s_channel[:, :] = s_extra.to(self.s_channel.device)
-            self.register_buffer("s_channel", s_extra.to(CPU))
+            self.register_buffer("s_group", packed_s_group)
+            self.register_buffer("s_channel", packed_s_channel)
         else:
-            # self.s_group = torch.tensor(
-            #     [], dtype=torch.float16, device=self.s_channel.device
-            # )
-            # self.register_buffer("s_group", torch.tensor(
-            #     [], dtype=torch.float16, device=CPU
-            # ))
-
-            #self.s_channel[:, :] = s.to(self.s_channel.device)
-            self.register_buffer("s_channel", s.to(CPU))
-        if linear.bias is not None:
-            if self.bias is not None:
-                # self.bias[:] = linear.bias.data.to(self.bias.device).to(torch.float16)
-                self.register_buffer("bias", linear.bias.data.to(self.bias.device).to(torch.float16))
-            # else:
-            #    # self.bias = linear.bias.clone().to(torch.float16)
-            #    self.register_buffer("bias", linear.bias.clone().to(torch.float16))
+            self.register_buffer("s_channel", packed_s_channel)
+        if linear.bias is not None and self.bias is not None:
+            self.register_buffer("bias", linear.bias.data.to(self.bias.device).to(torch.float16))
 
     # activation int8 quantization
     def dynamic_quant(self, x: torch.Tensor):

tests/test_qqq_jit.py

@@ -75,3 +75,110 @@ def test_qqq_forward_raises_runtime_error_when_jit_ops_missing(monkeypatch):
 
     with pytest.raises(ModuleNotFoundError, match="missing qqq jit ops"):
         module(torch.randn((1, module.in_features), dtype=torch.float16))
+
+
+def _build_grouped_parity_modules(
+    cases: list[tuple[int, float, int]],
+) -> tuple[qqq_module.QQQLinear, qqq_module.QQQTorchLinear]:
+    in_features = 256
+    out_features = 128
+    group_size = 128
+    linear = torch.nn.Linear(in_features, out_features, bias=False, dtype=torch.float16)
+    linear.weight.data.zero_()
+    scales = torch.ones((out_features, in_features // group_size), dtype=torch.float16)
+
+    for output_index, (code, scale, _) in enumerate(cases):
+        scales[output_index, 0] = scale
+        linear.weight.data[output_index, 0] = (code - 8) * scale
+
+    s_channel = torch.ones(out_features, dtype=torch.float32)
+    modules = []
+    for module_cls in (qqq_module.QQQLinear, qqq_module.QQQTorchLinear):
+        module = module_cls(
+            bits=4,
+            group_size=group_size,
+            sym=True,
+            desc_act=False,
+            in_features=in_features,
+            out_features=out_features,
+            bias=False,
+            register_buffers=True,
+        )
+        module.pack(linear, scales, s_channel)
+        module.post_init()
+        modules.append(module.eval())
+
+    return modules[0].cuda(), modules[1]
+
+
+@pytest.mark.cuda
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+def test_qqq_grouped_cuda_matches_torch_rounding_and_saturation():
+    if torch.cuda.get_device_capability()[0] < 8:
+        pytest.skip("QQQ CUDA requires compute capability >= 8.0")
+    if not qqq_module.qqq_runtime_available():
+        pytest.skip(qqq_module.qqq_runtime_error())
+
+    cases = [
+        (15, 22.0, 127),
+        (0, 20.0, -128),
+        (15, 18.0, 126),
+        (9, 0.5, 0),
+        (11, 0.5, 2),
+        (13, 0.5, 2),
+        (7, 0.5, 0),
+        (5, 0.5, -2),
+        (3, 0.5, -2),
+    ]
+    cuda_module, torch_module = _build_grouped_parity_modules(cases)
+    x = torch.zeros((1, cuda_module.in_features), dtype=torch.float16)
+    x[0, 0] = 1.0
+
+    expected = torch.tensor([case[2] for case in cases], dtype=torch.float16)
+    cuda_output = cuda_module(x.cuda()).cpu()[0, : len(cases)]
+    torch_output = torch_module(x)[0, : len(cases)]
+
+    torch.testing.assert_close(torch_output, expected, rtol=0, atol=0)
+    torch.testing.assert_close(cuda_output, expected, rtol=0, atol=0)
+
+
+@pytest.mark.cuda
+@pytest.mark.parametrize("tokens", [1, 17])
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+def test_qqq_grouped_cuda_matches_torch_for_regular_values(tokens, dtype):
+    if torch.cuda.get_device_capability()[0] < 8:
+        pytest.skip("QQQ CUDA requires compute capability >= 8.0")
+    if not qqq_module.qqq_runtime_available():
+        pytest.skip(qqq_module.qqq_runtime_error())
+
+    torch.manual_seed(42)
+    in_features = 256
+    out_features = 128
+    linear = torch.nn.Linear(in_features, out_features, bias=False, dtype=torch.float16)
+    linear.weight.data.normal_(0, 0.08)
+    scales = torch.rand((out_features, 2), dtype=torch.float16) * 0.02 + 0.01
+    s_channel = torch.rand(out_features, dtype=torch.float32) * 0.5 + 0.75
+
+    modules = []
+    for module_cls in (qqq_module.QQQLinear, qqq_module.QQQTorchLinear):
+        module = module_cls(
+            bits=4,
+            group_size=128,
+            sym=True,
+            desc_act=False,
+            in_features=in_features,
+            out_features=out_features,
+            bias=False,
+            register_buffers=True,
+        )
+        module.pack(linear, scales, s_channel)
+        module.post_init()
+        modules.append(module.eval())
+
+    x = torch.randn((tokens, in_features), dtype=dtype)
+    cuda_output = modules[0].cuda()(x.cuda()).cpu()
+    torch_output = modules[1](x)
+
+    assert cuda_output.dtype == dtype
+    torch.testing.assert_close(cuda_output, torch_output, rtol=0.02, atol=0.02)