diff --git a/bitsandbytes/backends/cuda/ops.py b/bitsandbytes/backends/cuda/ops.py index a0d9ffe83..b6badb885 100644 --- a/bitsandbytes/backends/cuda/ops.py +++ b/bitsandbytes/backends/cuda/ops.py @@ -87,7 +87,7 @@ def _(A: torch.Tensor, B: torch.Tensor, out: torch.Tensor): def _int8_linear_matmul_impl(A: torch.Tensor, B: torch.Tensor, out: torch.Tensor): - A, B = B, A + A, B = B.contiguous(), A.contiguous() shapeA = A.shape shapeB = B.shape @@ -183,6 +183,12 @@ def _( if col_stats.dtype != torch.float32: raise ValueError(f"col_stats must be float32, got {col_stats.dtype}") + A = A.contiguous() + row_stats = row_stats.contiguous() + col_stats = col_stats.contiguous() + if bias is not None: + bias = bias.contiguous() + # Note: cuda kernel only currently supports fp16 output. # We'll later cast to desired dtype if needed. out = torch.empty_like(A, dtype=torch.float16) @@ -217,6 +223,8 @@ def _(A: torch.Tensor, threshold=0.0): if threshold < 0.0: raise ValueError("threshold must be non-negative") + A = A.contiguous() + rows = A.numel() // A.shape[-1] cols = A.shape[-1] diff --git a/tests/test_linear8bitlt.py b/tests/test_linear8bitlt.py index b078e82b7..9616eaeae 100644 --- a/tests/test_linear8bitlt.py +++ b/tests/test_linear8bitlt.py @@ -364,3 +364,22 @@ def test_linear8bitlt_device_movement(device): # Accelerator outputs should match both times. torch.testing.assert_close(out_accelerator_2, out_accelerator, rtol=1e-8, atol=1e-8) + + +@pytest.mark.parametrize("device", get_available_devices()) +@pytest.mark.parametrize("threshold", [0.0, 6.0], ids=id_formatter("threshold")) +def test_linear8bitlt_non_contiguous_input(device, threshold): + # Pooling and last-token heads pass strided views of the hidden states to the layer, + # e.g. hidden_states[:, 0] (BERT-style pooler) or hidden_states[:, -1:, :]. + torch.manual_seed(0) + layer = Linear8bitLt(64, 32, has_fp16_weights=False, threshold=threshold).to(device) + hidden_states = torch.randn(8, 16, 64, dtype=torch.float16, device=device) + + for x in (hidden_states[:, 0], hidden_states[:, -1:, :]): + assert not x.is_contiguous() + with torch.no_grad(): + out = layer(x) + ref = layer(x.contiguous()) + # The bug reads unrelated rows, so the error is O(1); the tolerance only absorbs + # last-bit differences between strided and contiguous paths on some backends. + torch.testing.assert_close(out, ref, rtol=1e-2, atol=1e-2) diff --git a/tests/test_ops.py b/tests/test_ops.py index 4ca60f845..b38521bfb 100644 --- a/tests/test_ops.py +++ b/tests/test_ops.py @@ -472,3 +472,50 @@ def test_quantize_4bit_roundtrip_non_contiguous(self, device, dtype, quant_type, deq_c = torch.ops.bitsandbytes.dequantize_4bit(q_c, absmax_c, blocksize, quant_type, shape, dtype) torch.testing.assert_close(deq_nc, deq_c) + + @pytest.mark.parametrize("device", get_available_devices()) + @pytest.mark.parametrize("threshold", [0.0, 6.0], ids=id_formatter("threshold")) + def test_int8_vectorwise_quant_non_contiguous(self, device, threshold): + # Token pooling (hidden_states[:, 0]) gives a 2D view whose rows are not adjacent in memory. + A_full = torch.randn(8, 16, 64, dtype=torch.float16, device=device) + A_full[:, :, 3] = 10.0 # an outlier column, so the threshold path is exercised + A_noncontig = A_full[:, 0] + assert not A_noncontig.is_contiguous() + + A_contig = A_noncontig.contiguous() + + out_nc, stats_nc, outliers_nc = torch.ops.bitsandbytes.int8_vectorwise_quant(A_noncontig, threshold) + out_c, stats_c, outliers_c = torch.ops.bitsandbytes.int8_vectorwise_quant(A_contig, threshold) + + torch.testing.assert_close(stats_nc, stats_c) + torch.testing.assert_close(out_nc, out_c) + if threshold > 0.0: + torch.testing.assert_close(outliers_nc, outliers_c) + + @pytest.mark.parametrize("device", get_available_devices()) + def test_int8_linear_matmul_non_contiguous(self, device): + A_noncontig = torch.randint(-128, 127, (10, 3, 32), dtype=torch.int8, device=device)[:, 0] + B_noncontig = torch.randint(-128, 127, (32, 30), dtype=torch.int8, device=device).t() + assert not A_noncontig.is_contiguous() + assert not B_noncontig.is_contiguous() + + out_nc = torch.ops.bitsandbytes.int8_linear_matmul(A_noncontig, B_noncontig) + out_c = torch.ops.bitsandbytes.int8_linear_matmul(A_noncontig.contiguous(), B_noncontig.contiguous()) + + torch.testing.assert_close(out_nc, out_c) + + @pytest.mark.parametrize("device", get_available_devices()) + def test_int8_mm_dequant_non_contiguous(self, device): + A_noncontig = torch.randint(-128, 127, (4, 32, 64), dtype=torch.int32, device=device)[:, 0] + row_stats_noncontig = torch.rand(8, dtype=torch.float32, device=device)[::2] + col_stats_noncontig = torch.rand(128, dtype=torch.float32, device=device)[::2] + assert not A_noncontig.is_contiguous() + assert not row_stats_noncontig.is_contiguous() + assert not col_stats_noncontig.is_contiguous() + + out_nc = torch.ops.bitsandbytes.int8_mm_dequant(A_noncontig, row_stats_noncontig, col_stats_noncontig) + out_c = torch.ops.bitsandbytes.int8_mm_dequant( + A_noncontig.contiguous(), row_stats_noncontig.contiguous(), col_stats_noncontig.contiguous() + ) + + torch.testing.assert_close(out_nc, out_c)