diff --git a/bitsandbytes/backends/default/ops.py b/bitsandbytes/backends/default/ops.py index 9aa2d60c2..33bdd7ed3 100644 --- a/bitsandbytes/backends/default/ops.py +++ b/bitsandbytes/backends/default/ops.py @@ -1,6 +1,6 @@ from collections.abc import Sequence from functools import cache, wraps -from math import sqrt +from math import prod, sqrt from typing import Optional import torch @@ -196,7 +196,7 @@ def _(A: torch.Tensor, code: torch.Tensor, blocksize: int) -> tuple[torch.Tensor return q.reshape(A.shape), absmax -@_try_torch_compile(dynamic=True) +@_try_torch_compile(dynamic=False) def _dequantize_blockwise_compute( A_flat: torch.Tensor, absmax: torch.Tensor, code: torch.Tensor, blocksize: int, dtype: torch.dtype ): @@ -259,7 +259,7 @@ def _( return packed, absmax -@_try_torch_compile(dynamic=True) +@_try_torch_compile(dynamic=False) def _dequantize_4bit_compute( A_flat: torch.Tensor, absmax: torch.Tensor, @@ -268,9 +268,7 @@ def _dequantize_4bit_compute( shape: Sequence[int], dtype: torch.dtype, ): - n = 1 - for s in shape: - n *= s + n = prod(shape) out_dq = torch.empty(A_flat.size(0) * 2, dtype=torch.int32, device=A_flat.device) out_dq[1::2] = A_flat & 0xF out_dq[::2] = A_flat >> 4 diff --git a/bitsandbytes/backends/mps/ops.py b/bitsandbytes/backends/mps/ops.py index 0e3186663..04c3d4fda 100644 --- a/bitsandbytes/backends/mps/ops.py +++ b/bitsandbytes/backends/mps/ops.py @@ -1,38 +1,131 @@ -"""MPS backend for bitsandbytes 4-bit quantization ops. +"""MPS backend for bitsandbytes quantization ops. -Uses Metal kernels from kernels-community/bitsandbytes-mps via the -HuggingFace Kernels Hub. +Hub kernels (kernels-community/bitsandbytes-mps) are attempted lazily on +macOS 26+. On older macOS the hub kernel path is skipped entirely and +default fallbacks or MPS-specific pure PyTorch fallbacks are used for all ops. + +Note: not all ops have implementations on the Hub kernels. Those that do not are also +implemented using pure PyTorch fallbacks. """ from collections.abc import Sequence from math import prod +import platform from typing import Optional -from warnings import warn import torch from ..._ops import register_kernel -from ..default.ops import _gemm_4bit_default_impl +from ..default.ops import ( + _dequantize_4bit_compute, + _get_4bit_quantize_bounds, + _try_torch_compile, +) +from ..utils import _get_4bit_code -# --------------------------------------------------------------------------- -# Quant-type mapping: BnB uses strings, our Metal kernel uses ints. -# --------------------------------------------------------------------------- _QUANT_MAP = {"fp4": 1, "nf4": 2} + _kernel = None +_macos_major = int(platform.mac_ver()[0].split(".")[0]) if platform.mac_ver()[0] else 0 + +# Pre-set to True on macOS < 26 so _get_kernel() never attempts the import. +_kernel_load_failed = _macos_major < 26 + def _get_kernel(): - """Lazily load the bitsandbytes-mps kernel (local build or Hub).""" - global _kernel - if _kernel is None: + global _kernel, _kernel_load_failed + if _kernel_load_failed: + return None + if _kernel is not None: + return _kernel + try: from kernels import get_kernel - # TODO: use kernels-community/bitsandbytes-mps when it's available - _kernel = get_kernel("kernels-community/bitsandbytes-mps") + _kernel = get_kernel("kernels-community/bitsandbytes-mps", version=1) + except Exception: + _kernel_load_failed = True + return None return _kernel -# ============================= quantize_4bit ================================= +@_try_torch_compile(dynamic=True) +def _quantize_blockwise_compute( + A_flat: torch.Tensor, code: torch.Tensor, blocksize: int +) -> tuple[torch.Tensor, torch.Tensor]: + """ + On torch <= 2.12, torch.bucketize does not perform well. + Implements blockwise quantization using a binary search instead of using the default. + """ + n = A_flat.numel() + rem = n % blocksize + full = n - rem + blocks = full // blocksize + A_com = A_flat[:full].reshape(blocks, blocksize) + absmax = A_com.abs().max(dim=-1)[0] + scaled = torch.clamp(A_com * (1.0 / absmax.clamp(min=1e-38).view(-1, 1)), -1, 1).reshape(-1) + if rem: + am = A_flat[full:].abs().max().clamp(min=1e-38) + absmax = torch.cat([absmax, am.unsqueeze(0)]) + scaled = torch.cat([scaled, torch.clamp(A_flat[full:] / am, -1, 1)]) + bounds = (code[:-1] + code[1:]) / 2 + n_bounds = bounds.shape[0] + n_iters = n_bounds.bit_length() + lo = torch.zeros(scaled.shape, dtype=torch.int16, device=scaled.device) + hi = torch.full(scaled.shape, n_bounds, dtype=torch.int16, device=scaled.device) + for _ in range(n_iters): + mid = (lo + hi) >> 1 + val = bounds[mid.to(torch.int64)] + lo = torch.where(val < scaled, (mid + 1).to(torch.int16), lo) + hi = torch.where(val >= scaled, mid, hi) + return lo.to(torch.uint8), absmax + + +@register_kernel("bitsandbytes::quantize_blockwise", "mps") +def _(A: torch.Tensor, code: torch.Tensor, blocksize: int) -> tuple[torch.Tensor, torch.Tensor]: + q, absmax = _quantize_blockwise_compute(A.reshape(-1).float(), code.float(), blocksize) + return q.reshape(A.shape), absmax + + +@_try_torch_compile(dynamic=True) +def _quantize_4bit_compute( + A_flat: torch.Tensor, + blocksize: int, + bounds: torch.Tensor, + order: torch.Tensor, + nf4: bool, +) -> tuple[torch.Tensor, torch.Tensor]: + n = A_flat.numel() + rem = n % blocksize + full = n - rem + blocks = full // blocksize + A_com = A_flat[:full].reshape(blocks, blocksize) + absmax = A_com.abs().max(dim=-1)[0] + scaled = torch.clamp(A_com * (1.0 / absmax.clamp(min=1e-38).view(-1, 1)), -1, 1).reshape(-1) + if rem: + am = A_flat[full:].abs().max().clamp(min=1e-38) + absmax = torch.cat([absmax, am.unsqueeze(0)]) + scaled = torch.cat([scaled, torch.clamp(A_flat[full:] / am, -1, 1)]) + if scaled.numel() % 2: + scaled = torch.nn.functional.pad(scaled, (0, 1)) + idx = torch.zeros(scaled.shape, dtype=torch.int8, device=scaled.device) + for b in bounds: + idx = idx + (scaled > b).to(torch.int8) + if not nf4: + idx = order[idx.to(torch.int32)] + q8 = idx.to(torch.uint8) + return (q8[::2] << 4) | q8[1::2], absmax + + +def _quantize_4bit_fallback( + A: torch.Tensor, blocksize: int, quant_type: str, quant_storage: torch.dtype +) -> tuple[torch.Tensor, torch.Tensor]: + bounds, order = _get_4bit_quantize_bounds(quant_type, A.device) + packed, absmax = _quantize_4bit_compute(A.reshape(-1).float(), blocksize, bounds, order, quant_type == "nf4") + packed = packed.unsqueeze(1) + if quant_storage != torch.uint8: + packed = packed.squeeze().view(quant_storage).unsqueeze(1) + return packed, absmax @register_kernel("bitsandbytes::quantize_4bit", "mps") @@ -42,18 +135,11 @@ def _( quant_type: str, quant_storage: torch.dtype, ) -> tuple[torch.Tensor, torch.Tensor]: - if blocksize not in (64, 128, 256, 512): - raise ValueError(f"MPS backend only supports blocksize in (64, 128, 256, 512), got {blocksize}") - - k = _get_kernel() - packed, absmax = k.quantize_4bit(A.contiguous(), blocksize, _QUANT_MAP[quant_type]) - - packed = packed.view(quant_storage).unsqueeze(1) - - return packed, absmax - - -# ============================ dequantize_4bit ================================ + if blocksize in (64, 128, 256, 512) and (k := _get_kernel()) is not None: + packed, absmax = k.quantize_4bit(A.contiguous(), blocksize, _QUANT_MAP[quant_type]) + packed = packed.view(quant_storage).unsqueeze(1) + return packed, absmax + return _quantize_4bit_fallback(A, blocksize, quant_type, quant_storage) def _dequantize_4bit_impl( @@ -67,10 +153,15 @@ def _dequantize_4bit_impl( if A.dtype != torch.uint8: A = A.view(torch.uint8) - numel = prod(shape) - k = _get_kernel() - out = k.dequantize_4bit(A, absmax, blocksize, _QUANT_MAP[quant_type], numel, dtype) - return out.reshape(shape) + # Use HF Hub kernel when supported. + if blocksize in (64, 128, 256, 512) and (k := _get_kernel()) is not None: + numel = prod(shape) + out = k.dequantize_4bit(A, absmax, blocksize, _QUANT_MAP[quant_type], numel, dtype) + return out.reshape(shape) + + # Fallback to implementation from default backend. + code = _get_4bit_code(quant_type, A.device) + return _dequantize_4bit_compute(A.reshape(-1), absmax, code, blocksize, shape, dtype) @register_kernel("bitsandbytes::dequantize_4bit", "mps") @@ -82,8 +173,6 @@ def _( shape: Sequence[int], dtype: torch.dtype, ) -> torch.Tensor: - if blocksize not in (64, 128, 256, 512): - raise ValueError(f"MPS backend only supports blocksize in (64, 128, 256, 512), got {blocksize}") return _dequantize_4bit_impl(A, absmax, blocksize, quant_type, shape, dtype) @@ -101,9 +190,6 @@ def _( out.copy_(result) -# ================================ gemv_4bit ================================== - - def _gemv_4bit_impl( A: torch.Tensor, B: torch.Tensor, @@ -112,14 +198,18 @@ def _gemv_4bit_impl( code: torch.Tensor, blocksize: int, ) -> torch.Tensor: - if B.dtype != torch.uint8: - B = B.view(torch.uint8) + if blocksize in (64, 128, 256) and (k := _get_kernel()) is not None: + if B.dtype != torch.uint8: + B = B.view(torch.uint8) - quant_type_int = _QUANT_MAP["fp4"] if code[1] > 0 else _QUANT_MAP["nf4"] - output_features = shapeB[0] + output_features = shapeB[0] + quant_type_int = _QUANT_MAP["fp4"] if code[1] > 0 else _QUANT_MAP["nf4"] - k = _get_kernel() - return k.gemv_4bit(A, B, absmax, output_features, blocksize, quant_type_int) + return k.gemv_4bit(A, B, absmax, output_features, blocksize, quant_type_int) + + quant_type = "fp4" if code[1] > 0 else "nf4" + B_dq = _dequantize_4bit_impl(B, absmax, blocksize, quant_type, shapeB, A.dtype) + return torch.nn.functional.linear(A, B_dq) @register_kernel("bitsandbytes::gemv_4bit", "mps") @@ -165,36 +255,23 @@ def _( M = A.numel() // K N = shapeB[0] + # For nested absmax, we don't have a fused implementation yet. + # Dequantize the absmax values first. if absmax_8bit is not None: absmax = ( torch.ops.bitsandbytes.dequantize_blockwise.default(absmax_8bit, absmax, absmax_code, 256, torch.float32) + absmax_offset ) - if M == 1: - if K % blocksize == 0: - if B.dtype != torch.uint8: - B = B.view(torch.uint8) - - k = _get_kernel() - result = k.gemv_4bit(A, B, absmax.view(N, -1), N, blocksize, _QUANT_MAP[quant_type]) - - if bias is not None: - result = result + bias - return result - - warn( - f"inner dimension ({K}) is not aligned for fast kernel " - f"with blocksize={blocksize}, falling back to slower implementation.", - UserWarning, - ) - - return _gemm_4bit_default_impl( - A, - B, - shapeB, - absmax, - blocksize, - quant_type, - bias, - ) + # Use HF Hub kernel when supported for GEMV. + if M == 1 and blocksize in (64, 128, 256) and (k := _get_kernel()) is not None: + if B.dtype != torch.uint8: + B = B.view(torch.uint8) + result = k.gemv_4bit(A, B, absmax.view(N, -1), N, blocksize, _QUANT_MAP[quant_type]) + if bias is not None: + result = result + bias + return result + + # Fallback: dequantize + linear. + B_dq = _dequantize_4bit_impl(B, absmax, blocksize, quant_type, shapeB, A.dtype) + return torch.nn.functional.linear(A, B_dq, bias)