- | 🐧 Linux |
+ 🐧 Linux, glibc >= 2.24 |
| x86-64 |
◻️ CPU |
- |
+ AVX2 |
〰️ Partial Support |
|
- 🟩 NVIDIA GPU |
+ 🟩 NVIDIA GPU
cuda |
SM50+ minimum
SM75+ recommended |
- ✅ Full Support * |
+ ✅ Full Support |
|
- 🟥 AMD GPU |
- gfx90a, gfx942, gfx1100 |
+ 🟥 AMD GPU
cuda |
+
+ CDNA: gfx90a, gfx942
+ RDNA: gfx1100, gfx1200
+ |
🚧 In Development |
|
- 🟦 Intel XPU |
+ 🟦 Intel GPU
xpu |
- Data Center GPU Max Series (Ponte Vecchio)
- Arc A-Series (Alchemist)
+ Data Center GPU Max Series
+ Arc A-Series (Alchemist)
Arc B-Series (Battlemage)
|
🚧 In Development |
-
| aarch64 |
◻️ CPU |
@@ -82,12 +83,12 @@ bitsandbytes has the following minimum requirements for all platforms:
|
- 🟩 NVIDIA GPU |
+ 🟩 NVIDIA GPU
cuda |
SM75, SM80, SM90, SM100 |
- ✅ Full Support * |
+ ✅ Full Support |
- | 🪟 Windows |
+ 🪟 Windows 11 / Windows Server 2019+ |
| x86-64 |
@@ -97,13 +98,13 @@ bitsandbytes has the following minimum requirements for all platforms:
|
- 🟩 NVIDIA GPU |
+ 🟩 NVIDIA GPU
cuda |
SM50+ minimum
SM75+ recommended |
- ✅ Full Support * |
+ ✅ Full Support |
|
- 🟦 Intel XPU |
+ 🟦 Intel GPU
xpu |
Arc A-Series (Alchemist)
Arc B-Series (Battlemage)
@@ -111,19 +112,22 @@ bitsandbytes has the following minimum requirements for all platforms:
| 🚧 In Development |
- | 🍎 macOS |
+ 🍎 macOS 13.1+ |
| arm64 |
- ◻️ CPU / Metal |
+ ◻️ CPU |
Apple M1+ |
- ❌ Under consideration |
+ 🛣️ Future Roadmap |
+
+ |
+ ⬜ Metal
mps |
+ Apple M1+ |
+ 🛣️ Future Roadmap |
-\* Accelerated INT8 requires SM75+.
-
## :book: Documentation
* [Official Documentation](https://huggingface.co/docs/bitsandbytes/main)
* 🤗 [Transformers](https://huggingface.co/docs/transformers/quantization/bitsandbytes)
From 4def9590abb8a3f0ef789fce0b1659af729643e4 Mon Sep 17 00:00:00 2001
From: MISHANMAURYA <118961433+MISHANMAURYA@users.noreply.github.com>
Date: Thu, 22 May 2025 20:51:50 +0530
Subject: [PATCH 16/85] Update ops.py
---
bitsandbytes/backends/cuda/ops.py | 14 +++++---------
1 file changed, 5 insertions(+), 9 deletions(-)
diff --git a/bitsandbytes/backends/cuda/ops.py b/bitsandbytes/backends/cuda/ops.py
index 14f55847c..5b94c5349 100644
--- a/bitsandbytes/backends/cuda/ops.py
+++ b/bitsandbytes/backends/cuda/ops.py
@@ -11,7 +11,6 @@
from ...cextension import lib, HIP_ENVIRONMENT
-
@register_kernel("bitsandbytes::int8_linear_matmul", "cuda")
def _(A: torch.Tensor, B: torch.Tensor):
out = torch.empty((*A.shape[:-1], B.shape[0]), device=A.device, dtype=torch.int32)
@@ -78,12 +77,9 @@ def _int8_linear_matmul_impl(A: torch.Tensor, B: torch.Tensor, out: torch.Tensor
ldb = ct.c_int32(ldb)
ldc = ct.c_int32(ldc)
stream = _get_tensor_stream(A)
-
- if dtype == torch.int32:
- has_error = lib.cigemmlt_32(ctx, m, n, k, ptrA, ptrB, ptrC, ptrRowScale, lda, ldb, ldc, stream)
- else:
- has_error = lib.cigemmlt_8(ctx, m, n, k, ptrA, ptrB, ptrC, ptrRowScale, lda, ldb, ldc, stream)
-
+
+ has_error = lib.cigemmlt_32(ctx, m, n, k, ptrA, ptrB, ptrC, ptrRowScale, lda, ldb, ldc, stream)
+
if has_error:
if has_error == 100:
# `ERR_NOT_IMPLEMENTED` is defined as 100 in `ops.cu`
@@ -96,6 +92,7 @@ def _int8_linear_matmul_impl(A: torch.Tensor, B: torch.Tensor, out: torch.Tensor
return out
+
@register_kernel("bitsandbytes::int8_mm_dequant", "cuda")
def _(
A: torch.Tensor,
@@ -384,6 +381,7 @@ def _(
return out, absmax
+
@register_kernel("bitsandbytes::dequantize_4bit", "cuda")
def _(
A: torch.Tensor,
@@ -398,7 +396,6 @@ def _(
return out
-
@register_kernel("bitsandbytes::dequantize_4bit.out", "cuda")
def _(
A: torch.Tensor,
@@ -496,7 +493,6 @@ def _(
torch._check(out.dtype == A.dtype, lambda: f"Expected out.dtype == {A.dtype}, got {out.dtype}")
_gemv_4bit_impl(A, B, shapeB, absmax, code, blocksize, out=out)
-
def _gemv_4bit_impl(
A: torch.Tensor,
B: torch.Tensor,
From 0f318667aaf4de15cd29f8063dcaa4fd90d24783 Mon Sep 17 00:00:00 2001
From: MISHANMAURYA <118961433+MISHANMAURYA@users.noreply.github.com>
Date: Thu, 22 May 2025 21:31:55 +0530
Subject: [PATCH 17/85] Update functional.py
---
bitsandbytes/functional.py | 157 ++++++++++++-------------------------
1 file changed, 48 insertions(+), 109 deletions(-)
diff --git a/bitsandbytes/functional.py b/bitsandbytes/functional.py
index 7730f7182..3f0c1ff94 100644
--- a/bitsandbytes/functional.py
+++ b/bitsandbytes/functional.py
@@ -728,11 +728,9 @@ def quantize_blockwise(
nested=False,
) -> tuple[torch.Tensor, QuantState]:
"""Quantize a tensor in blocks of values.
-
The input tensor is quantized by dividing it into blocks of `blocksize` values.
The the absolute maximum value within these blocks is calculated for scaling
the non-linear quantization.
-
Args:
A (`torch.Tensor`): The input tensor. Supports `float16`, `bfloat16`, or `float32` datatypes.
code (`torch.Tensor`, *optional*):
@@ -744,10 +742,8 @@ def quantize_blockwise(
The size of the blocks. Defaults to 4096.
Valid values are 64, 128, 256, 512, 1024, 2048, and 4096.
nested (`bool`, *optional*): Whether to additionally quantize the absmax values. Defaults to False.
-
Raises:
ValueError: Raised when the input data type is not supported.
-
Returns:
`Tuple[torch.Tensor, QuantState]`: A tuple containing the quantization results.
- `torch.Tensor`: The quantized tensor.
@@ -759,61 +755,23 @@ def quantize_blockwise(
name2qmap["dynamic"] = create_dynamic_map().to(A.device)
code = name2qmap["dynamic"]
- if absmax is None:
- n = A.numel()
- blocks = -(n // -blocksize)
- absmax = torch.zeros((blocks,), device=A.device, dtype=torch.float32)
-
- if out is None:
- out = torch.zeros_like(A, dtype=torch.uint8)
-
device_type = A.device.type
-
- if device_type == "cpu":
- code = code.cpu()
- lib.cquantize_blockwise_cpu_fp32(
- get_ptr(code),
- get_ptr(A),
- get_ptr(absmax),
- get_ptr(out),
- ct.c_longlong(blocksize),
- ct.c_longlong(A.numel()),
- )
- elif device_type in ["cuda", "hip"]:
+ if device_type in ["cuda", "hip"]:
if not HIP_ENVIRONMENT:
- assert blocksize in [4096, 2048, 1024, 512, 256, 128, 64]
+ assert blocksize in [4096, 2048, 1024, 512, 256, 128, 64]
else:
- assert blocksize in [4096, 2048, 1024, 512, 256, 128]
-
- code = code.to(A.device)
-
- is_on_gpu([A, out, absmax])
+ assert blocksize in [4096, 2048, 1024, 512, 256, 128]
- with _cuda_device_of(A):
- args = (
- get_ptr(code),
- get_ptr(A),
- get_ptr(absmax),
- get_ptr(out),
- ct.c_int32(blocksize),
- ct.c_int(A.numel()),
- )
-
- if A.dtype == torch.float16:
- lib.cquantize_blockwise_fp16(*args)
- elif A.dtype == torch.bfloat16:
- lib.cquantize_blockwise_bf16(*args)
- elif A.dtype == torch.float32:
- lib.cquantize_blockwise_fp32(*args)
- else:
- raise ValueError(f"Blockwise quantization only supports 16/32-bit floats, but got {A.dtype}")
- else:
- raise RuntimeError(f"Device type {device_type} not supported for quantization")
+ _out, _absmax = torch.ops.bitsandbytes.quantize_blockwise.default(
+ A,
+ code.to(A.device),
+ blocksize,
+ )
if nested:
- offset = absmax.mean()
- absmax -= offset
- qabsmax, state2 = quantize_blockwise(absmax, blocksize=blocksize, nested=False)
+ offset = _absmax.mean()
+ _absmax -= offset
+ qabsmax, state2 = quantize_blockwise(_absmax, blocksize=blocksize, nested=False)
quant_state = QuantState(
absmax=qabsmax,
code=code,
@@ -823,11 +781,18 @@ def quantize_blockwise(
state2=state2,
)
else:
- quant_state = QuantState(absmax=absmax, code=code, blocksize=blocksize, dtype=A.dtype)
+ quant_state = QuantState(absmax=_absmax, code=code.to(A.device), blocksize=blocksize, dtype=A.dtype)
+
+ # TODO(matthewdouglas): Deprecate out kwarg
+ out = out.copy_(_out) if out is not None else _out
+
+ # TODO(matthewdouglas): Deprecate absmax kwarg
+ if absmax is not None:
+ quant_state.absmax = absmax.copy_(quant_state.absmax)
+
+ return out, quant_state
+
- return out, quant_state
-
-
def dequantize_blockwise(
A: torch.Tensor,
quant_state: Optional[QuantState] = None,
@@ -838,11 +803,9 @@ def dequantize_blockwise(
nested=False,
) -> torch.Tensor:
"""Dequantize a tensor in blocks of values.
-
The input tensor is dequantized by dividing it into blocks of `blocksize` values.
The the absolute maximum value within these blocks is used for scaling
the non-linear dequantization.
-
Args:
A (`torch.Tensor`): The quantized input tensor.
quant_state ([`QuantState`], *optional*):
@@ -860,10 +823,8 @@ def dequantize_blockwise(
The size of the blocks. Defaults to 4096.
Valid values are 64, 128, 256, 512, 1024, 2048, and 4096.
Ignored when `quant_state` is provided.
-
Raises:
ValueError: Raised when the input data type is not supported.
-
Returns:
`torch.Tensor`:
The dequantized tensor. The datatype is indicated by `quant_state.dtype` and defaults to `torch.float32`.
@@ -878,6 +839,16 @@ def dequantize_blockwise(
if quant_state is None:
quant_state = QuantState(absmax=absmax, code=code, blocksize=blocksize, dtype=torch.float32)
+ device_type = A.device.type
+ if device_type in ["cuda", "hip"]:
+ supported_blocksizes = [4096, 2048, 1024, 512, 256, 128, 64]
+ if HIP_ENVIRONMENT:
+ supported_blocksizes = supported_blocksizes[:-1]
+ if quant_state.blocksize not in supported_blocksizes:
+ raise ValueError(
+ f"The blocksize of {quant_state.blocksize} is not supported. Supported values: {supported_blocksizes}"
+ )
+
absmax = quant_state.absmax
if quant_state.nested:
absmax = dequantize_blockwise(quant_state.absmax, quant_state.state2)
@@ -885,56 +856,24 @@ def dequantize_blockwise(
if absmax.dtype != torch.float32:
absmax = absmax.float()
- if out is None:
- out = torch.empty(A.shape, dtype=quant_state.dtype, device=A.device)
-
- device_type = A.device.type
-
- if device_type == "cpu":
- code = quant_state.code.cpu()
- lib.cdequantize_blockwise_cpu_fp32(
- get_ptr(code),
- get_ptr(A),
- get_ptr(quant_state.absmax),
- get_ptr(out),
- ct.c_longlong(quant_state.blocksize),
- ct.c_longlong(A.numel()),
+ if out is not None:
+ torch.ops.bitsandbytes.dequantize_blockwise.out(
+ A,
+ absmax,
+ quant_state.code.to(A.device),
+ quant_state.blocksize,
+ quant_state.dtype,
+ out=out,
)
- elif device_type in ["cuda", "hip"]:
- code = quant_state.code.to(A.device)
- supported_blocksizes = [2048, 4096, 1024, 512, 256, 128, 64]
- if HIP_ENVIRONMENT:
- supported_blocksizes = supported_blocksizes[:-1]
- if quant_state.blocksize not in supported_blocksizes:
- raise ValueError(
- f"The blocksize of {quant_state.blocksize} is not supported. Supported values: {supported_blocksizes}",
- )
-
- is_on_gpu([A, absmax, out])
-
- with _cuda_device_of(A):
- args = (
- get_ptr(quant_state.code),
- get_ptr(A),
- get_ptr(absmax),
- get_ptr(out),
- ct.c_int(quant_state.blocksize),
- ct.c_int(A.numel()),
- _get_tensor_stream(A),
- )
-
- if out.dtype == torch.float16:
- lib.cdequantize_blockwise_fp16(*args)
- elif out.dtype == torch.bfloat16:
- lib.cdequantize_blockwise_bf16(*args)
- elif out.dtype == torch.float32:
- lib.cdequantize_blockwise_fp32(*args)
- else:
- raise ValueError(f"Blockwise quantization only supports 16/32-bit floats, but got {out.dtype}")
- else:
- raise RuntimeError(f"Device type {device_type} not supported for dequantization")
+ return out
- return out
+ return torch.ops.bitsandbytes.dequantize_blockwise.default(
+ A,
+ absmax,
+ quant_state.code.to(A.device),
+ quant_state.blocksize,
+ quant_state.dtype,
+ )
def get_4bit_type(typename, device=None, blocksize=64):
From 190faed7e96b8b27e033fe3c6ee5e3a6d5a4772a Mon Sep 17 00:00:00 2001
From: MISHANMAURYA <118961433+MISHANMAURYA@users.noreply.github.com>
Date: Thu, 22 May 2025 23:35:15 +0530
Subject: [PATCH 18/85] Update ops.py
---
bitsandbytes/backends/cuda/ops.py | 12 +++---------
1 file changed, 3 insertions(+), 9 deletions(-)
diff --git a/bitsandbytes/backends/cuda/ops.py b/bitsandbytes/backends/cuda/ops.py
index 5b94c5349..ff5e023cc 100644
--- a/bitsandbytes/backends/cuda/ops.py
+++ b/bitsandbytes/backends/cuda/ops.py
@@ -52,15 +52,9 @@ def _int8_linear_matmul_impl(A: torch.Tensor, B: torch.Tensor, out: torch.Tensor
# cuBLASLt does not support int8 matmul with inner dimensions that are not divisible by 4.
# We'll fall back to a slower fp32 calculation in this circumstance.
# Fortunately, this should not be very common.
-
- if lda % 4 != 0:
- result = torch.matmul(B.float(), A.float().t()).to(torch.int32)
- if out is not None:
- result = out.copy_(result)
- return result
-
- if out is None:
- out = torch.empty(shapeC, device=A.device, dtype=dtype)
+ if lda % 4 != 0:
+ result = torch.matmul(B.float(), A.float().t()).to(torch.int32)
+ return out.copy_(result)
is_on_gpu([A, B, out])
From d7f413b9b367b9b26b87180095ebcc7a561fdc26 Mon Sep 17 00:00:00 2001
From: MISHANMAURYA <118961433+MISHANMAURYA@users.noreply.github.com>
Date: Thu, 22 May 2025 23:52:39 +0530
Subject: [PATCH 19/85] Update ops.py
---
bitsandbytes/backends/cuda/ops.py | 30 +++++-------------------------
1 file changed, 5 insertions(+), 25 deletions(-)
diff --git a/bitsandbytes/backends/cuda/ops.py b/bitsandbytes/backends/cuda/ops.py
index ff5e023cc..b75f67d62 100644
--- a/bitsandbytes/backends/cuda/ops.py
+++ b/bitsandbytes/backends/cuda/ops.py
@@ -55,8 +55,6 @@ def _int8_linear_matmul_impl(A: torch.Tensor, B: torch.Tensor, out: torch.Tensor
if lda % 4 != 0:
result = torch.matmul(B.float(), A.float().t()).to(torch.int32)
return out.copy_(result)
-
- is_on_gpu([A, B, out])
with _cuda_device_of(A):
ctx = CUBLAS_Context.get_instance().get_context(A.device)
@@ -114,8 +112,6 @@ def _(
# TODO(matthewdouglas): Consider supporting bf16 fused bias
ptrBias = get_ptr(bias) if bias is not None and bias.dtype == torch.float16 else None
- is_on_gpu([A, row_stats, col_stats, out, bias])
-
with _cuda_device_of(A):
lib.cdequant_mm_int32_fp16(
ptrA, ptrRowStats, ptrColStats, ptrOut, ptrBias, numRows, numCols, _get_tensor_stream(A)
@@ -133,8 +129,6 @@ def _(A: torch.Tensor, threshold=0.0):
torch._check(A.dtype == torch.float16, lambda: f"A must be float16, got {A.dtype}")
torch._check(threshold >= 0.0, lambda: "threshold must be non-negative")
- is_on_gpu([A])
-
rows = prod(A.shape[:-1])
cols = A.shape[-1]
@@ -231,9 +225,7 @@ def _(A: torch.Tensor, code: torch.Tensor, blocksize: int) -> tuple[torch.Tensor
blocks = -(n // -blocksize)
absmax = torch.empty((blocks,), device=A.device, dtype=torch.float32)
out = torch.empty_like(A, dtype=torch.uint8)
-
- is_on_gpu([A, out, absmax])
-
+
if device_type == 'cuda' or (device_type == 'hip' or HIP_ENVIRONMENT):
with _cuda_device_of(A):
args = (
@@ -295,9 +287,7 @@ def _dequantize_blockwise_impl(
dtype in [torch.float16, torch.bfloat16, torch.float32],
lambda: f"Blockwise dequantization only supports 16bit/32bit floating types, got {dtype}",
)
-
- is_on_gpu([A, absmax, out])
-
+
if device_type == 'cuda' or (device_type == 'hip' or HIP_ENVIRONMENT):
with _cuda_device_of(A):
args = (
@@ -341,9 +331,7 @@ def _(
blocks = -(n // -blocksize)
absmax = torch.empty((blocks,), device=A.device, dtype=torch.float32)
out = torch.empty(((n + 1) // (quant_storage.itemsize * 2), 1), device=A.device, dtype=quant_storage)
-
- is_on_gpu([A, out, absmax])
-
+
if device_type == 'cuda' or (device_type == 'hip' or HIP_ENVIRONMENT):
with _cuda_device_of(A):
args = (
@@ -370,8 +358,6 @@ def _(
lib.cquantize_blockwise_fp32_fp4(*args)
else:
lib.cquantize_blockwise_fp32_nf4(*args)
- else:
- raise NotImplementedError(f"4-bit quantization not implemented for {device_type}")
return out, absmax
@@ -400,11 +386,11 @@ def _(
dtype: torch.dtype,
out: torch.Tensor,
) -> None:
-
torch._check(out.shape == shape, lambda: f"Expected out.shape == {shape}, got {out.shape}")
torch._check(out.dtype == dtype, lambda: f"Expected out.dtype == {dtype}, got {out.dtype}")
_dequantize_4bit_impl(A, absmax, blocksize, quant_type, dtype, out=out)
+
def _dequantize_4bit_impl(
A: torch.Tensor,
absmax: torch.Tensor,
@@ -426,9 +412,7 @@ def _dequantize_4bit_impl(
dtype in [torch.bfloat16, torch.float16, torch.float32],
lambda: f"Blockwise 4bit dequantization only supports 16/32-bit floats, but got {dtype}",
)
-
- is_on_gpu([A, absmax, out])
-
+
if device_type == 'cuda' or (device_type == 'hip' or HIP_ENVIRONMENT):
with _cuda_device_of(A):
args = (
@@ -456,8 +440,6 @@ def _dequantize_4bit_impl(
lib.cdequantize_blockwise_fp32_fp4(*args)
else:
lib.cdequantize_blockwise_fp32_nf4(*args)
- else:
- raise NotImplementedError(f"4-bit dequantization not implemented for {device_type}")
@register_kernel("bitsandbytes::gemv_4bit", "cuda")
@@ -520,8 +502,6 @@ def _gemv_4bit_impl(
ldb = ct.c_int32((A.shape[-1] + 1) // 2)
ldc = m
- is_on_gpu([B, A, out, absmax])
-
stream = _get_tensor_stream(A)
with _cuda_device_of(A):
From 3b6e68a001b0dce3b368129599335fcb569ac5cd Mon Sep 17 00:00:00 2001
From: MISHANMAURYA <118961433+MISHANMAURYA@users.noreply.github.com>
Date: Fri, 23 May 2025 00:05:43 +0530
Subject: [PATCH 20/85] Update ops.py
---
bitsandbytes/backends/cuda/ops.py | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/bitsandbytes/backends/cuda/ops.py b/bitsandbytes/backends/cuda/ops.py
index b75f67d62..156125c9f 100644
--- a/bitsandbytes/backends/cuda/ops.py
+++ b/bitsandbytes/backends/cuda/ops.py
@@ -5,7 +5,7 @@
import torch
-from bitsandbytes.functional import CUBLAS_Context, _cuda_device_of, _get_tensor_stream, get_ptr, is_on_gpu
+from bitsandbytes.functional import CUBLAS_Context, _cuda_device_of, _get_tensor_stream, get_ptr
from ..._ops import register_kernel
from ...cextension import lib, HIP_ENVIRONMENT
From 06740b1372a9c9751216b76dc4c8cc98514905dd Mon Sep 17 00:00:00 2001
From: MISHANMAURYA <118961433+MISHANMAURYA@users.noreply.github.com>
Date: Fri, 23 May 2025 01:53:30 +0530
Subject: [PATCH 21/85] Update ops.py
---
bitsandbytes/backends/cuda/ops.py | 989 +++++++++++++++---------------
1 file changed, 486 insertions(+), 503 deletions(-)
diff --git a/bitsandbytes/backends/cuda/ops.py b/bitsandbytes/backends/cuda/ops.py
index 156125c9f..48dc75135 100644
--- a/bitsandbytes/backends/cuda/ops.py
+++ b/bitsandbytes/backends/cuda/ops.py
@@ -1,325 +1,312 @@
-from collections.abc import Sequence
-import ctypes as ct
-from math import prod
-from typing import Optional
-
-import torch
-
-from bitsandbytes.functional import CUBLAS_Context, _cuda_device_of, _get_tensor_stream, get_ptr
-
-from ..._ops import register_kernel
-from ...cextension import lib, HIP_ENVIRONMENT
-
-
-@register_kernel("bitsandbytes::int8_linear_matmul", "cuda")
-def _(A: torch.Tensor, B: torch.Tensor):
- out = torch.empty((*A.shape[:-1], B.shape[0]), device=A.device, dtype=torch.int32)
- return _int8_linear_matmul_impl(A, B, out)
-
-
-@register_kernel("bitsandbytes::int8_linear_matmul.out", "cuda")
-def _(A: torch.Tensor, B: torch.Tensor, out: torch.Tensor):
- _int8_linear_matmul_impl(A, B, out)
-
-
-def _int8_linear_matmul_impl(A: torch.Tensor, B: torch.Tensor, out: torch.Tensor):
- A, B = B, A
-
- shapeA = A.shape
- shapeB = B.shape
-
- torch._check(A.dtype == torch.int8, lambda: "B must be int8")
- torch._check(B.dtype == torch.int8, lambda: "A must be int8")
- torch._check(A.ndim == 2, lambda: "Only two dimensional matrices are supported for argument B")
- torch._check(B.ndim in [2, 3], lambda: "Only two or three dimensional matrices are supported for argument A")
- torch._check(prod(shapeB) > 0, lambda: f"Input tensor dimensions need to be > 0: {shapeB}")
- torch._check(out.dtype == torch.int32)
-
- shapeC = (*shapeB[:-1], shapeA[0])
- torch._check(out.shape == shapeC, lambda: f"Output shape {out.shape} does not match expected shape {shapeC}")
-
- k, m = shapeA
- n = prod(shapeB[:-1])
- lda = shapeA[-1] # Weights (outputs, inputs)
- ldb = shapeB[-1] # Activations (batch, tokens, inputs)
- ldc = shapeC[-1] # Output (batch, tokens, outputs)
-
- torch._check(
- lda == ldb,
- lambda: f"int8_linear_matmul only supports B^T @ A. Inner dimensions do not match: B @ A = {shapeB} @ {shapeA}",
- )
-
- # cuBLASLt does not support int8 matmul with inner dimensions that are not divisible by 4.
- # We'll fall back to a slower fp32 calculation in this circumstance.
- # Fortunately, this should not be very common.
- if lda % 4 != 0:
- result = torch.matmul(B.float(), A.float().t()).to(torch.int32)
- return out.copy_(result)
-
- with _cuda_device_of(A):
- ctx = CUBLAS_Context.get_instance().get_context(A.device)
- ptrA = get_ptr(A)
- ptrB = get_ptr(B)
- ptrC = get_ptr(out)
- ptrRowScale = None
- m = ct.c_int32(m)
- n = ct.c_int32(n)
- k = ct.c_int32(k)
- lda = ct.c_int32(lda)
- ldb = ct.c_int32(ldb)
- ldc = ct.c_int32(ldc)
- stream = _get_tensor_stream(A)
-
- has_error = lib.cigemmlt_32(ctx, m, n, k, ptrA, ptrB, ptrC, ptrRowScale, lda, ldb, ldc, stream)
-
- if has_error:
- if has_error == 100:
- # `ERR_NOT_IMPLEMENTED` is defined as 100 in `ops.cu`
- # TODO: Warn and implement a fallback to fp32 compute?
- raise NotImplementedError("int8_linear_matmul not implemented!")
- else:
- raise RuntimeError(
- f"cublasLt ran into an error!\n\t{shapeA=}, {shapeB=}, {shapeC=}\n\t{(lda, ldb, ldc)=}\n\t{(m, n, k)=}"
- )
-
- return out
-
-
-@register_kernel("bitsandbytes::int8_mm_dequant", "cuda")
-def _(
- A: torch.Tensor,
- row_stats: torch.Tensor,
- col_stats: torch.Tensor,
- dtype: Optional[torch.dtype] = None,
- bias: Optional[torch.Tensor] = None,
-) -> torch.Tensor:
- torch._check(A.dtype == torch.int32, lambda: f"A must be int32, got {A.dtype}")
- torch._check(row_stats.dtype == torch.float32, lambda: f"row_stats must be float32, got {row_stats.dtype}")
- torch._check(col_stats.dtype == torch.float32, lambda: f"col_stats must be float32, got {col_stats.dtype}")
-
- # 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)
-
- ptrA = get_ptr(A)
- ptrOut = get_ptr(out)
- ptrRowStats = get_ptr(row_stats)
- ptrColStats = get_ptr(col_stats)
- numRows = ct.c_int32(prod(A.shape[:-1]))
- numCols = ct.c_int32(A.shape[-1])
-
- # Note: fused bias in the kernel is only supported for fp16
- # TODO(matthewdouglas): Consider supporting bf16 fused bias
- ptrBias = get_ptr(bias) if bias is not None and bias.dtype == torch.float16 else None
-
- with _cuda_device_of(A):
- lib.cdequant_mm_int32_fp16(
- ptrA, ptrRowStats, ptrColStats, ptrOut, ptrBias, numRows, numCols, _get_tensor_stream(A)
- )
-
- # Add bias separately if not fused in kernel
- if bias is not None and bias.dtype != torch.float16:
- out.add_(bias)
-
- return out.to(dtype or torch.float16)
-
-
-@register_kernel("bitsandbytes::int8_vectorwise_quant", "cuda")
-def _(A: torch.Tensor, threshold=0.0):
- torch._check(A.dtype == torch.float16, lambda: f"A must be float16, got {A.dtype}")
- torch._check(threshold >= 0.0, lambda: "threshold must be non-negative")
-
- rows = prod(A.shape[:-1])
- cols = A.shape[-1]
-
- row_stats = torch.empty(rows, device=A.device, dtype=torch.float32)
- out_row = torch.empty(A.shape, device=A.device, dtype=torch.int8)
-
- outlier_cols = None
-
- if threshold > 0.0:
- # TODO we could improve perf of this
- outliers = A.abs() >= threshold
-
- if outliers.any():
- outlier_cols = torch.argwhere(outliers.any(dim=0)).view(-1)
- else:
- # Needed for torch.compile support.
- outlier_cols = torch.empty(0, device=A.device, dtype=torch.int64)
-
- with _cuda_device_of(A):
- lib.cint8_vector_quant(
- get_ptr(A),
- get_ptr(out_row),
- get_ptr(row_stats),
- ct.c_float(threshold),
- ct.c_int32(rows),
- ct.c_int32(cols),
- _get_tensor_stream(A),
- )
-
- # Zero out values from outlier columns across all rows.
- # The kernel will handle this for outliers themselves, so we can optimize for rows=1.
- if rows > 1 and outlier_cols is not None:
- out_row[:, outlier_cols] = 0
-
- return out_row, row_stats, outlier_cols
-
-
-@register_kernel("bitsandbytes::int8_double_quant", "cuda")
-def _(
- A: torch.Tensor,
- threshold=0.0,
-) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, Optional[torch.Tensor]]:
- # Use CUDA kernel for rowwise and COO tensor
- quant_row, row_stats, outlier_cols = torch.ops.bitsandbytes.int8_vectorwise_quant.default(
- A,
- threshold=threshold,
- )
-
- # PyTorch impl for colwise
- col_stats, outlier_mask = _get_col_absmax(A, threshold=threshold)
- if threshold > 0.0 and outlier_mask is not None:
- A = A.masked_fill(outlier_mask, 0.0)
- quant_col = torch.round(A.mul(127.0) / col_stats.unsqueeze(0)).to(torch.int8)
-
- return quant_row, quant_col, row_stats, col_stats.flatten().float(), outlier_cols
-
-
-def _get_col_absmax(
- A: torch.Tensor,
- threshold=0.0,
-) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
- torch._check(A.is_floating_point())
-
- outlier_mask = None
-
- absA = A.abs().view(-1, A.shape[-1])
-
- if threshold > 0.0:
- # Filter outliers from stats when enabled
- outlier_mask = absA >= threshold
- absA.masked_fill_(outlier_mask, 0.0)
-
- # shape [cols]; unsqueeze(0) gives [1,cols]
- col_stats = absA.amax(dim=0, keepdim=False).float()
-
- return col_stats, outlier_mask
-
-
-@register_kernel("bitsandbytes::quantize_blockwise", "cuda")
-def _(A: torch.Tensor, code: torch.Tensor, blocksize: int) -> tuple[torch.Tensor, torch.Tensor]:
- torch._check_is_size(blocksize)
-
- device = A.device
- device_type = device.type
+from collections.abc import Sequence
+import ctypes as ct
+from math import prod
+from typing import Optional
+
+import torch
+
+from bitsandbytes.functional import CUBLAS_Context, _cuda_device_of, _get_tensor_stream, get_ptr
+
+from ..._ops import register_kernel
+from ...cextension import lib, HIP_ENVIRONMENT
+
+
+@register_kernel("bitsandbytes::int8_linear_matmul", "cuda")
+def _(A: torch.Tensor, B: torch.Tensor):
+ out = torch.empty((*A.shape[:-1], B.shape[0]), device=A.device, dtype=torch.int32)
+ return _int8_linear_matmul_impl(A, B, out)
+
+
+@register_kernel("bitsandbytes::int8_linear_matmul.out", "cuda")
+def _(A: torch.Tensor, B: torch.Tensor, out: torch.Tensor):
+ _int8_linear_matmul_impl(A, B, out)
+
+
+def _int8_linear_matmul_impl(A: torch.Tensor, B: torch.Tensor, out: torch.Tensor):
+ A, B = B, A
+
+ shapeA = A.shape
+ shapeB = B.shape
+
+ torch._check(A.dtype == torch.int8, lambda: "B must be int8")
+ torch._check(B.dtype == torch.int8, lambda: "A must be int8")
+ torch._check(A.ndim == 2, lambda: "Only two dimensional matrices are supported for argument B")
+ torch._check(B.ndim in [2, 3], lambda: "Only two or three dimensional matrices are supported for argument A")
+ torch._check(prod(shapeB) > 0, lambda: f"Input tensor dimensions need to be > 0: {shapeB}")
+ torch._check(out.dtype == torch.int32)
+
+ shapeC = (*shapeB[:-1], shapeA[0])
+ torch._check(out.shape == shapeC, lambda: f"Output shape {out.shape} does not match expected shape {shapeC}")
+
+ k, m = shapeA
+ n = prod(shapeB[:-1])
+ lda = shapeA[-1] # Weights (outputs, inputs)
+ ldb = shapeB[-1] # Activations (batch, tokens, inputs)
+ ldc = shapeC[-1] # Output (batch, tokens, outputs)
+
+ torch._check(
+ lda == ldb,
+ lambda: f"int8_linear_matmul only supports B^T @ A. Inner dimensions do not match: B @ A = {shapeB} @ {shapeA}",
+ )
+
+ # cuBLASLt does not support int8 matmul with inner dimensions that are not divisible by 4.
+ # We'll fall back to a slower fp32 calculation in this circumstance.
+ # Fortunately, this should not be very common.
+ if lda % 4 != 0:
+ result = torch.matmul(B.float(), A.float().t()).to(torch.int32)
+ return out.copy_(result)
- if device_type == 'cuda':
- torch._check(blocksize in [4096, 2048, 1024, 512, 256, 128, 64])
- elif device_type == 'hip' or HIP_ENVIRONMENT:
- torch._check(blocksize in [4096, 2048, 1024, 512, 256, 128])
+ with _cuda_device_of(A):
+ ctx = CUBLAS_Context.get_instance().get_context(A.device)
+ ptrA = get_ptr(A)
+ ptrB = get_ptr(B)
+ ptrC = get_ptr(out)
+ ptrRowScale = None
+ m = ct.c_int32(m)
+ n = ct.c_int32(n)
+ k = ct.c_int32(k)
+ lda = ct.c_int32(lda)
+ ldb = ct.c_int32(ldb)
+ ldc = ct.c_int32(ldc)
+ stream = _get_tensor_stream(A)
+
+ has_error = lib.cigemmlt_32(ctx, m, n, k, ptrA, ptrB, ptrC, ptrRowScale, lda, ldb, ldc, stream)
+
+ if has_error:
+ if has_error == 100:
+ # `ERR_NOT_IMPLEMENTED` is defined as 100 in `ops.cu`
+ # TODO: Warn and implement a fallback to fp32 compute?
+ raise NotImplementedError("int8_linear_matmul not implemented!")
+ else:
+ raise RuntimeError(
+ f"cublasLt ran into an error!\n\t{shapeA=}, {shapeB=}, {shapeC=}\n\t{(lda, ldb, ldc)=}\n\t{(m, n, k)=}"
+ )
+
+ return out
+
+
+@register_kernel("bitsandbytes::int8_mm_dequant", "cuda")
+def _(
+ A: torch.Tensor,
+ row_stats: torch.Tensor,
+ col_stats: torch.Tensor,
+ dtype: Optional[torch.dtype] = None,
+ bias: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+ torch._check(A.dtype == torch.int32, lambda: f"A must be int32, got {A.dtype}")
+ torch._check(row_stats.dtype == torch.float32, lambda: f"row_stats must be float32, got {row_stats.dtype}")
+ torch._check(col_stats.dtype == torch.float32, lambda: f"col_stats must be float32, got {col_stats.dtype}")
+
+ # 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)
+
+ ptrA = get_ptr(A)
+ ptrOut = get_ptr(out)
+ ptrRowStats = get_ptr(row_stats)
+ ptrColStats = get_ptr(col_stats)
+ numRows = ct.c_int32(prod(A.shape[:-1]))
+ numCols = ct.c_int32(A.shape[-1])
+
+ # Note: fused bias in the kernel is only supported for fp16
+ # TODO(matthewdouglas): Consider supporting bf16 fused bias
+ ptrBias = get_ptr(bias) if bias is not None and bias.dtype == torch.float16 else None
+
+ with _cuda_device_of(A):
+ lib.cdequant_mm_int32_fp16(
+ ptrA, ptrRowStats, ptrColStats, ptrOut, ptrBias, numRows, numCols, _get_tensor_stream(A)
+ )
+
+ # Add bias separately if not fused in kernel
+ if bias is not None and bias.dtype != torch.float16:
+ out.add_(bias)
+
+ return out.to(dtype or torch.float16)
+
+
+@register_kernel("bitsandbytes::int8_vectorwise_quant", "cuda")
+def _(A: torch.Tensor, threshold=0.0):
+ torch._check(A.dtype == torch.float16, lambda: f"A must be float16, got {A.dtype}")
+ torch._check(threshold >= 0.0, lambda: "threshold must be non-negative")
+ rows = prod(A.shape[:-1])
+ cols = A.shape[-1]
+
+ row_stats = torch.empty(rows, device=A.device, dtype=torch.float32)
+ out_row = torch.empty(A.shape, device=A.device, dtype=torch.int8)
+
+ outlier_cols = None
+
+ if threshold > 0.0:
+ # TODO we could improve perf of this
+ outliers = A.abs() >= threshold
+
+ if outliers.any():
+ outlier_cols = torch.argwhere(outliers.any(dim=0)).view(-1)
+ else:
+ # Needed for torch.compile support.
+ outlier_cols = torch.empty(0, device=A.device, dtype=torch.int64)
+
+ with _cuda_device_of(A):
+ lib.cint8_vector_quant(
+ get_ptr(A),
+ get_ptr(out_row),
+ get_ptr(row_stats),
+ ct.c_float(threshold),
+ ct.c_int32(rows),
+ ct.c_int32(cols),
+ _get_tensor_stream(A),
+ )
+
+ # Zero out values from outlier columns across all rows.
+ # The kernel will handle this for outliers themselves, so we can optimize for rows=1.
+ if rows > 1 and outlier_cols is not None:
+ out_row[:, outlier_cols] = 0
+
+ return out_row, row_stats, outlier_cols
+
+
+@register_kernel("bitsandbytes::int8_double_quant", "cuda")
+def _(
+ A: torch.Tensor,
+ threshold=0.0,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, Optional[torch.Tensor]]:
+ # Use CUDA kernel for rowwise and COO tensor
+ quant_row, row_stats, outlier_cols = torch.ops.bitsandbytes.int8_vectorwise_quant.default(
+ A,
+ threshold=threshold,
+ )
+
+ # PyTorch impl for colwise
+ col_stats, outlier_mask = _get_col_absmax(A, threshold=threshold)
+ if threshold > 0.0 and outlier_mask is not None:
+ A = A.masked_fill(outlier_mask, 0.0)
+ quant_col = torch.round(A.mul(127.0) / col_stats.unsqueeze(0)).to(torch.int8)
+
+ return quant_row, quant_col, row_stats, col_stats.flatten().float(), outlier_cols
+
+
+def _get_col_absmax(
+ A: torch.Tensor,
+ threshold=0.0,
+) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
+ torch._check(A.is_floating_point())
+
+ outlier_mask = None
+
+ absA = A.abs().view(-1, A.shape[-1])
+
+ if threshold > 0.0:
+ # Filter outliers from stats when enabled
+ outlier_mask = absA >= threshold
+ absA.masked_fill_(outlier_mask, 0.0)
+
+ # shape [cols]; unsqueeze(0) gives [1,cols]
+ col_stats = absA.amax(dim=0, keepdim=False).float()
+
+ return col_stats, outlier_mask
+
+
+@register_kernel("bitsandbytes::quantize_blockwise", "cuda")
+def _(A: torch.Tensor, code: torch.Tensor, blocksize: int) -> tuple[torch.Tensor, torch.Tensor]:
+ torch._check_is_size(blocksize)
+
+ if HIP_ENVIRONMENT:
+ torch._check(blocksize in [4096, 2048, 1024, 512, 256, 128])
+ else:
+ torch._check(blocksize in [4096, 2048, 1024, 512, 256, 128, 64])
+
torch._check(code.dtype == torch.float32, lambda: f"code must be float32, got {code.dtype}")
n = A.numel()
blocks = -(n // -blocksize)
absmax = torch.empty((blocks,), device=A.device, dtype=torch.float32)
out = torch.empty_like(A, dtype=torch.uint8)
-
- if device_type == 'cuda' or (device_type == 'hip' or HIP_ENVIRONMENT):
- with _cuda_device_of(A):
- args = (
- get_ptr(code),
- get_ptr(A),
- get_ptr(absmax),
- get_ptr(out),
- ct.c_int32(blocksize),
- ct.c_int(A.numel()),
- )
- if A.dtype == torch.float16:
- lib.cquantize_blockwise_fp16(*args)
- elif A.dtype == torch.bfloat16:
- lib.cquantize_blockwise_bf16(*args)
- elif A.dtype == torch.float32:
- lib.cquantize_blockwise_fp32(*args)
- else:
- raise ValueError(f"Blockwise quantization only supports 16/32-bit floats, but got {A.dtype}")
-
- return out, absmax
-
-
-@register_kernel("bitsandbytes::dequantize_blockwise", "cuda")
-def _(A: torch.Tensor, absmax: torch.Tensor, code: torch.Tensor, blocksize: int, dtype: torch.dtype) -> torch.Tensor:
- out = torch.empty_like(A, dtype=dtype)
- _dequantize_blockwise_impl(A, absmax, code, blocksize, dtype, out=out)
- return out
-
-
-@register_kernel("bitsandbytes::dequantize_blockwise.out", "cuda")
-def _(
- A: torch.Tensor,
- absmax: torch.Tensor,
- code: torch.Tensor,
- blocksize: int,
- dtype: torch.dtype,
- out: torch.Tensor,
-) -> None:
- torch._check(out.dtype == dtype, lambda: f"Expected out.dtype == {dtype}, got {out.dtype}")
- torch._check(out.shape == A.shape, lambda: f"Expected out.shape == {A.shape}, got {out.shape}")
- _dequantize_blockwise_impl(A, absmax, code, blocksize, dtype, out=out)
-
-
-def _dequantize_blockwise_impl(
- A: torch.Tensor, absmax: torch.Tensor, code: torch.Tensor, blocksize: int, dtype: torch.dtype, out: torch.Tensor
-) -> None:
-
- device = A.device
- device_type = device.type
-
- if device_type == 'cuda':
- torch._check(blocksize in [4096, 2048, 1024, 512, 256, 128, 64])
- elif device_type == 'hip' or HIP_ENVIRONMENT:
+ with _cuda_device_of(A):
+ args = (
+ get_ptr(code),
+ get_ptr(A),
+ get_ptr(absmax),
+ get_ptr(out),
+ ct.c_int32(blocksize),
+ ct.c_int(A.numel()),
+ )
+
+ if A.dtype == torch.float16:
+ lib.cquantize_blockwise_fp16(*args)
+ elif A.dtype == torch.bfloat16:
+ lib.cquantize_blockwise_bf16(*args)
+ elif A.dtype == torch.float32:
+ lib.cquantize_blockwise_fp32(*args)
+ else:
+ raise ValueError(f"Blockwise quantization only supports 16/32-bit floats, but got {A.dtype}")
+
+ return out, absmax
+
+
+@register_kernel("bitsandbytes::dequantize_blockwise", "cuda")
+def _(A: torch.Tensor, absmax: torch.Tensor, code: torch.Tensor, blocksize: int, dtype: torch.dtype) -> torch.Tensor:
+ out = torch.empty_like(A, dtype=dtype)
+ _dequantize_blockwise_impl(A, absmax, code, blocksize, dtype, out=out)
+ return out
+
+
+@register_kernel("bitsandbytes::dequantize_blockwise.out", "cuda")
+def _(
+ A: torch.Tensor,
+ absmax: torch.Tensor,
+ code: torch.Tensor,
+ blocksize: int,
+ dtype: torch.dtype,
+ out: torch.Tensor,
+) -> None:
+ torch._check(out.dtype == dtype, lambda: f"Expected out.dtype == {dtype}, got {out.dtype}")
+ torch._check(out.shape == A.shape, lambda: f"Expected out.shape == {A.shape}, got {out.shape}")
+ _dequantize_blockwise_impl(A, absmax, code, blocksize, dtype, out=out)
+
+
+def _dequantize_blockwise_impl(
+ A: torch.Tensor, absmax: torch.Tensor, code: torch.Tensor, blocksize: int, dtype: torch.dtype, out: torch.Tensor
+) -> None:
+ if HIP_ENVIRONMENT:
torch._check(blocksize in [4096, 2048, 1024, 512, 256, 128])
+ else:
+ torch._check(blocksize in [4096, 2048, 1024, 512, 256, 128, 64])
torch._check(A.dtype == torch.uint8, lambda: f"A must be uint8, got {A.dtype}")
torch._check(
dtype in [torch.float16, torch.bfloat16, torch.float32],
lambda: f"Blockwise dequantization only supports 16bit/32bit floating types, got {dtype}",
)
-
- if device_type == 'cuda' or (device_type == 'hip' or HIP_ENVIRONMENT):
- with _cuda_device_of(A):
- args = (
- get_ptr(code),
- get_ptr(A),
- get_ptr(absmax),
- get_ptr(out),
- ct.c_int(blocksize),
- ct.c_int(A.numel()),
- _get_tensor_stream(A),
- )
- if dtype == torch.float16:
- lib.cdequantize_blockwise_fp16(*args)
- elif dtype == torch.bfloat16:
- lib.cdequantize_blockwise_bf16(*args)
- elif dtype == torch.float32:
- lib.cdequantize_blockwise_fp32(*args)
-
-
-@register_kernel("bitsandbytes::quantize_4bit", "cuda")
-def _(
- A: torch.Tensor, blocksize: int, quant_type: str, quant_storage: torch.dtype
-) -> tuple[torch.Tensor, torch.Tensor]:
-
- device = A.device
- device_type = device.type
-
- if device_type == 'cuda':
- torch._check(blocksize in [4096, 2048, 1024, 512, 256, 128, 64])
- elif device_type == 'hip' or HIP_ENVIRONMENT:
+ with _cuda_device_of(A):
+ args = (
+ get_ptr(code),
+ get_ptr(A),
+ get_ptr(absmax),
+ get_ptr(out),
+ ct.c_int(blocksize),
+ ct.c_int(A.numel()),
+ _get_tensor_stream(A),
+ )
+
+ if dtype == torch.float16:
+ lib.cdequantize_blockwise_fp16(*args)
+ elif dtype == torch.bfloat16:
+ lib.cdequantize_blockwise_bf16(*args)
+ elif dtype == torch.float32:
+ lib.cdequantize_blockwise_fp32(*args)
+
+
+@register_kernel("bitsandbytes::quantize_4bit", "cuda")
+def _(
+ A: torch.Tensor, blocksize: int, quant_type: str, quant_storage: torch.dtype
+) -> tuple[torch.Tensor, torch.Tensor]:
+ if HIP_ENVIRONMENT:
torch._check(blocksize in [4096, 2048, 1024, 512, 256, 128])
+ else:
+ torch._check(blocksize in [4096, 2048, 1024, 512, 256, 128, 64])
torch._check(quant_type in ["fp4", "nf4"])
torch._check(
@@ -331,66 +318,65 @@ def _(
blocks = -(n // -blocksize)
absmax = torch.empty((blocks,), device=A.device, dtype=torch.float32)
out = torch.empty(((n + 1) // (quant_storage.itemsize * 2), 1), device=A.device, dtype=quant_storage)
-
- if device_type == 'cuda' or (device_type == 'hip' or HIP_ENVIRONMENT):
- with _cuda_device_of(A):
- args = (
- None,
- get_ptr(A),
- get_ptr(absmax),
- get_ptr(out),
- ct.c_int32(blocksize),
- ct.c_int(n),
- )
- if A.dtype == torch.bfloat16:
- if quant_type == "fp4":
- lib.cquantize_blockwise_bf16_fp4(*args)
- else:
- lib.cquantize_blockwise_bf16_nf4(*args)
- elif A.dtype == torch.float16:
- if quant_type == "fp4":
- lib.cquantize_blockwise_fp16_fp4(*args)
- else:
- lib.cquantize_blockwise_fp16_nf4(*args)
- elif A.dtype == torch.float32:
- if quant_type == "fp4":
- lib.cquantize_blockwise_fp32_fp4(*args)
- else:
- lib.cquantize_blockwise_fp32_nf4(*args)
+ with _cuda_device_of(A):
+ args = (
+ None,
+ get_ptr(A),
+ get_ptr(absmax),
+ get_ptr(out),
+ ct.c_int32(blocksize),
+ ct.c_int(n),
+ )
+
+ if A.dtype == torch.bfloat16:
+ if quant_type == "fp4":
+ lib.cquantize_blockwise_bf16_fp4(*args)
+ else:
+ lib.cquantize_blockwise_bf16_nf4(*args)
+ elif A.dtype == torch.float16:
+ if quant_type == "fp4":
+ lib.cquantize_blockwise_fp16_fp4(*args)
+ else:
+ lib.cquantize_blockwise_fp16_nf4(*args)
+ elif A.dtype == torch.float32:
+ if quant_type == "fp4":
+ lib.cquantize_blockwise_fp32_fp4(*args)
+ else:
+ lib.cquantize_blockwise_fp32_nf4(*args)
return out, absmax
-
-
-@register_kernel("bitsandbytes::dequantize_4bit", "cuda")
-def _(
- A: torch.Tensor,
- absmax: torch.Tensor,
- blocksize: int,
- quant_type: str,
- shape: Sequence[int],
- dtype: torch.dtype,
-) -> torch.Tensor:
- out = torch.empty(shape, dtype=dtype, device=A.device)
- _dequantize_4bit_impl(A, absmax, blocksize, quant_type, dtype, out=out)
- return out
-
-
-@register_kernel("bitsandbytes::dequantize_4bit.out", "cuda")
-def _(
- A: torch.Tensor,
- absmax: torch.Tensor,
- blocksize: int,
- quant_type: str,
- shape: Sequence[int],
- dtype: torch.dtype,
- out: torch.Tensor,
-) -> None:
- torch._check(out.shape == shape, lambda: f"Expected out.shape == {shape}, got {out.shape}")
- torch._check(out.dtype == dtype, lambda: f"Expected out.dtype == {dtype}, got {out.dtype}")
- _dequantize_4bit_impl(A, absmax, blocksize, quant_type, dtype, out=out)
-
-
+
+
+@register_kernel("bitsandbytes::dequantize_4bit", "cuda")
+def _(
+ A: torch.Tensor,
+ absmax: torch.Tensor,
+ blocksize: int,
+ quant_type: str,
+ shape: Sequence[int],
+ dtype: torch.dtype,
+) -> torch.Tensor:
+ out = torch.empty(shape, dtype=dtype, device=A.device)
+ _dequantize_4bit_impl(A, absmax, blocksize, quant_type, dtype, out=out)
+ return out
+
+
+@register_kernel("bitsandbytes::dequantize_4bit.out", "cuda")
+def _(
+ A: torch.Tensor,
+ absmax: torch.Tensor,
+ blocksize: int,
+ quant_type: str,
+ shape: Sequence[int],
+ dtype: torch.dtype,
+ out: torch.Tensor,
+) -> None:
+ torch._check(out.shape == shape, lambda: f"Expected out.shape == {shape}, got {out.shape}")
+ torch._check(out.dtype == dtype, lambda: f"Expected out.dtype == {dtype}, got {out.dtype}")
+ _dequantize_4bit_impl(A, absmax, blocksize, quant_type, dtype, out=out)
+
+
def _dequantize_4bit_impl(
A: torch.Tensor,
absmax: torch.Tensor,
@@ -399,157 +385,154 @@ def _dequantize_4bit_impl(
dtype: torch.dtype,
out: torch.Tensor,
) -> None:
- device = A.device
- device_type = device.type
-
- if device_type == 'cuda':
- torch._check(blocksize in [4096, 2048, 1024, 512, 256, 128, 64])
- elif device_type == 'hip' or HIP_ENVIRONMENT:
+ if HIP_ENVIRONMENT:
torch._check(blocksize in [4096, 2048, 1024, 512, 256, 128])
+ else:
+ torch._check(blocksize in [4096, 2048, 1024, 512, 256, 128, 64])
torch._check(quant_type in ["fp4", "nf4"])
torch._check(
dtype in [torch.bfloat16, torch.float16, torch.float32],
lambda: f"Blockwise 4bit dequantization only supports 16/32-bit floats, but got {dtype}",
)
+
+ with _cuda_device_of(A):
+ args = (
+ None,
+ get_ptr(A),
+ get_ptr(absmax),
+ get_ptr(out),
+ ct.c_int(blocksize),
+ ct.c_int(out.numel()),
+ _get_tensor_stream(A),
+ )
+
+ if out.dtype == torch.bfloat16:
+ if quant_type == "fp4":
+ lib.cdequantize_blockwise_bf16_fp4(*args)
+ else:
+ lib.cdequantize_blockwise_bf16_nf4(*args)
+ elif out.dtype == torch.float16:
+ if quant_type == "fp4":
+ lib.cdequantize_blockwise_fp16_fp4(*args)
+ else:
+ lib.cdequantize_blockwise_fp16_nf4(*args)
+ elif out.dtype == torch.float32:
+ if quant_type == "fp4":
+ lib.cdequantize_blockwise_fp32_fp4(*args)
+ else:
+ lib.cdequantize_blockwise_fp32_nf4(*args)
+
+
+@register_kernel("bitsandbytes::gemv_4bit", "cuda")
+def _(
+ A: torch.Tensor, B: torch.Tensor, shapeB: Sequence[int], absmax: torch.Tensor, code: torch.Tensor, blocksize: int
+) -> torch.Tensor:
+ shape = (*A.shape[:-1], shapeB[0])
+ out = torch.empty(shape, device=A.device, dtype=A.dtype)
+ _gemv_4bit_impl(A, B, shapeB, absmax, code, blocksize, out=out)
+ return out
+
+
+@register_kernel("bitsandbytes::gemv_4bit.out", "cuda")
+def _(
+ A: torch.Tensor,
+ B: torch.Tensor,
+ shapeB: Sequence[int],
+ absmax: torch.Tensor,
+ code: torch.Tensor,
+ blocksize: int,
+ out: torch.Tensor,
+) -> None:
+ torch._check(
+ out.shape == (*A.shape[:-1], shapeB[0]),
+ lambda: f"Expected out.shape == {(*A.shape[:-1], shapeB[0])}, got {out.shape}",
+ )
+ torch._check(out.dtype == A.dtype, lambda: f"Expected out.dtype == {A.dtype}, got {out.dtype}")
+ _gemv_4bit_impl(A, B, shapeB, absmax, code, blocksize, out=out)
+
+
+def _gemv_4bit_impl(
+ A: torch.Tensor,
+ B: torch.Tensor,
+ shapeB: Sequence[int],
+ absmax: torch.Tensor,
+ code: torch.Tensor,
+ blocksize: int,
+ out: torch.Tensor,
+) -> None:
+ torch._check_is_size(blocksize)
+ torch._check(
+ A.numel() == A.size(-1),
+ lambda: f"A must be a vector with leading dimensions of 1, got {A.shape}",
+ )
+ torch._check(
+ A.dtype in [torch.float16, torch.bfloat16, torch.float32],
+ lambda: f"A must be float16, bfloat16, or float32, got {A.dtype}",
+ )
+ torch._check(
+ B.dtype in [torch.uint8, torch.bfloat16, torch.float16, torch.float32],
+ lambda: f"B must be backed by storage of type uint8, bfloat16, float16, or float32, got {B.dtype}",
+ )
+ torch._check(absmax.dtype == torch.float32, lambda: f"absmax must be float32, got {absmax.dtype}")
+ torch._check(code.dtype == torch.float32, lambda: f"code must be float32, got {code.dtype}")
+
+ m = ct.c_int32(shapeB[0])
+ n = ct.c_int32(1)
+ k = ct.c_int32(shapeB[1])
+
+ lda = m
+ ldb = ct.c_int32((A.shape[-1] + 1) // 2)
+ ldc = m
- if device_type == 'cuda' or (device_type == 'hip' or HIP_ENVIRONMENT):
- with _cuda_device_of(A):
- args = (
- None,
+ stream = _get_tensor_stream(A)
+
+ with _cuda_device_of(A):
+ if A.dtype == torch.float16:
+ lib.cgemm_4bit_inference_naive_fp16(
+ m,
+ n,
+ k,
get_ptr(A),
+ get_ptr(B),
get_ptr(absmax),
+ get_ptr(code),
get_ptr(out),
- ct.c_int(blocksize),
- ct.c_int(out.numel()),
- _get_tensor_stream(A),
+ lda,
+ ldb,
+ ldc,
+ ct.c_int32(blocksize),
+ stream,
+ )
+ elif A.dtype == torch.bfloat16:
+ lib.cgemm_4bit_inference_naive_bf16(
+ m,
+ n,
+ k,
+ get_ptr(A),
+ get_ptr(B),
+ get_ptr(absmax),
+ get_ptr(code),
+ get_ptr(out),
+ lda,
+ ldb,
+ ldc,
+ ct.c_int32(blocksize),
+ stream,
+ )
+ elif A.dtype == torch.float32:
+ lib.cgemm_4bit_inference_naive_fp32(
+ m,
+ n,
+ k,
+ get_ptr(A),
+ get_ptr(B),
+ get_ptr(absmax),
+ get_ptr(code),
+ get_ptr(out),
+ lda,
+ ldb,
+ ldc,
+ ct.c_int32(blocksize),
+ stream,
)
-
- if out.dtype == torch.bfloat16:
- if quant_type == "fp4":
- lib.cdequantize_blockwise_bf16_fp4(*args)
- else:
- lib.cdequantize_blockwise_bf16_nf4(*args)
- elif out.dtype == torch.float16:
- if quant_type == "fp4":
- lib.cdequantize_blockwise_fp16_fp4(*args)
- else:
- lib.cdequantize_blockwise_fp16_nf4(*args)
- elif out.dtype == torch.float32:
- if quant_type == "fp4":
- lib.cdequantize_blockwise_fp32_fp4(*args)
- else:
- lib.cdequantize_blockwise_fp32_nf4(*args)
-
-
-@register_kernel("bitsandbytes::gemv_4bit", "cuda")
-def _(
- A: torch.Tensor, B: torch.Tensor, shapeB: Sequence[int], absmax: torch.Tensor, code: torch.Tensor, blocksize: int
-) -> torch.Tensor:
- shape = (*A.shape[:-1], shapeB[0])
- out = torch.empty(shape, device=A.device, dtype=A.dtype)
- _gemv_4bit_impl(A, B, shapeB, absmax, code, blocksize, out=out)
- return out
-
-
-@register_kernel("bitsandbytes::gemv_4bit.out", "cuda")
-def _(
- A: torch.Tensor,
- B: torch.Tensor,
- shapeB: Sequence[int],
- absmax: torch.Tensor,
- code: torch.Tensor,
- blocksize: int,
- out: torch.Tensor,
-) -> None:
- torch._check(
- out.shape == (*A.shape[:-1], shapeB[0]),
- lambda: f"Expected out.shape == {(*A.shape[:-1], shapeB[0])}, got {out.shape}",
- )
- torch._check(out.dtype == A.dtype, lambda: f"Expected out.dtype == {A.dtype}, got {out.dtype}")
- _gemv_4bit_impl(A, B, shapeB, absmax, code, blocksize, out=out)
-
-def _gemv_4bit_impl(
- A: torch.Tensor,
- B: torch.Tensor,
- shapeB: Sequence[int],
- absmax: torch.Tensor,
- code: torch.Tensor,
- blocksize: int,
- out: torch.Tensor,
-) -> None:
- torch._check_is_size(blocksize)
- torch._check(
- A.numel() == A.size(-1),
- lambda: f"A must be a vector with leading dimensions of 1, got {A.shape}",
- )
- torch._check(
- A.dtype in [torch.float16, torch.bfloat16, torch.float32],
- lambda: f"A must be float16, bfloat16, or float32, got {A.dtype}",
- )
- torch._check(
- B.dtype in [torch.uint8, torch.bfloat16, torch.float16, torch.float32],
- lambda: f"B must be backed by storage of type uint8, bfloat16, float16, or float32, got {B.dtype}",
- )
- torch._check(absmax.dtype == torch.float32, lambda: f"absmax must be float32, got {absmax.dtype}")
- torch._check(code.dtype == torch.float32, lambda: f"code must be float32, got {code.dtype}")
-
- m = ct.c_int32(shapeB[0])
- n = ct.c_int32(1)
- k = ct.c_int32(shapeB[1])
-
- lda = m
- ldb = ct.c_int32((A.shape[-1] + 1) // 2)
- ldc = m
-
- stream = _get_tensor_stream(A)
-
- with _cuda_device_of(A):
- if A.dtype == torch.float16:
- lib.cgemm_4bit_inference_naive_fp16(
- m,
- n,
- k,
- get_ptr(A),
- get_ptr(B),
- get_ptr(absmax),
- get_ptr(code),
- get_ptr(out),
- lda,
- ldb,
- ldc,
- ct.c_int32(blocksize),
- stream,
- )
- elif A.dtype == torch.bfloat16:
- lib.cgemm_4bit_inference_naive_bf16(
- m,
- n,
- k,
- get_ptr(A),
- get_ptr(B),
- get_ptr(absmax),
- get_ptr(code),
- get_ptr(out),
- lda,
- ldb,
- ldc,
- ct.c_int32(blocksize),
- stream,
- )
- elif A.dtype == torch.float32:
- lib.cgemm_4bit_inference_naive_fp32(
- m,
- n,
- k,
- get_ptr(A),
- get_ptr(B),
- get_ptr(absmax),
- get_ptr(code),
- get_ptr(out),
- lda,
- ldb,
- ldc,
- ct.c_int32(blocksize),
- stream,
- )
From 9fe67efada457a759d1d8193265243209e784e2c Mon Sep 17 00:00:00 2001
From: MISHANMAURYA <118961433+MISHANMAURYA@users.noreply.github.com>
Date: Fri, 23 May 2025 02:11:31 +0530
Subject: [PATCH 22/85] Update functional.py
---
bitsandbytes/functional.py | 28 +++++++++++++---------------
1 file changed, 13 insertions(+), 15 deletions(-)
diff --git a/bitsandbytes/functional.py b/bitsandbytes/functional.py
index 3f0c1ff94..237aa3e54 100644
--- a/bitsandbytes/functional.py
+++ b/bitsandbytes/functional.py
@@ -754,13 +754,11 @@ def quantize_blockwise(
if "dynamic" not in name2qmap:
name2qmap["dynamic"] = create_dynamic_map().to(A.device)
code = name2qmap["dynamic"]
-
- device_type = A.device.type
- if device_type in ["cuda", "hip"]:
- if not HIP_ENVIRONMENT:
- assert blocksize in [4096, 2048, 1024, 512, 256, 128, 64]
- else:
- assert blocksize in [4096, 2048, 1024, 512, 256, 128]
+
+ if HIP_ENVIRONMENT:
+ assert blocksize in [4096, 2048, 1024, 512, 256, 128]
+ else:
+ assert blocksize in [4096, 2048, 1024, 512, 256, 128, 64]
_out, _absmax = torch.ops.bitsandbytes.quantize_blockwise.default(
A,
@@ -839,15 +837,15 @@ def dequantize_blockwise(
if quant_state is None:
quant_state = QuantState(absmax=absmax, code=code, blocksize=blocksize, dtype=torch.float32)
- device_type = A.device.type
- if device_type in ["cuda", "hip"]:
+ if HIP_ENVIRONMENT:
+ supported_blocksizes = [4096, 2048, 1024, 512, 256, 128]
+ else:
supported_blocksizes = [4096, 2048, 1024, 512, 256, 128, 64]
- if HIP_ENVIRONMENT:
- supported_blocksizes = supported_blocksizes[:-1]
- if quant_state.blocksize not in supported_blocksizes:
- raise ValueError(
- f"The blocksize of {quant_state.blocksize} is not supported. Supported values: {supported_blocksizes}"
- )
+
+ if quant_state.blocksize not in supported_blocksizes:
+ raise ValueError(
+ f"The blocksize of {quant_state.blocksize} is not supported. Supported values: {supported_blocksizes}"
+ )
absmax = quant_state.absmax
if quant_state.nested:
From d97fdce654129ca156f0cb47555529d4f4941778 Mon Sep 17 00:00:00 2001
From: MISHANMAURYA <118961433+MISHANMAURYA@users.noreply.github.com>
Date: Fri, 23 May 2025 02:18:37 +0530
Subject: [PATCH 23/85] Update functional.py
---
bitsandbytes/functional.py | 4 ++--
1 file changed, 2 insertions(+), 2 deletions(-)
diff --git a/bitsandbytes/functional.py b/bitsandbytes/functional.py
index 237aa3e54..1cee234ea 100644
--- a/bitsandbytes/functional.py
+++ b/bitsandbytes/functional.py
@@ -858,8 +858,8 @@ def dequantize_blockwise(
torch.ops.bitsandbytes.dequantize_blockwise.out(
A,
absmax,
- quant_state.code.to(A.device),
- quant_state.blocksize,
+ code.to(A.device),
+ blocksize,
quant_state.dtype,
out=out,
)
From e99ac0a1fb4336a53250c6d74e4908de0883fb52 Mon Sep 17 00:00:00 2001
From: Matthew Douglas <38992547+matthewdouglas@users.noreply.github.com>
Date: Fri, 23 May 2025 16:05:21 -0400
Subject: [PATCH 24/85] Optimizer backwards compatibility (#1647)
---
bitsandbytes/optim/optimizer.py | 10 +++++-----
1 file changed, 5 insertions(+), 5 deletions(-)
diff --git a/bitsandbytes/optim/optimizer.py b/bitsandbytes/optim/optimizer.py
index 4bed9a7c3..9c20f9376 100644
--- a/bitsandbytes/optim/optimizer.py
+++ b/bitsandbytes/optim/optimizer.py
@@ -303,9 +303,9 @@ def get_config(self, gindex, pindex, group):
config["eps"] = group["eps"]
config["weight_decay"] = group["weight_decay"]
config["lr"] = group["lr"]
- config["alpha"] = group.get("alpha")
- config["t_alpha"] = group.get("t_alpha")
- config["t_beta3"] = group.get("t_beta3")
+ config["alpha"] = group.get("alpha", 0.0)
+ config["t_alpha"] = group.get("t_alpha", 0)
+ config["t_beta3"] = group.get("t_beta3", 0)
config["optim_bits"] = self.args.optim_bits
config["min_8bit_size"] = self.args.min_8bit_size
config["percentile_clipping"] = self.args.percentile_clipping
@@ -530,7 +530,7 @@ def update_step(self, group, p, gindex, pindex):
state["state2"],
config["betas"][1],
config["betas"][2] if len(config["betas"]) >= 3 else 0.0,
- config["alpha"],
+ config.get("alpha", 0.0),
config["weight_decay"],
gnorm_scale,
state["unorm_vec"] if config["max_unorm"] > 0.0 else None,
@@ -575,7 +575,7 @@ def update_step(self, group, p, gindex, pindex):
config["betas"][0],
config["betas"][1],
config["betas"][2] if len(config["betas"]) >= 3 else 0.0,
- config["alpha"],
+ config.get("alpha", 0.0),
config["eps"],
step,
config["lr"],
From 503d243e047d64bae3b0cd3e713cf28530e9af43 Mon Sep 17 00:00:00 2001
From: Matthew Douglas <38992547+matthewdouglas@users.noreply.github.com>
Date: Fri, 23 May 2025 20:19:09 -0400
Subject: [PATCH 25/85] General cleanup & test improvements (#1646)
* General cleanup & test improvements
* Tests: WA numpy 2 compat issue for torch<2.3
* Tests: update aarch64 cpu min torch version
* Tests: update aarch64 cpu min torch version
* Tests: update aarch64 cpu min torch version
---
.github/workflows/tests.yml | 28 ++-
benchmarking/int8/row_scale_benchmark.py | 70 -------
deploy.sh | 237 -----------------------
environment-bnb.yml | 21 --
environment.yml | 46 -----
tests/test_autograd.py | 7 +
tests/test_functional.py | 8 +-
tests/test_ops.py | 33 ++--
tests/test_triton.py | 2 +-
9 files changed, 57 insertions(+), 395 deletions(-)
delete mode 100644 benchmarking/int8/row_scale_benchmark.py
delete mode 100644 deploy.sh
delete mode 100644 environment-bnb.yml
delete mode 100644 environment.yml
diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
index 5d2a2708b..f1a5dca69 100644
--- a/.github/workflows/tests.yml
+++ b/.github/workflows/tests.yml
@@ -93,24 +93,32 @@ jobs:
path: output/${{ matrix.os }}/${{ matrix.arch }}/*
retention-days: 7
- cpu-tests:
+ test-cpu:
if: github.repository == 'bitsandbytes-foundation/bitsandbytes'
needs: build-cpu
strategy:
fail-fast: false
matrix:
os: [ubuntu-22.04, ubuntu-22.04-arm, windows-2025, macos-15]
- torch_version: ["2.6.0", "2.7.0"]
+ # Test with the oldest supported torch version and the two newest.
+ torch_version: ["2.2.2", "2.6.0", "2.7.0"]
include:
- os: ubuntu-22.04
arch: x86_64
runner: banb-aws-general-8-plus-use1-public-80
- os: ubuntu-22.04-arm
arch: aarch64
+ - os: ubuntu-22.04-arm
+ arch: aarch64
+ torch_version: "2.5.1"
- os: windows-2025
arch: x86_64
- os: macos-15
arch: arm64
+ exclude:
+ - os: ubuntu-22.04-arm
+ torch_version: "2.2.2"
+
runs-on: ${{ matrix.runner || matrix.os }}
env:
BNB_TEST_DEVICE: cpu
@@ -135,6 +143,11 @@ jobs:
pip install -e ".[test]"
pip install pytest-cov
+ # We need to downgrade to numpy<2 for torch<2.3 compatibility.
+ - name: Downgrade NumPy
+ if: startsWith(matrix.torch_version, '2.2.')
+ run: pip install "numpy<2"
+
- name: Show installed packages
run: pip list
@@ -144,7 +157,7 @@ jobs:
- name: Run tests
run: pytest --durations=100
- # cuda-aarch64-tests:
+ # test-cuda-aarch64:
# if: github.repository == 'bitsandbytes-foundation/bitsandbytes'
# needs: build-cuda
# strategy:
@@ -167,7 +180,7 @@ jobs:
- cuda-tests:
+ test-cuda:
if: github.repository == 'bitsandbytes-foundation/bitsandbytes'
needs: build-cuda
strategy:
@@ -179,7 +192,7 @@ jobs:
cuda_version: ["11.8.0", "12.6.3", "12.8.1"]
include:
- cuda_version: "11.8.0"
- torch_version: "2.4.1"
+ torch_version: "2.2.2"
pypi_index: "https://download.pytorch.org/whl/cu118"
- cuda_version: "12.6.3"
torch_version: "2.6.0"
@@ -238,6 +251,11 @@ jobs:
pip install -e ".[test]"
pip install pytest-cov
+ # We need to downgrade to numpy<2 for torch<2.3 compatibility.
+ - name: Downgrade NumPy
+ if: startsWith(matrix.torch_version, '2.2.')
+ run: pip install "numpy<2"
+
- name: Show installed packages
run: pip list
diff --git a/benchmarking/int8/row_scale_benchmark.py b/benchmarking/int8/row_scale_benchmark.py
deleted file mode 100644
index 98d2496de..000000000
--- a/benchmarking/int8/row_scale_benchmark.py
+++ /dev/null
@@ -1,70 +0,0 @@
-"""
-Extracted from tests/test_functional.py
-
-Note: This feature is currently unused! It is kept here for archival purposes.
-
-Usage: pytest benchmarking/int8/row_scale_benchmark.py
-"""
-
-import time
-
-import pytest
-import torch
-
-from bitsandbytes import functional as F
-
-k = 20
-torch.set_printoptions(precision=5, sci_mode=False, linewidth=120, edgeitems=20, threshold=10000)
-
-
-@pytest.mark.parametrize(
- ("dim1", "dim4", "inner"),
- [
- pytest.param(1024, 12288 * 4, 12288, id="1024, 12288*4, 12288"),
- pytest.param(2048, 4096 * 4, 4096, id="2048, 4096*4, 4096"),
- ],
-)
-@pytest.mark.skip("Row scale has some bugs for ampere")
-@pytest.mark.benchmark
-def test_row_scale_bench(dim1, dim4, inner):
- formatB = F.get_special_format_str()
- err1, err2, err3 = [], [], []
- relerr1, relerr2 = [], []
- scale = 1
- A = torch.randn(dim1, inner, device="cuda").half()
- B = torch.randn(dim4, inner, device="cuda").half()
- torch.nn.init.xavier_uniform_(B)
- # warmpup
- for i in range(k):
- C1 = torch.matmul(A, B.t())
-
- torch.cuda.synchronize()
- t0 = time.time()
- for i in range(k):
- C1 = torch.matmul(A, B.t())
- torch.cuda.synchronize()
- print("16", time.time() - t0)
-
- C1a, C1b, stats1a, stats1b, coo_tensor = F.int8_double_quant(A)
- CB, absmaxB = F.vectorwise_quant(B, quant_type="linear")
- A2, SA = F.nvidia_transform(C1a, "col32")
- B2, SB = F.nvidia_transform(CB, formatB)
- A1, maxA = F.vectorwise_quant(A, dim=1)
-
- c = 10.0 * inner * scale
- row_scale = maxA / c
- torch.cuda.synchronize()
- t0 = time.time()
- for i in range(k):
- outC32 = F.int8_linear_matmul(A2, B2, dtype=torch.int8, row_scale=row_scale)
- torch.cuda.synchronize()
- print("row-wise", time.time() - t0)
-
- C2a, C2b, stats2a, stats2b, coo_tensor = F.int8_double_quant(B)
- B2, SB = F.nvidia_transform(C2a, formatB)
- torch.cuda.synchronize()
- t0 = time.time()
- for i in range(k):
- outC32 = F.int8_linear_matmul(A2, B2)
- torch.cuda.synchronize()
- print("vector-wise", time.time() - t0)
diff --git a/deploy.sh b/deploy.sh
deleted file mode 100644
index e60373627..000000000
--- a/deploy.sh
+++ /dev/null
@@ -1,237 +0,0 @@
-#!/bin/bash
-BASE_PATH=$1
-
-echo "MAKE SURE LD_LIBRARY_PATH IS EMPTY!"
-echo $LD_LIBRARY_PATH
-
-if [[ ! -z "${LD_LIBRARY_PATH}" ]]; then
- echo "Compilation unsuccessful!" 1>&2
- exit 64
-fi
-
-
-module unload cuda && echo "no module function available. Probably not on a slurm cluster."
-module unload gcc && echo "no module function available. Probably not on a slurm cluster."
-
-rm -rf dist build
-make cleaneggs
-make cleanlibs
-
-rm -rf build/*
-export CUDA_HOME=
-export CUDA_VERSION=
-make cpuonly CUDA_VERSION="CPU"
-
-if [ ! -f "./bitsandbytes/libbitsandbytes_cpu.so" ]; then
- # Control will enter here if $DIRECTORY doesn't exist.
- echo "Compilation unsuccessful!" 1>&2
- exit 64
-fi
-
-rm -rf build/*
-export CUDA_HOME=$BASE_PATH/cuda-11.0
-make cuda110 CUDA_VERSION=110
-
-if [ ! -f "./bitsandbytes/libbitsandbytes_cuda110.so" ]; then
- # Control will enter here if $DIRECTORY doesn't exist.
- echo "Compilation unsuccessful!" 1>&2
- exit 64
-fi
-
-rm -rf build/*
-export CUDA_HOME=$BASE_PATH/cuda-11.1
-make cuda11x CUDA_VERSION=111
-
-if [ ! -f "./bitsandbytes/libbitsandbytes_cuda111.so" ]; then
- # Control will enter here if $DIRECTORY doesn't exist.
- echo "Compilation unsuccessful!" 1>&2
- exit 64
-fi
-
-rm -rf build/*
-export CUDA_HOME=$BASE_PATH/cuda-11.4
-make cuda11x CUDA_VERSION=114
-
-if [ ! -f "./bitsandbytes/libbitsandbytes_cuda114.so" ]; then
- # Control will enter here if $DIRECTORY doesn't exist.
- echo "Compilation unsuccessful!" 1>&2
- exit 64
-fi
-
-rm -rf build/*
-export CUDA_HOME=$BASE_PATH/cuda-11.5
-make cuda11x CUDA_VERSION=115
-
-if [ ! -f "./bitsandbytes/libbitsandbytes_cuda115.so" ]; then
- # Control will enter here if $DIRECTORY doesn't exist.
- echo "Compilation unsuccessful!" 1>&2
- exit 64
-fi
-
-rm -rf build/*
-export CUDA_HOME=$BASE_PATH/cuda-11.7
-make cuda11x CUDA_VERSION=117
-
-if [ ! -f "./bitsandbytes/libbitsandbytes_cuda117.so" ]; then
- # Control will enter here if $DIRECTORY doesn't exist.
- echo "Compilation unsuccessful!" 1>&2
- exit 64
-fi
-
-rm -rf build/*
-export CUDA_HOME=$BASE_PATH/cuda-11.8
-make cuda118 CUDA_VERSION=118
-
-if [ ! -f "./bitsandbytes/libbitsandbytes_cuda118.so" ]; then
- # Control will enter here if $DIRECTORY doesn't exist.
- echo "Compilation unsuccessful!" 1>&2
- exit 64
-fi
-
-rm -rf build/*
-export CUDA_HOME=$BASE_PATH/cuda-12.0
-make cuda12x CUDA_VERSION=120
-
-if [ ! -f "./bitsandbytes/libbitsandbytes_cuda120.so" ]; then
- # Control will enter here if $DIRECTORY doesn't exist.
- echo "Compilation unsuccessful!" 1>&2
- exit 64
-fi
-
-rm -rf build/*
-export CUDA_HOME=$BASE_PATH/cuda-12.1
-make cuda12x CUDA_VERSION=121
-
-if [ ! -f "./bitsandbytes/libbitsandbytes_cuda121.so" ]; then
- # Control will enter here if $DIRECTORY doesn't exist.
- echo "Compilation unsuccessful!" 1>&2
- exit 64
-fi
-
-rm -rf build/*
-export CUDA_HOME=$BASE_PATH/cuda-12.2
-make cuda12x CUDA_VERSION=122
-
-if [ ! -f "./bitsandbytes/libbitsandbytes_cuda122.so" ]; then
- # Control will enter here if $DIRECTORY doesn't exist.
- echo "Compilation unsuccessful!" 1>&2
- exit 64
-fi
-
-rm -rf build/*
-export CUDA_HOME=$BASE_PATH/cuda-12.3
-make cuda12x CUDA_VERSION=123
-
-if [ ! -f "./bitsandbytes/libbitsandbytes_cuda123.so" ]; then
- # Control will enter here if $DIRECTORY doesn't exist.
- echo "Compilation unsuccessful!" 1>&2
- exit 64
-fi
-
-############################# START NO CUBLASLT #############################################
-# binaries without 8-bit matmul support START HERE
-# ###########################################################################################
-
-rm -rf build/*
-export CUDA_HOME=$BASE_PATH/cuda-11.0
-make cuda110_nomatmul CUDA_VERSION=110
-
-if [ ! -f "./bitsandbytes/libbitsandbytes_cuda110_nocublaslt.so" ]; then
- # Control will enter here if $DIRECTORY doesn't exist.
- echo "Compilation unsuccessful!" 1>&2
- exit 64
-fi
-
-
-rm -rf build/*
-export CUDA_HOME=$BASE_PATH/cuda-11.1
-make cuda11x_nomatmul CUDA_VERSION=111
-
-if [ ! -f "./bitsandbytes/libbitsandbytes_cuda111_nocublaslt.so" ]; then
- # Control will enter here if $DIRECTORY doesn't exist.
- echo "Compilation unsuccessful!" 1>&2
- exit 64
-fi
-
-rm -rf build/*
-export CUDA_HOME=$BASE_PATH/cuda-11.4
-make cuda11x_nomatmul CUDA_VERSION=114
-
-if [ ! -f "./bitsandbytes/libbitsandbytes_cuda114_nocublaslt.so" ]; then
- # Control will enter here if $DIRECTORY doesn't exist.
- echo "Compilation unsuccessful!" 1>&2
- exit 64
-fi
-
-rm -rf build/*
-export CUDA_HOME=$BASE_PATH/cuda-11.5
-make cuda11x_nomatmul CUDA_VERSION=115
-
-if [ ! -f "./bitsandbytes/libbitsandbytes_cuda115_nocublaslt.so" ]; then
- # Control will enter here if $DIRECTORY doesn't exist.
- echo "Compilation unsuccessful!" 1>&2
- exit 64
-fi
-
-rm -rf build/*
-export CUDA_HOME=$BASE_PATH/cuda-11.7
-make cuda11x_nomatmul CUDA_VERSION=117
-
-if [ ! -f "./bitsandbytes/libbitsandbytes_cuda117_nocublaslt.so" ]; then
- # Control will enter here if $DIRECTORY doesn't exist.
- echo "Compilation unsuccessful!" 1>&2
- exit 64
-fi
-
-rm -rf build/*
-export CUDA_HOME=$BASE_PATH/cuda-11.8
-make cuda118_nomatmul CUDA_VERSION=118
-
-if [ ! -f "./bitsandbytes/libbitsandbytes_cuda118_nocublaslt.so" ]; then
- # Control will enter here if $DIRECTORY doesn't exist.
- echo "Compilation unsuccessful!" 1>&2
- exit 64
-fi
-
-rm -rf build/*
-export CUDA_HOME=$BASE_PATH/cuda-12.0
-make cuda12x_nomatmul CUDA_VERSION=120
-
-if [ ! -f "./bitsandbytes/libbitsandbytes_cuda120_nocublaslt.so" ]; then
- # Control will enter here if $DIRECTORY doesn't exist.
- echo "Compilation unsuccessful!" 1>&2
- exit 64
-fi
-
-rm -rf build/*
-export CUDA_HOME=$BASE_PATH/cuda-12.1
-make cuda12x_nomatmul CUDA_VERSION=121
-
-if [ ! -f "./bitsandbytes/libbitsandbytes_cuda121_nocublaslt.so" ]; then
- # Control will enter here if $DIRECTORY doesn't exist.
- echo "Compilation unsuccessful!" 1>&2
- exit 64
-fi
-
-rm -rf build/*
-export CUDA_HOME=$BASE_PATH/cuda-12.2
-make cuda12x_nomatmul CUDA_VERSION=122
-
-if [ ! -f "./bitsandbytes/libbitsandbytes_cuda122_nocublaslt.so" ]; then
- # Control will enter here if $DIRECTORY doesn't exist.
- echo "Compilation unsuccessful!" 1>&2
- exit 64
-fi
-
-rm -rf build/*
-export CUDA_HOME=$BASE_PATH/cuda-12.3
-make cuda12x_nomatmul CUDA_VERSION=123
-
-if [ ! -f "./bitsandbytes/libbitsandbytes_cuda123_nocublaslt.so" ]; then
- # Control will enter here if $DIRECTORY doesn't exist.
- echo "Compilation unsuccessful!" 1>&2
- exit 64
-fi
-
-python -m build
-python -m twine upload dist/* --verbose
diff --git a/environment-bnb.yml b/environment-bnb.yml
deleted file mode 100644
index 1214f7930..000000000
--- a/environment-bnb.yml
+++ /dev/null
@@ -1,21 +0,0 @@
-# for cmake build
-name: bnb
-channels:
- - pytorch
- - nvidia
- - conda-forge
-
-dependencies:
- - python
- #- accelerate
- #- einops
- - scipy
- #- transformers
- - pytest
- - pytest-cases
- - ipython
- - debugpy
- - yapf
- - monkeytype
- - rich
- - pytest-sugar
diff --git a/environment.yml b/environment.yml
deleted file mode 100644
index af421b3c6..000000000
--- a/environment.yml
+++ /dev/null
@@ -1,46 +0,0 @@
-name: bnb
-channels:
- - pytorch
- - nvidia
- - conda-forge
-
-dependencies:
- # Base
- - conda-forge::python=3.8
- - pytorch::pytorch=>2.1
- - pytorch::pytorch-cuda=11.8
- - nvidia::cuda=11.8
- # Libraries
- - conda-forge::accelerate
- - conda-forge::einops
- - conda-forge::scipy
- - conda-forge::transformers
- # Development
- - conda-forge::pytest
- - conda-forge::build # build Python packages
- - conda-forge::twine # upload Python packages
- - conda-forge::pytest-cases # more readable and composable parametrized tests
- - conda-forge::ipython # better interactive shell
- - conda-forge::debugpy # debugger-support for VSCode
- - conda-forge::ruff # linting
- - conda-forge::yapf # code formatting
- - conda-forge::monkeytype # infer type annotations
- - conda-forge::rich # better, colored tracebacks, etc
- - conda-forge::pytest-sugar # better pytest output
- # - conda-forge::nodejs # for `doc-builder preview` (optional)
-
-## ENV CREATION - steps to reproduce:
-# mamba env remove -n bnb
-# mamba create -y -n bnb python=3.8 # creating an empty env bypasses conda
-# # and leads to much faster env resolution in the next step https://github.com/mamba-org/mamba/issues/633#issuecomment-812272143
-# mamba env update -n bnb -f environment.yml
-# mamba activate bnb
-
-## PIP dependencies (install *after* ENV CREATION):
-# pip install --no-cache-dir --no-deps lion_pytorch triton hf-doc-builder watchdog
-## NOTE: conda peft is not up to date, so we install from pip
-# cd pip install -e . ## installs bitsandbytes as editable development install from within repo root dir
-
-## ENV UPDATE:
-# # add new packages to environment.yml, then:
-# mamba env update -n bnb -f environment.yml
diff --git a/tests/test_autograd.py b/tests/test_autograd.py
index b6ba284c9..fc2e7aa6f 100644
--- a/tests/test_autograd.py
+++ b/tests/test_autograd.py
@@ -49,6 +49,10 @@ def test_matmullt(
req_grad = list(req_grad)
req_grad[2] = False
+ if device == "cpu" and dtype != torch.float32 and has_fp16_weights and any(req_grad):
+ if torch.__version__ < (2, 6):
+ pytest.xfail("mse_loss bf16/fp16 on CPU is not supported in torch < 2.6")
+
for i in range(3):
# normal multiply
if funcs[0] in [torch.mm, torch.matmul]:
@@ -185,6 +189,9 @@ def test_matmul_4bit(
req_grad = list(req_grad)
req_grad[2] = False
+ if device == "cpu" and dtype != torch.float32 and any(req_grad) and torch.__version__ < (2, 6):
+ pytest.xfail("mse_loss fp16 on CPU is not supported in torch < 2.6")
+
for i in range(3):
# normal multiply
if funcs[0] in [torch.mm, torch.matmul]:
diff --git a/tests/test_functional.py b/tests/test_functional.py
index 0b9390aaa..8568d45f0 100644
--- a/tests/test_functional.py
+++ b/tests/test_functional.py
@@ -1342,8 +1342,12 @@ def test_gemv_4bit(self, device, dim, dtype, storage_type, quant_storage, double
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32], ids=describe_dtype)
@pytest.mark.parametrize("double_quant", [False], ids=["DQ_True"])
def test_gemv_eye_4bit(self, device, storage_type, dtype, double_quant):
- if device == "cpu" and storage_type != "nf4":
- pytest.xfail("fp4 quantization is not supported on CPU")
+ if device == "cpu":
+ if storage_type != "nf4":
+ pytest.xfail("fp4 quantization is not supported on CPU")
+
+ if dtype == torch.bfloat16 and torch.__version__ < (2, 3):
+ pytest.xfail("eye doe not support bfloat16 on CPU in torch < 2.3")
dims = 10
torch.random.manual_seed(np.random.randint(0, 412424242))
diff --git a/tests/test_ops.py b/tests/test_ops.py
index 4da1663f0..e85bc0ef0 100644
--- a/tests/test_ops.py
+++ b/tests/test_ops.py
@@ -6,6 +6,13 @@
import bitsandbytes
from tests.helpers import TRUE_FALSE, get_available_devices, id_formatter
+# torch.library.opcheck is only available in torch 2.4 and later.
+# When testing with older versions, we will skip it as a no-op.
+if torch.__version__ >= (2, 4):
+ opcheck = torch.library.opcheck
+else:
+ opcheck = lambda *args, **kwargs: None
+
class TestLLMInt8Ops:
@pytest.mark.parametrize("device", get_available_devices())
@@ -18,7 +25,7 @@ def test_int8_linear_matmul(self, device):
assert out.dtype == torch.int32
assert out.device == A.device
- torch.library.opcheck(torch.ops.bitsandbytes.int8_linear_matmul.default, (A, B))
+ opcheck(torch.ops.bitsandbytes.int8_linear_matmul.default, (A, B))
@pytest.mark.parametrize("device", get_available_devices())
def test_int8_linear_matmul_out(self, device):
@@ -32,7 +39,7 @@ def test_int8_linear_matmul_out(self, device):
assert out.dtype == torch.int32
assert out.device == A.device
- torch.library.opcheck(torch.ops.bitsandbytes.int8_linear_matmul.out, (A, B, out))
+ opcheck(torch.ops.bitsandbytes.int8_linear_matmul.out, (A, B, out))
@pytest.mark.parametrize("threshold", [0.0, 6.0])
@pytest.mark.parametrize("device", get_available_devices())
@@ -57,9 +64,8 @@ def test_int8_vectorwise_quant(self, threshold, device):
else:
assert outlier_cols is None
- torch.library.opcheck(torch.ops.bitsandbytes.int8_vectorwise_quant, (A,))
-
- torch.library.opcheck(torch.ops.bitsandbytes.int8_vectorwise_quant, (A, threshold))
+ opcheck(torch.ops.bitsandbytes.int8_vectorwise_quant, (A,))
+ opcheck(torch.ops.bitsandbytes.int8_vectorwise_quant, (A, threshold))
@pytest.mark.parametrize("device", get_available_devices())
def test_int8_mm_dequant(self, device):
@@ -72,7 +78,7 @@ def test_int8_mm_dequant(self, device):
assert out.dtype == torch.float16
assert out.device == A.device
- torch.library.opcheck(torch.ops.bitsandbytes.int8_mm_dequant, (A, row_stats, col_stats))
+ opcheck(torch.ops.bitsandbytes.int8_mm_dequant, (A, row_stats, col_stats))
@pytest.mark.parametrize("device", get_available_devices())
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32], ids=id_formatter("dtype"))
@@ -89,7 +95,7 @@ def test_int8_scaled_mm(self, device, dtype, has_bias):
assert out.dtype == dtype
assert out.device == A.device
- torch.library.opcheck(torch.ops.bitsandbytes.int8_scaled_mm, (A, B, row_stats, col_stats, bias, dtype))
+ opcheck(torch.ops.bitsandbytes.int8_scaled_mm, (A, B, row_stats, col_stats, bias, dtype))
class TestInt8BlockwiseQuantOps:
@@ -115,7 +121,7 @@ def test_quantize_blockwise(self, device, dtype, blocksize):
assert absmax.device == A.device
assert absmax.dtype == torch.float32
- torch.library.opcheck(torch.ops.bitsandbytes.quantize_blockwise, (A, code, blocksize))
+ opcheck(torch.ops.bitsandbytes.quantize_blockwise, (A, code, blocksize))
@pytest.mark.parametrize("device", get_available_devices())
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32], ids=id_formatter("dtype"))
@@ -137,7 +143,7 @@ def test_dequantize_blockwise(self, device, dtype, blocksize):
assert out.dtype == dtype
assert out.device == A.device
- torch.library.opcheck(torch.ops.bitsandbytes.dequantize_blockwise.default, (A, absmax, code, blocksize, dtype))
+ opcheck(torch.ops.bitsandbytes.dequantize_blockwise.default, (A, absmax, code, blocksize, dtype))
class Test4bitBlockwiseQuantOps:
@@ -163,7 +169,7 @@ def test_quantize_4bit(self, device, dtype, storage_dtype, quant_type, blocksize
assert absmax.device == A.device
assert absmax.dtype == torch.float32
- torch.library.opcheck(torch.ops.bitsandbytes.quantize_4bit, (A, blocksize, quant_type, storage_dtype))
+ opcheck(torch.ops.bitsandbytes.quantize_4bit, (A, blocksize, quant_type, storage_dtype))
@pytest.mark.parametrize("device", get_available_devices())
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32], ids=id_formatter("dtype"))
@@ -198,8 +204,9 @@ def test_dequantize_4bit(self, device, dtype, storage_dtype, quant_type, blocksi
assert out.device == A.device
assert out.shape == shape
- torch.library.opcheck(
- torch.ops.bitsandbytes.dequantize_4bit.default, (A, absmax, blocksize, quant_type, shape, dtype)
+ opcheck(
+ torch.ops.bitsandbytes.dequantize_4bit.default,
+ (A, absmax, blocksize, quant_type, shape, dtype),
)
@pytest.mark.parametrize("device", get_available_devices())
@@ -226,4 +233,4 @@ def test_gemv_4bit(self, device, dtype, storage_dtype, quant_type, blocksize):
assert out.shape == (1, 1, out_features)
assert out.isreal().all()
- torch.library.opcheck(torch.ops.bitsandbytes.gemv_4bit.default, (A, B_q, B.shape, absmax, code, blocksize))
+ opcheck(torch.ops.bitsandbytes.gemv_4bit.default, (A, B_q, B.shape, absmax, code, blocksize))
diff --git a/tests/test_triton.py b/tests/test_triton.py
index 70656a56f..b245e534a 100644
--- a/tests/test_triton.py
+++ b/tests/test_triton.py
@@ -11,7 +11,7 @@
not is_triton_available() or not torch.cuda.is_available() or not torch.cuda.get_device_capability()[0] >= 8,
reason="This test requires triton and a GPU with compute capability 8.0 or higher.",
)
-@pytest.mark.skip("No longer supported.")
+@pytest.mark.deprecated
@pytest.mark.parametrize("vector_wise_quantization", TRUE_FALSE)
def test_switchback(vector_wise_quantization):
for dim in [83]:
From 9f85829479ac2299dcf692e913a83911a1069ad4 Mon Sep 17 00:00:00 2001
From: Matthew Douglas <38992547+matthewdouglas@users.noreply.github.com>
Date: Fri, 23 May 2025 23:45:49 -0400
Subject: [PATCH 26/85] Add torch.compile tests (#1648)
* Add torch.compile tests
* Tests: WA aarch64 CPU regressions for torch 2.6.0; add Windows torch==2.7.0+cu118 test config
* Tests: skip torch.compile for cuda on windows
---
.github/workflows/tests.yml | 30 +++++++++++-
bitsandbytes/functional.py | 4 +-
bitsandbytes/nn/modules.py | 2 +-
tests/test_linear4bit.py | 92 ++++++++++++++++++++++++++++++++++++-
tests/test_linear8bitlt.py | 66 ++++++++++++++++++++++++++
5 files changed, 188 insertions(+), 6 deletions(-)
diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
index f1a5dca69..b93bff4f0 100644
--- a/.github/workflows/tests.yml
+++ b/.github/workflows/tests.yml
@@ -137,6 +137,10 @@ jobs:
with:
python-version: 3.9
+ - name: Setup MSVC
+ if: startsWith(matrix.os, 'windows')
+ uses: ilammy/msvc-dev-cmd@v1.13.0 # to use cl for torch.compile
+
- name: Install dependencies
run: |
pip install torch==${{ matrix.torch_version }} --index-url https://download.pytorch.org/whl/cpu
@@ -201,18 +205,40 @@ jobs:
torch_version: "2.7.0"
pypi_index: "https://download.pytorch.org/whl/cu128"
- # L40S runners
+
+ # Linux L40S runners
- os: ubuntu-22.04
gpu: L40S
runner: bandb-aws-g6e-4xlarge-plus-use1-public-80
- # T4 runners
+ # Linux T4 runners
- os: ubuntu-22.04
gpu: T4
runner: bandb-aws-g4dn-4xlarge-plus-use1-public-80
+
+ # Specific Windows runners using cu118
+ - os: windows-2025
+ arch: x86_64
+ gpu: T4
+ runner: CUDA-Windows-x64
+ cuda_version: "11.8.0"
+ torch_version: "2.2.0"
+ pypi_index: "https://download.pytorch.org/whl/cu118"
- os: windows-2025
+ arch: x86_64
+ gpu: T4
+ runner: CUDA-Windows-x64
+ cuda_version: "11.8.0"
+ torch_version: "2.6.0"
+ pypi_index: "https://download.pytorch.org/whl/cu118"
+ - os: windows-2025
+ arch: x86_64
gpu: T4
runner: CUDA-Windows-x64
+ cuda_version: "11.8.0"
+ torch_version: "2.7.0"
+ pypi_index: "https://download.pytorch.org/whl/cu118"
+
exclude:
# Our current T4 Windows runner has a driver too old (471.11)
# and cannot support CUDA 12+. Skip for now.
diff --git a/bitsandbytes/functional.py b/bitsandbytes/functional.py
index b0092ffd1..f84f16c21 100644
--- a/bitsandbytes/functional.py
+++ b/bitsandbytes/functional.py
@@ -771,14 +771,14 @@ def quantize_blockwise(
qabsmax, state2 = quantize_blockwise(_absmax, blocksize=blocksize, nested=False)
quant_state = QuantState(
absmax=qabsmax,
- code=code,
+ code=code.to(A.device, copy=True),
blocksize=blocksize,
dtype=A.dtype,
offset=offset,
state2=state2,
)
else:
- quant_state = QuantState(absmax=_absmax, code=code.to(A.device), blocksize=blocksize, dtype=A.dtype)
+ quant_state = QuantState(absmax=_absmax, code=code.to(A.device, copy=True), blocksize=blocksize, dtype=A.dtype)
# TODO(matthewdouglas): Deprecate out kwarg
out = out.copy_(_out) if out is not None else _out
diff --git a/bitsandbytes/nn/modules.py b/bitsandbytes/nn/modules.py
index 937084cf1..500102ab1 100644
--- a/bitsandbytes/nn/modules.py
+++ b/bitsandbytes/nn/modules.py
@@ -493,7 +493,7 @@ def forward(self, x: torch.Tensor):
bias = None if self.bias is None else self.bias.to(self.compute_dtype)
- return bnb.matmul_4bit(x, self.weight.data.t(), bias=bias, quant_state=self.weight.quant_state).to(inp_dtype)
+ return bnb.matmul_4bit(x, self.weight.t(), bias=bias, quant_state=self.weight.quant_state).to(inp_dtype)
class LinearFP4(Linear4bit):
diff --git a/tests/test_linear4bit.py b/tests/test_linear4bit.py
index 67b61cb05..f3673797c 100644
--- a/tests/test_linear4bit.py
+++ b/tests/test_linear4bit.py
@@ -1,13 +1,21 @@
import copy
import os
import pickle
+import platform
from tempfile import TemporaryDirectory
import pytest
import torch
import bitsandbytes as bnb
-from tests.helpers import TRUE_FALSE, get_available_devices, id_formatter, torch_load_from_buffer, torch_save_to_buffer
+from tests.helpers import (
+ TRUE_FALSE,
+ describe_dtype,
+ get_available_devices,
+ id_formatter,
+ torch_load_from_buffer,
+ torch_save_to_buffer,
+)
storage = {
"uint8": torch.uint8,
@@ -275,3 +283,85 @@ def test_params4bit_real_serialization(device, quant_type, blocksize, compress_s
# there was a bug where deepcopy would modify the original object
assert dict_keys_before == dict_keys_after
assert dict_keys_before == dict_keys_deserialized
+
+
+@pytest.mark.parametrize("device", get_available_devices())
+@pytest.mark.parametrize("quant_type", ["nf4", "fp4"])
+@pytest.mark.parametrize("compute_dtype", [torch.bfloat16, torch.float32], ids=describe_dtype)
+@pytest.mark.parametrize("compress_statistics", TRUE_FALSE, ids=id_formatter("compress_statistics"))
+@pytest.mark.parametrize("bias", TRUE_FALSE, ids=id_formatter("bias"))
+@pytest.mark.parametrize("fullgraph", TRUE_FALSE, ids=id_formatter("fullgraph"))
+@pytest.mark.parametrize("mode", ["default", "reduce-overhead"], ids=id_formatter("mode"))
+@pytest.mark.skipif(torch.__version__ < (2, 4), reason="Not supported in torch < 2.4")
+def test_linear4bit_torch_compile(device, quant_type, compute_dtype, compress_statistics, bias, fullgraph, mode):
+ if device == "cpu" and quant_type == "fp4":
+ pytest.skip("FP4 is not supported for CPU")
+
+ if fullgraph and torch.__version__ < (2, 8):
+ pytest.skip("fullgraph mode requires torch 2.8 or higher")
+
+ if device == "cuda" and platform.system() == "Windows":
+ pytest.skip("Triton is not officially supported on Windows")
+
+ # Has a strange regression on Linux aarch64 CPU in torch==2.6.0 when fullgraph=False.
+ if (
+ not fullgraph
+ and device == "cpu"
+ and platform.machine() == "aarch64"
+ and platform.system() == "Linux"
+ and ((2, 7) > torch.__version__ >= (2, 6))
+ ):
+ pytest.xfail("Regression in torch==2.6.0 on Linux aarch64 CPU")
+
+ dim = 256
+ batch_size = 16
+
+ torch.compiler.reset()
+
+ # Create a small network with Linear4bit layers
+ net = torch.nn.Sequential(
+ *[
+ bnb.nn.Linear4bit(
+ dim,
+ dim,
+ bias=bias,
+ compute_dtype=compute_dtype,
+ compress_statistics=compress_statistics,
+ quant_type=quant_type,
+ )
+ for _ in range(4)
+ ]
+ ).to(device)
+
+ # Create input tensor
+ x = torch.randn(batch_size, dim, dtype=compute_dtype, device=device)
+
+ # Get reference output before compilation
+ with torch.no_grad():
+ ref_output = net(x)
+
+ # Compile the model
+ compiled_net = torch.compile(net, fullgraph=fullgraph, mode=mode)
+
+ # Get output from compiled model
+ with torch.no_grad():
+ compiled_output = compiled_net(x)
+
+ # Check outputs match
+ assert compiled_output.shape == ref_output.shape
+ assert compiled_output.device == ref_output.device
+ assert compiled_output.dtype == ref_output.dtype
+ torch.testing.assert_close(compiled_output, ref_output)
+
+ # Test with gradients
+ x.requires_grad_(True)
+ y1 = net(x).sum()
+ y1.backward()
+ grad_ref = x.grad.clone()
+
+ x.grad = None
+ y2 = compiled_net(x).sum()
+ y2.backward()
+ grad_compiled = x.grad.clone()
+
+ torch.testing.assert_close(grad_compiled, grad_ref)
diff --git a/tests/test_linear8bitlt.py b/tests/test_linear8bitlt.py
index 8c08cfa2c..a77c693e0 100644
--- a/tests/test_linear8bitlt.py
+++ b/tests/test_linear8bitlt.py
@@ -2,6 +2,7 @@
import copy
import os
import pickle
+import platform
from tempfile import TemporaryDirectory
import pytest
@@ -224,3 +225,68 @@ def test_linear8bit_serialization(linear8bit):
# check for a bug where SCB and CB were not copied
assert (linear8bit.weight.SCB == deserialized.weight.SCB).all()
assert (linear8bit.weight.CB == deserialized.weight.CB).all()
+
+
+@pytest.mark.parametrize("device", get_available_devices())
+@pytest.mark.parametrize("threshold", [0.0, 6.0], ids=id_formatter("threshold"))
+@pytest.mark.parametrize("bias", TRUE_FALSE, ids=id_formatter("bias"))
+@pytest.mark.parametrize("fullgraph", TRUE_FALSE, ids=id_formatter("fullgraph"))
+@pytest.mark.parametrize("mode", ["default", "reduce-overhead"], ids=id_formatter("mode"))
+@pytest.mark.skipif(torch.__version__ < (2, 4), reason="Not supported in torch < 2.4")
+def test_linear8bitlt_torch_compile(device, threshold, bias, fullgraph, mode):
+ if device == "cuda" and platform.system() == "Windows":
+ pytest.skip("Triton is not officially supported on Windows")
+
+ dim = 256
+ batch_size = 16
+
+ torch.compiler.reset()
+
+ # Create a small network with Linear8bitLt layers
+ net = torch.nn.Sequential(
+ *[bnb.nn.Linear8bitLt(dim, dim, bias=bias, has_fp16_weights=False, threshold=threshold) for _ in range(4)]
+ ).to(device)
+
+ dynamic_output_shapes = fullgraph and threshold > 0
+ with torch._dynamo.config.patch("capture_dynamic_output_shape_ops", dynamic_output_shapes):
+ # Create input tensor
+ x = torch.randn(batch_size, dim, dtype=torch.float16, device=device)
+
+ # Get reference output before compilation
+ with torch.no_grad():
+ ref_output = net(x)
+
+ # Compile the model
+ compiled_net = torch.compile(net, fullgraph=fullgraph, mode=mode)
+
+ # Get output from compiled model
+ with torch.no_grad():
+ compiled_output = compiled_net(x)
+
+ # Check outputs match
+ assert compiled_output.shape == ref_output.shape
+ assert compiled_output.device == ref_output.device
+ assert compiled_output.dtype == ref_output.dtype
+ torch.testing.assert_close(compiled_output, ref_output)
+
+ # Test with gradients. Currently only works with threshold=0.
+ # Has a strange regression on Linux aarch64 CPU in torch==2.6.0.
+ is_broken_platform = (
+ device == "cpu"
+ and platform.machine() == "aarch64"
+ and platform.system() == "Linux"
+ and ((2, 7) > torch.__version__ >= (2, 6))
+ )
+
+ if threshold == 0 and not is_broken_platform:
+ x.requires_grad_(True)
+ y1 = net(x).sum()
+ y1.backward()
+ grad_ref = x.grad.clone()
+
+ x.grad = None
+ y2 = compiled_net(x).sum()
+ y2.backward()
+ grad_compiled = x.grad.clone()
+
+ torch.testing.assert_close(grad_compiled, grad_ref)
From 198d08fc475a7ee7653dd5a7f15d532647c455a1 Mon Sep 17 00:00:00 2001
From: Matthew Douglas <38992547+matthewdouglas@users.noreply.github.com>
Date: Fri, 23 May 2025 23:55:45 -0400
Subject: [PATCH 27/85] Documentation Cleanup (#1644)
* Start cleaning up docs
* Remove page
* Minor update
* correction
* Minor doc revisions
* Update installation.mdx
* Update _toctree.yml
---
docs/source/_toctree.yml | 13 +-
docs/source/algorithms.mdx | 12 --
docs/source/contributing.mdx | 3 +-
docs/source/faqs.mdx | 2 -
docs/source/installation.mdx | 229 ++++++++-------------------
docs/source/non_cuda_backends.mdx | 44 -----
docs/source/reference/functional.mdx | 5 -
7 files changed, 73 insertions(+), 235 deletions(-)
delete mode 100644 docs/source/algorithms.mdx
delete mode 100644 docs/source/non_cuda_backends.mdx
diff --git a/docs/source/_toctree.yml b/docs/source/_toctree.yml
index 5fa353d6d..0f46fe6b0 100644
--- a/docs/source/_toctree.yml
+++ b/docs/source/_toctree.yml
@@ -2,18 +2,15 @@
sections:
- local: index
title: bitsandbytes
- - local: quickstart
- title: Quickstart
- local: installation
title: Installation
-- title: Guides
+ - local: quickstart
+ title: Quickstart
+
+- title: Usage Guides
sections:
- local: optimizers
title: 8-bit optimizers
- - local: algorithms
- title: Algorithms
- - local: non_cuda_backends
- title: Non-CUDA compute backends
- local: fsdp_qlora
title: FSDP-QLoRA
- local: integrations
@@ -56,7 +53,7 @@
title: RMSprop
- local: reference/optim/sgd
title: SGD
- - title: k-bit quantizers
+ - title: Modules
sections:
- local: reference/nn/linear8bit
title: LLM.int8()
diff --git a/docs/source/algorithms.mdx b/docs/source/algorithms.mdx
deleted file mode 100644
index 65e5567a4..000000000
--- a/docs/source/algorithms.mdx
+++ /dev/null
@@ -1,12 +0,0 @@
-# Other algorithms
-_WIP: Still incomplete... Community contributions would be greatly welcome!_
-
-This is an overview of the `bnb.functional` API in `bitsandbytes` that we think would also be useful as standalone entities.
-
-## Using Int8 Matrix Multiplication
-
-For straight Int8 matrix multiplication without mixed precision decomposition you can use ``bnb.matmul(...)``. To enable mixed precision decomposition, use the threshold parameter:
-
-```py
-bnb.matmul(..., threshold=6.0)
-```
diff --git a/docs/source/contributing.mdx b/docs/source/contributing.mdx
index 5da42961e..464f92164 100644
--- a/docs/source/contributing.mdx
+++ b/docs/source/contributing.mdx
@@ -1,5 +1,4 @@
-# Contributors guidelines
-... still under construction ... (feel free to propose materials, `bitsandbytes` is a community project)
+# Contribution Guide
## Setup
diff --git a/docs/source/faqs.mdx b/docs/source/faqs.mdx
index b95a1d799..c81257451 100644
--- a/docs/source/faqs.mdx
+++ b/docs/source/faqs.mdx
@@ -3,5 +3,3 @@
Please submit your questions in [this Github Discussion thread](https://github.com/bitsandbytes-foundation/bitsandbytes/discussions/1013) if you feel that they will likely affect a lot of other users and that they haven't been sufficiently covered in the documentation.
We'll pick the most generally applicable ones and post the QAs here or integrate them into the general documentation (also feel free to submit doc PRs, please).
-
-# ... under construction ...
diff --git a/docs/source/installation.mdx b/docs/source/installation.mdx
index 704d7aacc..11dfbf5ea 100644
--- a/docs/source/installation.mdx
+++ b/docs/source/installation.mdx
@@ -1,91 +1,65 @@
# Installation Guide
-Welcome to the installation guide for the `bitsandbytes` library! This document provides step-by-step instructions to install `bitsandbytes` across various platforms and hardware configurations. The library primarily supports CUDA-based GPUs, but the team is actively working on enabling support for additional backends like AMD ROCm, Intel, and Apple Silicon.
-
-> [!TIP]
-> For a high-level overview of backend support and compatibility, see the [Multi-backend Support](#multi-backend) section.
+Welcome to the installation guide for the `bitsandbytes` library! This document provides step-by-step instructions to install `bitsandbytes` across various platforms and hardware configurations. The library primarily supports CUDA-based GPUs, but the team is actively working on enabling support for additional backends like CPU, AMD ROCm, Intel XPU, and Gaudi HPU.
## Table of Contents
- [CUDA](#cuda)
- [Installation via PyPI](#cuda-pip)
- [Compile from Source](#cuda-compile)
-- [Multi-backend Support (Alpha Release)](#multi-backend)
+ - [Preview Wheels from `main`](#cuda-preview)
+- [Multi-Backend Preview](#multi-backend)
- [Supported Backends](#multi-backend-supported-backends)
- [Pre-requisites](#multi-backend-pre-requisites)
- [Installation](#multi-backend-pip)
- [Compile from Source](#multi-backend-compile)
-- [PyTorch CUDA Versions](#pytorch-cuda-versions)
## CUDA[[cuda]]
-`bitsandbytes` is currently only supported on CUDA GPUs for CUDA versions **11.0 - 12.8**. However, there's an ongoing multi-backend effort under development, which is currently in alpha. If you're interested in providing feedback or testing, check out [the multi-backend section below](#multi-backend).
-
-### Supported CUDA Configurations[[cuda-pip]]
-
-The latest version of the distributed `bitsandbytes` package is built with the following configurations:
-
-| **OS** | **CUDA Toolkit** | **Host Compiler** |
-|-------------|------------------|----------------------|
-| **Linux** | 11.8 - 12.3 | GCC 11.4 |
-| | 12.4 - 12.8 | GCC 13.2 |
-| **Windows** | 11.8 - 12.8 | MSVC 19.42+ (VS2022) |
-
-For CUDA systems, ensure your hardware meets the following requirements:
+`bitsandbytes` is currently supported on NVIDIA GPUs with [Compute Capability](https://developer.nvidia.com/cuda-gpus) 5.0+.
+The library can be built using CUDA Toolkit versions as old as **11.6** on Windows and **11.4** on Linux.
-| **Feature** | **Minimum Hardware Requirement** |
-|---------------------------------|---------------------------------------------------------------|
-| LLM.int8() | NVIDIA Turing (RTX 20 series, T4) or newer GPUs |
-| 8-bit optimizers/quantization | NVIDIA Maxwell (GTX 900 series, TITAN X, M40) or newer GPUs * |
-| NF4/FP4 quantization | NVIDIA Maxwell (GTX 900 series, TITAN X, M40) or newer GPUs * |
+| **Feature** | **CC Required** | **Example Hardware Requirement** |
+|---------------------------------|-----------------|---------------------------------------------|
+| LLM.int8() | 7.5+ | Turing (RTX 20 series, T4) or newer GPUs |
+| 8-bit optimizers/quantization | 5.0+ | Maxwell (GTX 900 series, TITAN X, M40) or newer GPUs |
+| NF4/FP4 quantization | 5.0+ | Maxwell (GTX 900 series, TITAN X, M40) or newer GPUs |
> [!WARNING]
-> `bitsandbytes >= 0.45.0` no longer supports Kepler GPUs.
->
> Support for Maxwell GPUs is deprecated and will be removed in a future release. For the best results, a Turing generation device or newer is recommended.
-```bash
-pip install bitsandbytes
-```
-
-### `pip install` pre-built wheel from latest `main` commit
-
-If you would like to use new feature even before they are officially released and help us test them, feel free to install the wheel directly from our CI (*the wheel links will remain stable!*):
+### Installation via PyPI[[cuda-pip]]
-