- | 🐧 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)
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/bitsandbytes/__init__.py b/bitsandbytes/__init__.py
index 917cd0b6a..c747398ce 100644
--- a/bitsandbytes/__init__.py
+++ b/bitsandbytes/__init__.py
@@ -34,6 +34,9 @@
if torch.cuda.is_available():
from .backends.cuda import ops as cuda_ops
+if hasattr(torch, "xpu") and torch.xpu.is_available():
+ from .backends.xpu import ops as xpu_ops
+
def _import_backends():
"""
@@ -64,4 +67,4 @@ def _import_backends():
"optim.optimizer.MockArgs": False,
}
-__version__ = "0.46.0.dev0"
+__version__ = "0.47.0.dev0"
diff --git a/bitsandbytes/_ops.py b/bitsandbytes/_ops.py
index 9a3ac46ac..a260852f5 100644
--- a/bitsandbytes/_ops.py
+++ b/bitsandbytes/_ops.py
@@ -4,6 +4,8 @@
import torch
+from .cextension import ipex_cpu, ipex_xpu
+
_IS_TORCH_GTE_24 = False
if hasattr(torch.library, "register_fake"):
@@ -327,3 +329,22 @@ def _(
)
torch._check(out.device == A.device, lambda: f"Expected out.device == {A.device}, got {out.device}")
torch._check(out.dtype == A.dtype, lambda: f"Expected out.dtype == {A.dtype}, got {out.dtype}")
+
+
+if ipex_cpu or ipex_xpu:
+ # Register the dequantize_nf4_ipex implementation
+ torch.library.define(
+ "bitsandbytes::dequantize_nf4_ipex",
+ "(Tensor A, Tensor absmax, int blocksize, int[] shape, ScalarType dtype) -> Tensor",
+ )
+
+ @register_fake("bitsandbytes::dequantize_nf4_ipex")
+ def _(
+ A: torch.Tensor,
+ absmax: torch.Tensor,
+ blocksize: int,
+ shape: Sequence[int],
+ dtype: torch.dtype,
+ ) -> torch.Tensor:
+ torch._check_is_size(blocksize)
+ return torch.empty(shape, dtype=dtype, device=A.device)
diff --git a/bitsandbytes/autograd/_functions.py b/bitsandbytes/autograd/_functions.py
index c7ad3a82c..746d6c1ec 100644
--- a/bitsandbytes/autograd/_functions.py
+++ b/bitsandbytes/autograd/_functions.py
@@ -8,6 +8,7 @@
from typing_extensions import deprecated
import bitsandbytes.functional as F
+from bitsandbytes.functional import ipex_cpu, ipex_xpu
# The inverse transformation for the colTuring and colAmpere format were contributed by Alex Borzunov:
# https://github.com/bigscience-workshop/petals/blob/main/src/petals/utils/linear8bitlt_patch.py
@@ -298,6 +299,63 @@ def backward(ctx: torch.autograd.function.FunctionCtx, grad_output: torch.Tensor
return grad_A, grad_B, None, grad_bias, None
+class MatMul8bitFp(torch.autograd.Function):
+ # For Intel CPU and XPU MatMul8bitFp is much faster (~3x) than MatMul8bitLt in finetune.
+ # Because the MatMul8bitLt has more mechanisms in computing grad.
+ # We don't have fast kernel for quant/dequant 8bit in CPU/XPU, so it's very slow.
+ # We'd like to use dequant + matmul to run finetune with good performance.
+
+ @staticmethod
+ def forward(ctx, A, B, out=None, bias=None, state=MatmulLtState):
+ if state.has_fp16_weights or state.CB is None:
+ has_grad = getattr(B, "grad", None) is not None
+ is_transposed = not B.is_contiguous() and B.shape[0] == B.stride(1)
+ if is_transposed:
+ B = B.contiguous()
+
+ if (state.is_training and not has_grad) or state.CB is None or state.SCB is None:
+ state.reset_grads()
+ state.CB, state.SCB, _ = F.int8_vectorwise_quant(B.to(torch.float16))
+ B = state.CB
+
+ CB = state.CB.data.to(A.dtype).mul_(state.SCB.unsqueeze(1).mul(1.0 / 127.0))
+ output = torch.nn.functional.linear(A, CB, bias)
+ # to pass the test: tests/test_modules.py::test_linear8bitlt_no_fp16_weights[2.0-xpu]
+ state.idx = False
+ ctx.state = state
+ ctx.dtype_A = A.dtype
+ ctx.grad_shape = A.shape
+ ctx.A = A
+ ctx.dtype_bias = None if bias is None else bias.dtype
+ return output
+
+ @staticmethod
+ def backward(ctx, grad_output):
+ req_gradA, req_gradB, _, req_gradBias, _ = ctx.needs_input_grad
+ A = ctx.A
+ state = ctx.state
+ grad_A = grad_B = grad_bias = None
+ if req_gradBias:
+ # compute grad_bias first before changing grad_output dtype
+ grad_bias = grad_output.sum(0, dtype=ctx.dtype_bias)
+
+ # Cast grad_output to fp16
+ if len(grad_output.shape) == 3:
+ grad_output = grad_output.reshape(-1, grad_output.shape[-1]).contiguous()
+
+ if req_gradB:
+ grad_B = torch.matmul(A.t(), grad_output).t()
+
+ if req_gradA:
+ if state.CB is not None:
+ CB = state.CB.to(ctx.dtype_A, copy=True).mul_(state.SCB.unsqueeze(1).mul(1.0 / 127.0))
+ grad_A = torch.matmul(grad_output.to(ctx.dtype_A), CB).view(ctx.grad_shape)
+ else:
+ raise Exception("State must contain CB matrix for backward")
+
+ return grad_A, grad_B, None, grad_bias, None
+
+
class MatMul4Bit(torch.autograd.Function):
# forward is the same, but we added the fallback for pre-turing GPUs
# backward is mostly the same, but adds one extra clause (see "elif state.CxB is not None")
@@ -366,6 +424,10 @@ def matmul(
state = state or MatmulLtState()
if threshold > 0.0:
state.threshold = threshold
+ # MatMul8bitLt is slower because no fast kernel for quant/dequant 8bit in CPU/XPU
+ if state.is_training:
+ if (A.device.type == "cpu" and ipex_cpu) or (A.device.type == "xpu" and ipex_xpu):
+ return MatMul8bitFp.apply(A, B, out, bias, state)
return MatMul8bitLt.apply(A, B, out, bias, state)
@@ -378,6 +440,17 @@ def matmul_4bit(
):
assert quant_state is not None
+ if A.device.type in ("cpu", "xpu") and A.requires_grad == False:
+ if getattr(quant_state, "ipex", False):
+ # IPEX CPU will change weight to 4D so don't need transpose
+ B = B.t() if B.dim() == 2 else B
+ out = F.gemv_4bit(A, B, out, state=quant_state)
+ if bias is not None:
+ out += bias
+ return out
+ else:
+ return MatMul4Bit.apply(A, B, out, bias, quant_state)
+
if A.numel() == A.shape[-1] and A.requires_grad == False:
if A.shape[-1] % quant_state.blocksize != 0:
warn(
diff --git a/bitsandbytes/backends/cpu/ops.py b/bitsandbytes/backends/cpu/ops.py
index d5ab9aa88..5f009ea40 100644
--- a/bitsandbytes/backends/cpu/ops.py
+++ b/bitsandbytes/backends/cpu/ops.py
@@ -7,6 +7,7 @@
from ..._ops import register_kernel
from ...cextension import lib
+from ..utils import ipex_cpu
# torch._int_mm for s8@s8->s32 is supported on CPU from torch 2.4+.
# However, we can overflow if we use this without AVX512_VNNI support.
@@ -26,22 +27,42 @@ def _(A: torch.Tensor, B: torch.Tensor):
@register_kernel("bitsandbytes::quantize_blockwise", "cpu")
def _(A: torch.Tensor, code: torch.Tensor, blocksize: int) -> tuple[torch.Tensor, torch.Tensor]:
torch._check_is_size(blocksize)
- torch._check(A.dtype == torch.float32, lambda: f"A must be float32 on cpu, got {A.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)
-
- lib.cquantize_blockwise_cpu_fp32(
- get_ptr(code),
- get_ptr(A),
- get_ptr(absmax),
- get_ptr(out),
- ct.c_longlong(blocksize),
- ct.c_longlong(n),
- )
+ # Only FP32 has c++ kernrl
+ if A.dtype == torch.float32:
+ blocks = -(n // -blocksize)
+
+ absmax = torch.empty((blocks,), device=A.device, dtype=torch.float32)
+ out = torch.empty_like(A, dtype=torch.uint8)
+
+ lib.cquantize_blockwise_cpu_fp32(
+ get_ptr(code),
+ get_ptr(A),
+ get_ptr(absmax),
+ get_ptr(out),
+ ct.c_longlong(blocksize),
+ ct.c_longlong(n),
+ )
+ else:
+ rem = n % blocksize
+ has_rem = rem > 0
+ blocks = n // blocksize + has_rem
+ absmax = torch.zeros((blocks,), device=A.device, dtype=torch.float32)
+ A_reshaped = A.reshape(n)
+ A_com = A_reshaped[: n - rem]
+ A_com_reshaped = A_com.reshape(n // blocksize, blocksize)
+ absmax[: blocks - has_rem] = torch.abs(A_com_reshaped).max(dim=-1)[0]
+ scaled_A = torch.clamp(A_com_reshaped * (1 / absmax[: blocks - has_rem].view(-1, 1)), -1, 1)
+ scaled_A = scaled_A.reshape(-1)
+ if has_rem:
+ absmax[-1] = torch.abs(A_reshaped[n - rem :]).max()
+ scaled_A_rem = torch.clamp(A_reshaped[n - rem :] * (1 / absmax[-1]), -1, 1)
+ scaled_A = torch.cat([scaled_A, scaled_A_rem], dim=0)
+
+ diff = torch.abs(scaled_A.unsqueeze(-1) - code.to(scaled_A.device))
+ out = torch.argmin(diff, dim=-1).to(torch.uint8).to(scaled_A.device).reshape(A.shape)
return out, absmax
@@ -50,144 +71,50 @@ def _(A: torch.Tensor, code: torch.Tensor, blocksize: int) -> tuple[torch.Tensor
def _(A: torch.Tensor, absmax: torch.Tensor, code: torch.Tensor, blocksize: int, dtype: torch.dtype) -> torch.Tensor:
torch._check_is_size(blocksize)
torch._check(A.dtype == torch.uint8, lambda: f"A must be uint8, got {A.dtype}")
- torch._check(dtype == torch.float32, lambda: f"dtype must be float32 on cpu, got {dtype}")
-
- out = torch.empty_like(A, dtype=dtype)
- lib.cdequantize_blockwise_cpu_fp32(
- get_ptr(code),
- get_ptr(A),
- get_ptr(absmax),
- get_ptr(out),
- ct.c_longlong(blocksize),
- ct.c_longlong(A.numel()),
- )
+ # Only FP32 has c++ kernrl
+ if dtype == torch.float32:
+ out = torch.empty_like(A, dtype=dtype)
+
+ lib.cdequantize_blockwise_cpu_fp32(
+ get_ptr(code),
+ get_ptr(A),
+ get_ptr(absmax),
+ get_ptr(out),
+ ct.c_longlong(blocksize),
+ ct.c_longlong(A.numel()),
+ )
+ else:
+ out = code[A.reshape(-1).int()]
+ blocks = out.shape[-1] // blocksize
+ res = out.shape[-1] % blocksize
+ if res != 0:
+ out = torch.nn.functional.pad(out, (0, blocksize - res), mode="constant", value=0)
+ out = (out.view(-1, blocksize) * absmax.view(-1, 1)).to(dtype).reshape(-1)
+ out = out[: blocks * blocksize + res]
+ out = out.reshape(A.shape)
return out
-_NF4_QUANT_TABLE = torch.tensor(
- [
- -1.0,
- -0.6961928009986877,
- -0.5250730514526367,
- -0.39491748809814453,
- -0.28444138169288635,
- -0.18477343022823334,
- -0.09105003625154495,
- 0.0,
- 0.07958029955625534,
- 0.16093020141124725,
- 0.24611230194568634,
- 0.33791524171829224,
- 0.44070982933044434,
- 0.5626170039176941,
- 0.7229568362236023,
- 1.0,
- ],
- dtype=torch.float32,
- device="cpu",
-)
-
-
-@register_kernel("bitsandbytes::quantize_4bit", "cpu")
-def _(
- A: torch.Tensor, blocksize: int, quant_type: str, quant_storage: torch.dtype
-) -> tuple[torch.Tensor, torch.Tensor]:
- torch._check_is_size(blocksize)
- torch._check(quant_type == "nf4", lambda: f"quant_type must be nf4 on CPU, got {quant_type}")
- torch._check(
- A.dtype in [torch.bfloat16, torch.float16, torch.float32],
- lambda: f"Blockwise 4bit quantization only supports 16/32-bit floats, but got {A.dtype}",
- )
-
- n = A.numel()
-
- # TODO: Support when weight matrix is not divisible by blocksize
- torch._check(n % blocksize == 0, lambda: f"n must be divisible by blocksize, got {n} and {blocksize}")
-
- # Divide into blocks and normalize
- blocks = A.reshape(-1, blocksize)
- absmax = blocks.abs().max(dim=1).values.float()
- scaled = blocks / absmax.unsqueeze(-1)
-
- # Quantize with the lookup table
- quantized = torch.argmin(torch.abs(scaled.view(-1, 1) - _NF4_QUANT_TABLE), dim=-1, keepdim=True).to(torch.uint8)
-
- # Pack two quantized values per byte
- packed = quantized[::2] << 4 | quantized[1::2]
-
- if quant_storage != torch.uint8:
- packed = packed.squeeze().view(quant_storage).unsqueeze(1)
-
- return packed, absmax.float()
-
-
-@register_kernel("bitsandbytes::dequantize_4bit", "cpu")
-def _(
- A: torch.Tensor,
- absmax: torch.Tensor,
- blocksize: int,
- quant_type: str,
- shape: Sequence[int],
- dtype: torch.dtype,
-) -> torch.Tensor:
- torch._check_is_size(blocksize)
- torch._check(quant_type == "nf4", lambda: f"quant_type must be nf4 on CPU, got {quant_type}")
- torch._check(
- dtype in [torch.bfloat16, torch.float16, torch.float32],
- lambda: f"Blockwise 4bit dequantization only supports 16/32-bit floats, but got {dtype}",
- )
- torch._check(
- A.dtype == torch.uint8,
- lambda: f"Blockwise 4bit dequantization on CPU only supports uint8 storage, got {A.dtype}",
- )
-
- A = A.view(-1, 1)
-
- # Grab upper and lower nibbles. Using int64 for indexing in the LUT.
- upper = (A >> 4).to(torch.int64)
- lower = (A & 0x0F).to(torch.int64)
-
- # Expand to blocks
- blocks = torch.cat((upper, lower), dim=1).reshape(-1, blocksize)
-
- # Dequantize
- blocks = _NF4_QUANT_TABLE[blocks] * absmax[:, None]
-
- # Reshape to original shape
- blocks = blocks.reshape(-1, *shape[1:])
-
- return blocks.to(dtype)
-
-
-@register_kernel("bitsandbytes::gemv_4bit", "cpu")
-def _(
- A: torch.Tensor,
- B: torch.Tensor,
- shapeB: Sequence[int],
- absmax: torch.Tensor,
- code: torch.Tensor,
- blocksize: int,
-) -> torch.Tensor:
- # TODO: We need to determine whether `code` is NF4, FP4, or other.
- # Right now we assume NF4, as this is the only one supported on CPU.
-
- B_dq = torch.ops.bitsandbytes.dequantize_4bit.default(
- B,
- absmax,
- blocksize,
- "nf4",
- shape=shapeB,
- dtype=A.dtype,
- )
-
- # User called gemv with B.t(), so we need to transpose it back.
- # if B.shape[0] == 1:
- # B_dq = B_dq.t()
-
- return torch.nn.functional.linear(
- A,
- B_dq,
- bias=None,
- )
+if ipex_cpu:
+ from bitsandbytes.utils import _reverse_4bit_compress_format
+
+ @register_kernel("bitsandbytes::dequantize_nf4_ipex", "cpu")
+ def _(
+ A: torch.Tensor,
+ absmax: torch.Tensor,
+ blocksize: int,
+ shape: Sequence[int],
+ dtype: torch.dtype,
+ ) -> torch.Tensor:
+ ipex_weight = torch.ops.ipex_prepack.woq_linear_unpack_weight(A, "nf4", shape, 2)
+ A = _reverse_4bit_compress_format(ipex_weight.reshape(-1)).reshape(1, -1)
+ return torch.ops.bitsandbytes.dequantize_4bit.default(
+ A,
+ absmax,
+ blocksize,
+ "nf4",
+ shape,
+ dtype,
+ )
diff --git a/bitsandbytes/backends/cuda/ops.py b/bitsandbytes/backends/cuda/ops.py
index efdef2871..13359bbd8 100644
--- a/bitsandbytes/backends/cuda/ops.py
+++ b/bitsandbytes/backends/cuda/ops.py
@@ -8,7 +8,7 @@
from bitsandbytes.functional import CUBLAS_Context, _cuda_device_of, _get_tensor_stream, get_ptr
from ..._ops import register_kernel
-from ...cextension import lib
+from ...cextension import HIP_ENVIRONMENT, lib
@register_kernel("bitsandbytes::int8_linear_matmul", "cuda")
@@ -210,7 +210,12 @@ def _get_col_absmax(
@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)
- torch._check(blocksize in [4096, 2048, 1024, 512, 256, 128, 64])
+
+ 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()
@@ -264,7 +269,11 @@ def _(
def _dequantize_blockwise_impl(
A: torch.Tensor, absmax: torch.Tensor, code: torch.Tensor, blocksize: int, dtype: torch.dtype, out: torch.Tensor
) -> None:
- torch._check(blocksize in [4096, 2048, 1024, 512, 256, 128, 64])
+ 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],
@@ -294,7 +303,11 @@ def _dequantize_blockwise_impl(
def _(
A: torch.Tensor, blocksize: int, quant_type: str, quant_storage: torch.dtype
) -> tuple[torch.Tensor, torch.Tensor]:
- torch._check(blocksize in [4096, 2048, 1024, 512, 256, 128, 64])
+ 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(
A.dtype in [torch.bfloat16, torch.float16, torch.float32],
@@ -372,7 +385,11 @@ def _dequantize_4bit_impl(
dtype: torch.dtype,
out: torch.Tensor,
) -> None:
- torch._check(blocksize in [4096, 2048, 1024, 512, 256, 128, 64])
+ 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],
@@ -445,20 +462,22 @@ def _gemv_4bit_impl(
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}")
+
+ # Note: these checks are not strictly necessary, and cost more than they are worth, so they are commented out for now.
+ # 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)
diff --git a/bitsandbytes/backends/default/ops.py b/bitsandbytes/backends/default/ops.py
index 729c2b047..ce5926979 100644
--- a/bitsandbytes/backends/default/ops.py
+++ b/bitsandbytes/backends/default/ops.py
@@ -1,9 +1,11 @@
+from collections.abc import Sequence
from math import prod
from typing import Optional
import torch
from ..._ops import register_kernel
+from ..utils import CODE
@register_kernel("bitsandbytes::int8_mm_dequant", "default")
@@ -142,3 +144,160 @@ def _(A: torch.Tensor, threshold=0.0):
A[outliers] = outlier_restore
return out_row, row_stats, outlier_cols
+
+
+@register_kernel("bitsandbytes::quantize_blockwise", "default")
+def _(A: torch.Tensor, code: torch.Tensor, blocksize: int) -> tuple[torch.Tensor, torch.Tensor]:
+ torch._check_is_size(blocksize)
+
+ n = A.numel()
+ rem = n % blocksize
+ has_rem = rem > 0
+ blocks = n // blocksize + has_rem
+ absmax = torch.zeros((blocks,), device=A.device, dtype=torch.float32)
+ A_reshaped = A.reshape(n)
+ A_com = A_reshaped[: n - rem]
+ A_com_reshaped = A_com.reshape(n // blocksize, blocksize)
+ absmax[: blocks - has_rem] = torch.abs(A_com_reshaped).max(dim=-1)[0]
+ scaled_A = torch.clamp(A_com_reshaped * (1 / absmax[: blocks - has_rem].view(-1, 1)), -1, 1)
+ scaled_A = scaled_A.reshape(-1)
+ if has_rem:
+ absmax[-1] = torch.abs(A_reshaped[n - rem :]).max()
+ scaled_A_rem = torch.clamp(A_reshaped[n - rem :] * (1 / absmax[-1]), -1, 1)
+ scaled_A = torch.cat([scaled_A, scaled_A_rem], dim=0)
+
+ diff = torch.abs(scaled_A.unsqueeze(-1) - code.to(scaled_A.device))
+ out = torch.argmin(diff, dim=-1).to(torch.uint8).to(scaled_A.device).reshape(A.shape)
+
+ return out, absmax
+
+
+@register_kernel("bitsandbytes::dequantize_blockwise", "default")
+def _(A: torch.Tensor, absmax: torch.Tensor, code: torch.Tensor, blocksize: int, dtype: torch.dtype) -> torch.Tensor:
+ torch._check_is_size(blocksize)
+ torch._check(A.dtype == torch.uint8, lambda: f"A must be uint8, got {A.dtype}")
+
+ out = code[A.reshape(-1).int()]
+ blocks = out.shape[-1] // blocksize
+ res = out.shape[-1] % blocksize
+ if res != 0:
+ out = torch.nn.functional.pad(out, (0, blocksize - res), mode="constant", value=0)
+ out = (out.view(-1, blocksize) * absmax.view(-1, 1)).to(dtype).reshape(-1)
+ out = out[: blocks * blocksize + res]
+ out = out.reshape(A.shape)
+
+ return out
+
+
+@register_kernel("bitsandbytes::quantize_4bit", "default")
+def _(
+ A: torch.Tensor, blocksize: int, quant_type: str, quant_storage: torch.dtype
+) -> tuple[torch.Tensor, torch.Tensor]:
+ torch._check_is_size(blocksize)
+ torch._check(quant_type in ("nf4", "fp4"), lambda: f"quant_type must be nf4 or fp4, got {quant_type}")
+ torch._check(
+ A.dtype in [torch.bfloat16, torch.float16, torch.float32],
+ lambda: f"Blockwise 4bit quantization only supports 16/32-bit floats, but got {A.dtype}",
+ )
+
+ n = A.numel()
+ full_blocks = n // blocksize
+ rem = n % blocksize
+ blocks = full_blocks + 1 if rem else full_blocks
+ absmax = torch.zeros((blocks,), device=A.device, dtype=torch.float32)
+ A_flattened = A.reshape(n)
+
+ # Scale full blocks of the tensor to [-1, 1]
+ A_full_blocks = A_flattened[: n - rem].reshape(n // blocksize, blocksize)
+ absmax[:full_blocks] = torch.abs(A_full_blocks).max(dim=-1)[0]
+ scaled = torch.clamp(A_full_blocks * (1 / absmax[:full_blocks].view(-1, 1)), -1, 1).reshape(-1)
+
+ # Scale any partial block
+ if rem:
+ A_rem = A_flattened[-rem:]
+ absmax[-1] = torch.abs(A_rem).max()
+ scaled_rem = torch.clamp(A_rem * (1 / absmax[-1]), -1, 1)
+ scaled = torch.cat([scaled, scaled_rem], dim=0)
+
+ # Quantize with the lookup table
+ code = CODE[quant_type].to(scaled.device).to(scaled.dtype)
+ quantized = torch.argmin(torch.abs(scaled.view(-1, 1) - code), dim=-1, keepdim=True).to(torch.uint8)
+
+ # Pack two quantized values per byte
+ packed = quantized[::2] << 4 | quantized[1::2]
+
+ if quant_storage != torch.uint8:
+ packed = packed.squeeze().view(quant_storage).unsqueeze(1)
+
+ return packed, absmax.float()
+
+
+@register_kernel("bitsandbytes::dequantize_4bit", "default")
+def _(
+ A: torch.Tensor,
+ absmax: torch.Tensor,
+ blocksize: int,
+ quant_type: str,
+ shape: Sequence[int],
+ dtype: torch.dtype,
+) -> torch.Tensor:
+ torch._check_is_size(blocksize)
+ torch._check(quant_type in ("nf4", "fp4"), lambda: f"quant_type must be nf4 or fp4, got {quant_type}")
+ torch._check(
+ dtype in [torch.bfloat16, torch.float16, torch.float32],
+ lambda: f"Blockwise 4bit dequantization only supports 16/32-bit floats, but got {dtype}",
+ )
+
+ # Enable non uint8 dtype
+ if A.dtype != torch.uint8:
+ A = A.view(torch.uint8)
+
+ A = A.reshape(-1)
+ # Map nf4 to [-1, 1]
+ out_dq = torch.empty(A.size(0) * 2, dtype=torch.int32, device=A.device)
+ n = out_dq.numel()
+ out_dq[1::2] = A & 0xF
+ out_dq[::2] = A >> 4
+ # code is fp32, cast to dtype to avoid the mismatch issue
+ code = CODE[quant_type].to(dtype).to(A.device)
+ out_dq = code[out_dq]
+
+ # Apply scales
+ if out_dq.numel() != n:
+ assert out_dq.numel() == n + 1
+ out_dq = torch.narrow(out_dq, 0, 0, n)
+ blocks = n // blocksize
+ blocks += 1 if n % blocksize > 0 else 0
+ rem = n % blocksize
+ has_rem = rem > 0
+
+ out = torch.empty(shape, dtype=dtype, device=A.device).reshape(-1)
+ if has_rem:
+ out[: n - rem] = (out_dq[: n - rem].view(-1, blocksize) * absmax[: blocks - has_rem].view(-1, 1)).reshape(-1)
+ out[n - rem :] = out_dq[n - rem :] * absmax[-1]
+ else:
+ out = out_dq.view(-1, blocksize) * absmax.view(-1, 1)
+
+ out = out.reshape(-1, *shape[1:]).to(dtype)
+
+ return out
+
+
+@register_kernel("bitsandbytes::gemv_4bit", "default")
+def _(
+ A: torch.Tensor,
+ B: torch.Tensor,
+ shapeB: Sequence[int],
+ absmax: torch.Tensor,
+ code: torch.Tensor,
+ blocksize: int,
+) -> torch.Tensor:
+ # Applied from dequantize_4bit
+ quant_type = "fp4" if code[1] > 0 else "nf4"
+ B_dq = torch.ops.bitsandbytes.dequantize_4bit.default(B, absmax, blocksize, quant_type, shapeB, A.dtype)
+
+ return torch.nn.functional.linear(
+ A,
+ B_dq,
+ bias=None,
+ )
diff --git a/bitsandbytes/backends/utils.py b/bitsandbytes/backends/utils.py
new file mode 100755
index 000000000..bf277e7ea
--- /dev/null
+++ b/bitsandbytes/backends/utils.py
@@ -0,0 +1,61 @@
+import torch
+
+try:
+ # to support Intel CPU/XPU (IPEX) backend
+ import intel_extension_for_pytorch as ipex
+
+ ipex_cpu = ipex if ipex._C._has_cpu() else None
+ ipex_xpu = ipex if ipex._C._has_xpu() else None
+except BaseException:
+ ipex_cpu = None
+ ipex_xpu = None
+
+_NF4_QUANT_TABLE = torch.tensor(
+ [
+ -1.0,
+ -0.6961928009986877,
+ -0.5250730514526367,
+ -0.39491748809814453,
+ -0.28444138169288635,
+ -0.18477343022823334,
+ -0.09105003625154495,
+ 0.0,
+ 0.07958029955625534,
+ 0.16093020141124725,
+ 0.24611230194568634,
+ 0.33791524171829224,
+ 0.44070982933044434,
+ 0.5626170039176941,
+ 0.7229568362236023,
+ 1.0,
+ ],
+ dtype=torch.float32,
+ device="xpu"
+ if hasattr(torch, "xpu") and torch.xpu.is_available()
+ else "cpu", # Only cpu/xpu use this table for now.
+)
+_FP4_QUANT_TABLE = torch.tensor(
+ [
+ 0.0000,
+ 0.0052,
+ 0.6667,
+ 1.0000,
+ 0.3333,
+ 0.5000,
+ 0.1667,
+ 0.2500,
+ 0.0000,
+ -0.0052,
+ -0.6667,
+ -1.0000,
+ -0.3333,
+ -0.5000,
+ -0.1667,
+ -0.2500,
+ ],
+ dtype=torch.float32,
+ device="xpu"
+ if hasattr(torch, "xpu") and torch.xpu.is_available()
+ else "cpu", # Only cpu/xpu use this table for now.
+)
+CODE = {"nf4": _NF4_QUANT_TABLE, "fp4": _FP4_QUANT_TABLE}
diff --git a/bitsandbytes/backends/xpu/__init__.py b/bitsandbytes/backends/xpu/__init__.py
new file mode 100755
index 000000000..e69de29bb
diff --git a/bitsandbytes/backends/xpu/ops.py b/bitsandbytes/backends/xpu/ops.py
new file mode 100755
index 000000000..47a3bd009
--- /dev/null
+++ b/bitsandbytes/backends/xpu/ops.py
@@ -0,0 +1,51 @@
+from collections.abc import Sequence
+
+import torch
+
+from ..._ops import register_kernel
+from ..utils import ipex_xpu
+
+if torch.__version__ >= (2, 7):
+
+ @register_kernel("bitsandbytes::int8_linear_matmul", "xpu")
+ def _(A: torch.Tensor, B: torch.Tensor):
+ return torch._int_mm(
+ A.reshape(-1, A.shape[-1]),
+ B.t(),
+ ).reshape(*A.shape[:-1], B.shape[0])
+
+
+if ipex_xpu:
+
+ @register_kernel("bitsandbytes::dequantize_nf4_ipex", "xpu")
+ def _(
+ A: torch.Tensor,
+ absmax: torch.Tensor,
+ blocksize: int,
+ shape: Sequence[int],
+ dtype: torch.dtype,
+ ) -> torch.Tensor:
+ return torch.ops.torch_ipex.dequantize_4bit(A, "nf4", shape, absmax, None, blocksize).t().to(dtype)
+
+ @register_kernel("bitsandbytes::dequantize_blockwise", "xpu")
+ def _(
+ A: torch.Tensor,
+ absmax: torch.Tensor,
+ code: torch.Tensor,
+ blocksize: int,
+ dtype: torch.dtype,
+ ) -> torch.Tensor:
+ shape = A.shape
+ out = torch.empty(A.reshape(-1).shape, dtype=dtype, device=A.device)
+ # void cdequantize_blockwise_fp32(
+ # float *code, unsigned char *A, float *absmax, float *out, int blocksize, const int n, cudaStream_t stream)
+ if dtype == torch.float16:
+ ipex_xpu.xpu.bitsandbytes.cdequantize_blockwise_fp16(code, A, absmax, out, blocksize, A.numel())
+ elif dtype == torch.bfloat16:
+ ipex_xpu.xpu.bitsandbytes.cdequantize_blockwise_bf16(code, A, absmax, out, blocksize, A.numel())
+ elif dtype == torch.float32:
+ ipex_xpu.xpu.bitsandbytes.cdequantize_blockwise_fp32(code, A, absmax, out, blocksize, A.numel())
+ else:
+ raise ValueError(f"Blockwise quantization only supports 16/32-bit floats, but got {out.dtype}")
+
+ return out.reshape(shape)
diff --git a/bitsandbytes/cextension.py b/bitsandbytes/cextension.py
index c8b02fb22..7f5483531 100644
--- a/bitsandbytes/cextension.py
+++ b/bitsandbytes/cextension.py
@@ -1,4 +1,5 @@
import ctypes as ct
+import functools
import logging
import os
from pathlib import Path
@@ -8,7 +9,7 @@
import torch
from bitsandbytes.consts import DYNAMIC_LIBRARY_SUFFIX, PACKAGE_DIR
-from bitsandbytes.cuda_specs import CUDASpecs, get_cuda_specs, get_cuda_version_tuple
+from bitsandbytes.cuda_specs import CUDASpecs, get_cuda_specs, get_cuda_version_tuple, get_rocm_gpu_arch
logger = logging.getLogger(__name__)
@@ -35,10 +36,8 @@ def get_cuda_bnb_library_path(cuda_specs: CUDASpecs) -> Path:
)
logger.warning(
f"WARNING: BNB_CUDA_VERSION={override_value} environment variable detected; loading {library_name}.\n"
- "This can be used to load a bitsandbytes version that is different from the PyTorch CUDA version.\n"
+ "This can be used to load a bitsandbytes version built with a CUDA version that is different from the PyTorch CUDA version.\n"
"If this was unintended set the BNB_CUDA_VERSION variable to an empty string: export BNB_CUDA_VERSION=\n"
- "If you use the manual override make sure the right libcudart.so is in your LD_LIBRARY_PATH\n"
- "For example by adding the following to your .bashrc: export LD_LIBRARY_PATH=$LD_LIBRARY_PATH: