Fix CUDA INT8 quantization of non-contiguous inputs

bitsandbytes/backends/cuda/ops.py

@@ -217,6 +217,9 @@ def _(A: torch.Tensor, threshold=0.0):
     if threshold < 0.0:
         raise ValueError("threshold must be non-negative")
 
+    # The CUDA kernel reads rows directly without accounting for tensor strides.
+    A = A.contiguous()
+
     rows = A.numel() // A.shape[-1]
     cols = A.shape[-1]
 

tests/test_linear8bitlt.py

@@ -22,6 +22,32 @@
 HAS_FP16_WEIGHTS_NON_DEPRECATED = [False]
 
 
+@pytest.mark.parametrize("device", get_available_devices())
+@pytest.mark.parametrize("threshold", [0.0, 6.0])
+@pytest.mark.parametrize("layout", ["transpose", "column_slice"])
+@pytest.mark.parametrize("requires_grad", TRUE_FALSE)
+def test_linear8bit_non_contiguous_input(device, threshold, layout, requires_grad):
+    linear = Linear8bitLt(64, 32, has_fp16_weights=False, threshold=threshold).half().to(device).eval()
+    if layout == "transpose":
+        x = torch.randn(64, 4, device=device, dtype=torch.float16).t()
+    else:
+        x = torch.randn(4, 128, device=device, dtype=torch.float16)[:, 1::2]
+    x[1, 3] = 10.0
+    assert not x.is_contiguous()
+    x_ref = x.contiguous().requires_grad_(requires_grad)
+    x.requires_grad_(requires_grad)
+
+    expected = linear(x_ref)
+    actual = linear(x)
+    torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+
+    if requires_grad:
+        grad = torch.randn_like(actual)
+        expected.backward(grad)
+        actual.backward(grad)
+        torch.testing.assert_close(x.grad, x_ref.grad, rtol=0, atol=0)
+
+
 # contributed by Alex Borzunov, see:
 # https://github.com/bigscience-workshop/petals/blob/main/tests/test_linear8bitlt.py
 @pytest.mark.parametrize("device", get_available_devices())

tests/test_ops.py

@@ -62,6 +62,43 @@ def test_int8_vectorwise_quant(self, threshold, device):
         opcheck(torch.ops.bitsandbytes.int8_vectorwise_quant, (A,))
         opcheck(torch.ops.bitsandbytes.int8_vectorwise_quant, (A, threshold))
 
+    @pytest.mark.parametrize("device", get_available_devices())
+    @pytest.mark.parametrize("threshold", [0.0, 6.0])
+    @pytest.mark.parametrize("layout", ["transpose", "row_slice", "column_slice"])
+    @pytest.mark.parametrize("double_quant", TRUE_FALSE)
+    def test_int8_quant_non_contiguous(self, device, threshold, layout, double_quant):
+        if double_quant and device != "cuda":
+            pytest.skip("int8_double_quant is only implemented for CUDA/ROCm")
+
+        if layout == "transpose":
+            A = torch.randn(64, 8, device=device, dtype=torch.float16).t()
+        elif layout == "row_slice":
+            A = torch.randn(16, 64, device=device, dtype=torch.float16)[1::2]
+        else:
+            A = torch.randn(8, 128, device=device, dtype=torch.float16)[:, 1::2]
+        A[1, 3] = 10.0
+        A[6, 17] = -10.0
+        assert not A.is_contiguous()
+        original = A.clone()
+
+        quantize = (
+            bitsandbytes.functional.int8_double_quant
+            if double_quant
+            else bitsandbytes.functional.int8_vectorwise_quant
+        )
+        expected = quantize(A.contiguous(), threshold=threshold)
+        actual = quantize(A, threshold=threshold)
+
+        for result, reference in zip(actual, expected):
+            if reference is None:
+                assert result is None
+            else:
+                torch.testing.assert_close(result, reference, rtol=0, atol=0)
+        torch.testing.assert_close(A, original, rtol=0, atol=0)
+
+        if device == "cuda" and not double_quant:
+            opcheck(torch.ops.bitsandbytes.int8_vectorwise_quant, (A, threshold))
+
     @pytest.mark.parametrize("device", get_available_devices())
     def test_int8_mm_dequant(self, device):
         A = torch.randint(-128, 127, (256, 256), dtype=torch.int32, device=device)