bitsandbytes 0.50.0__py3-none-win_arm64.whl

bitsandbytes/autograd/_functions.py

@@ -0,0 +1,491 @@
+from dataclasses import dataclass
+import logging
+from math import prod
+from typing import Optional
+import warnings
+from warnings import warn
+
+import torch
+
+import bitsandbytes.functional as F
+
+logger = logging.getLogger(__name__)
+
+# 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
+
+
+"""
+    This class pools outlier dimensions across layers.
+    This is particularly important for small models where outlier features
+    are less systematic and occur with low frequency.
+"""
+
+
+class GlobalOutlierPooler:
+    _instance = None
+
+    def __init__(self):
+        raise RuntimeError("Call get_instance() instead")
+
+    def initialize(self):
+        self.outliers = set()
+        self.model_dim = None
+
+    @classmethod
+    def get_instance(cls):
+        if cls._instance is None:
+            cls._instance = cls.__new__(cls)
+            cls._instance.initialize()
+        return cls._instance
+
+    def add_outliers(self, outlier_idx, feature_dim):
+        if self.model_dim is None:
+            self.model_dim = feature_dim
+        if feature_dim != self.model_dim:
+            return  # we do not encode outliers for the 2nd FFN layer
+
+        self.outliers.update(outlier_idx.tolist())
+
+    def get_current_outlier_idx(self):
+        return torch.Tensor(list(self.outliers)).to(torch.int64)
+
+
+_is_compiling = torch.compiler.is_compiling
+
+
+@dataclass
+class MatmulLtState:
+    force_no_igemmlt: bool = False
+
+    CB: Optional[torch.Tensor] = None
+    SB: Optional[torch.Tensor] = None
+    SCB: Optional[torch.Tensor] = None
+
+    SBt: Optional[torch.Tensor] = None
+    CBt: Optional[torch.Tensor] = None
+
+    subB: Optional[torch.Tensor] = None
+
+    outlier_pool: Optional[GlobalOutlierPooler] = None
+    has_accumulated_gradients = False
+    threshold = 0.0
+    idx: Optional[torch.Tensor] = None
+    is_training = True
+    has_fp16_weights = True
+    use_pool = False
+
+    # Deprecated attributes kept for downstream compatibility (TGI, vLLM).
+    # These are always None and will be fully removed in the next release.
+    _deprecated_fields = frozenset({"CxB", "CxBt", "formatB", "_tile_indices"})
+
+    def __getattr__(self, name):
+        if name in MatmulLtState._deprecated_fields:
+            warnings.warn(
+                f"MatmulLtState.{name} is deprecated and will be removed in the next bitsandbytes release.",
+                FutureWarning,
+                stacklevel=2,
+            )
+            return None
+        raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'")
+
+    def reset_grads(self):
+        self.CB = None
+        self.SB = None
+        self.SCB = None
+
+        self.SBt = None
+        self.CBt = None
+
+
+class MatMul8bitLt(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx: torch.autograd.function.FunctionCtx,
+        A: torch.Tensor,
+        B: torch.Tensor,
+        out: Optional[torch.Tensor] = None,
+        bias: Optional[torch.Tensor] = None,
+        state: Optional[MatmulLtState] = None,
+    ):
+        state = state or MatmulLtState()
+
+        # default of pytorch behavior if inputs are empty
+        ctx.is_empty = False
+        if prod(A.shape) == 0:
+            ctx.is_empty = True
+            ctx.A = A
+            ctx.B = B
+            ctx.bias = bias
+            if A.shape[-1] == B.shape[0]:
+                return torch.empty(A.shape[:-1] + B.shape[1:], dtype=A.dtype, device=A.device)
+            else:
+                return torch.empty(A.shape[:-1] + B.shape[:1], dtype=A.dtype, device=A.device)
+
+        input_shape = A.shape
+
+        # Cast A to fp16
+        if A.dtype != torch.float16 and not _is_compiling():
+            logger.warning("MatMul8bitLt: inputs will be cast from %s to float16 during quantization", A.dtype)
+
+        if len(A.shape) == 3:
+            A = A.reshape(-1, A.shape[-1])
+
+        # 1. Quantize A. Note that as a side-effect, outliers are suppressed in CA/CAt.
+        if ctx.needs_input_grad[1]:
+            # Slower path
+            CA, CAt, SCA, SCAt, outlier_cols = F.int8_double_quant(A.to(torch.float16), threshold=state.threshold)
+        else:
+            # Fast path
+            CA, SCA, outlier_cols = F.int8_vectorwise_quant(A.to(torch.float16), threshold=state.threshold)
+            CAt = SCAt = None
+
+        has_grad = False
+
+        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()
+
+                # 2. Quantize B
+                state.CB, state.SCB, _ = F.int8_vectorwise_quant(B.to(torch.float16))
+
+        # Handle sparse decomposition
+        if state.threshold > 0.0:
+            state.idx = outlier_cols
+
+            # Mixed Int8 Matmul + Dequant + Bias
+            output, subA = torch.ops.bitsandbytes.int8_mixed_scaled_mm(
+                A,
+                CA,
+                state.CB,
+                SCA,
+                state.SCB,
+                outlier_cols,
+                bias,
+            )
+
+        else:
+            # Int8 Matmul + Dequant + Bias
+            output = torch.ops.bitsandbytes.int8_scaled_mm.default(
+                CA, state.CB, SCA, state.SCB, bias=bias, dtype=A.dtype
+            )
+            subA = None
+
+        # 5. Save state
+        ctx.state = state
+
+        ctx.grad_shape = input_shape
+        ctx.dtype_A = A.dtype
+        ctx.dtype_bias = None if bias is None else bias.dtype
+
+        if any(ctx.needs_input_grad[:2]):
+            ctx.tensors = (CAt, subA, A)
+            ctx.tensor_states = (SCAt, state.idx)
+        else:
+            ctx.tensors = [None, None, None]
+            ctx.tensor_states = (None, None)
+            ctx.save_for_backward(None, None)
+
+        output_shape = (*input_shape[:-1], state.CB.shape[0])
+
+        if len(input_shape) == 3:
+            return output.reshape(output_shape)
+
+        return output
+
+    @staticmethod
+    def backward(ctx: torch.autograd.function.FunctionCtx, grad_output: torch.Tensor):
+        if ctx.is_empty:
+            bias_grad = None if ctx.bias is None else torch.zeros_like(ctx.bias)
+            return torch.zeros_like(ctx.A), torch.zeros_like(ctx.B), None, bias_grad, None
+
+        req_gradA, req_gradB, _, req_gradBias, _ = ctx.needs_input_grad
+        CAt, subA, _A = ctx.tensors
+        SCAt, idx = ctx.tensor_states
+        state: MatmulLtState = 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:
+            Cgrad, _, _, SCgradt, _ = F.int8_double_quant(grad_output.to(torch.float16))
+
+            grad_B = torch.ops.bitsandbytes.int8_scaled_mm.default(
+                Cgrad.t().contiguous(),
+                CAt.t(),
+                SCgradt,
+                SCAt,
+                dtype=torch.float16,
+            )
+
+            if state.threshold > 0.0 and subA is not None and subA.numel() > 0:
+                grad_B[:, idx] += torch.matmul(grad_output.t(), subA)
+
+        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 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)
+        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
+
+    @staticmethod
+    def forward(ctx, A, B, out=None, bias=None, quant_state: Optional[F.QuantState] = None):
+        # default of pytorch behavior if inputs are empty
+        ctx.is_empty = False
+        if A.numel() == 0:
+            ctx.is_empty = True
+            ctx.A = A
+            ctx.B = B
+            ctx.bias = bias
+            B_shape = quant_state.shape
+            if A.shape[-1] == B_shape[0]:
+                return torch.empty(A.shape[:-1] + B_shape[1:], dtype=A.dtype, device=A.device)
+            else:
+                return torch.empty(A.shape[:-1] + B_shape[:1], dtype=A.dtype, device=A.device)
+
+        # Normalize to canonical [(N*K+1)//2, 1]. Packed weights are always contiguous
+        # in this orientation (B.t() callers get strides [1,1], still compatible).
+        # quant_state.shape is the source of truth for N and K.
+        B = B.view(-1, 1)
+
+        if not quant_state.nested:
+            output = torch.ops.bitsandbytes.gemm_4bit.default(
+                A,
+                B,
+                quant_state.shape,
+                quant_state.absmax,
+                quant_state.blocksize,
+                quant_state.quant_type,
+                bias=bias,
+            )
+        elif quant_state.state2.blocksize == 256:
+            output = torch.ops.bitsandbytes.gemm_4bit.default(
+                A,
+                B,
+                quant_state.shape,
+                quant_state.state2.absmax,
+                quant_state.blocksize,
+                quant_state.quant_type,
+                bias=bias,
+                absmax_8bit=quant_state.absmax,
+                absmax_code=quant_state.state2.code,
+                absmax_offset=quant_state.offset,
+            )
+        else:
+            raise NotImplementedError("nested quantization with state2.blocksize != 256 is not supported")
+
+        if out is not None:
+            out.copy_(output)
+            output = out
+
+        # 3. Save state
+        ctx.state = quant_state
+        ctx.dtype_A, ctx.dtype_B, ctx.dtype_bias = A.dtype, B.dtype, None if bias is None else bias.dtype
+
+        if any(ctx.needs_input_grad[:2]):
+            ctx.tensors = (None, B)
+        else:
+            ctx.tensors = (None, None)
+
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        if ctx.is_empty:
+            bias_grad = None if ctx.bias is None else torch.zeros_like(ctx.bias)
+            return torch.zeros_like(ctx.A), torch.zeros_like(ctx.B), None, bias_grad, None
+
+        req_gradA, _, _, req_gradBias, _ = ctx.needs_input_grad
+        _, B = ctx.tensors
+
+        grad_A, grad_B, grad_bias = None, None, None
+
+        if req_gradBias:
+            # compute grad_bias first before changing grad_output dtype
+            grad_bias = grad_output.sum(0, dtype=ctx.dtype_bias)
+
+        # not supported by PyTorch. TODO: create work-around
+        # if req_gradB: grad_B = torch.matmul(grad_output.t(), A)
+        if req_gradA:
+            # B in ctx.tensors is already in canonical [(N*K+1)//2, 1] form (normalized in forward).
+            # dequantize returns [N, K]; matmul(grad_output[M,N], [N,K]) = grad_A[M,K].
+            grad_A = torch.matmul(grad_output, F.dequantize_4bit(B, ctx.state).to(grad_output.dtype))
+
+        return grad_A, grad_B, None, grad_bias, None
+
+
+def matmul(
+    A: torch.Tensor,
+    B: torch.Tensor,
+    out: Optional[torch.Tensor] = None,
+    state: Optional[MatmulLtState] = None,
+    threshold=0.0,
+    bias: Optional[torch.Tensor] = None,
+):
+    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 in ("cpu", "xpu"):
+            return MatMul8bitFp.apply(A, B, out, bias, state)
+    return MatMul8bitLt.apply(A, B, out, bias, state)
+
+
+def matmul_4bit(
+    A: torch.Tensor,
+    B: torch.Tensor,
+    quant_state: F.QuantState,
+    out: Optional[torch.Tensor] = None,
+    bias: Optional[torch.Tensor] = None,
+):
+    if quant_state is None:
+        raise ValueError("quant_state is required")
+    if len(quant_state.shape) != 2:
+        raise ValueError("matmul_4bit: quant_state.shape must be 2D [N, K]")
+
+    # packing_format_for_cpu uses a different memory layout optimized for AVX512BF16.
+    # This flag is only set for inference (weight conversion happens at eval time).
+    # The underlying kernel supports any M via tiled GEMM despite the gemv name.
+    if A.device.type == "cpu" and getattr(quant_state, "packing_format_for_cpu", False):
+        result = F.gemv_4bit(A, B, out=out, state=quant_state)
+        if bias is not None:
+            result += bias
+        return result
+
+    # Normalize B to canonical [(N*K+1)//2, 1]. Packed weights are always contiguous
+    # in this orientation (B.t() callers get strides [1,1], still compatible).
+    # quant_state.shape is the source of truth for N and K.
+    B = B.view(-1, 1)
+
+    K = A.shape[-1]
+
+    # Weight is in [K, N] orientation when A's inner dim matches shape[0] not shape[1].
+    # Square weights (K==N) are ambiguous and treated as [N, K].
+    if K == quant_state.shape[0] and K != quant_state.shape[1]:
+        if not _is_compiling():
+            warn(
+                f"matmul_4bit: weight was quantized from a [K, N] tensor (quant_state.shape={list(quant_state.shape)}). "
+                "Re-quantize from the weight in [N, K] (out_features, in_features) orientation. "
+                "This will be an error in a future version.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
+        B_dq = F.dequantize_4bit(B, quant_state).to(A.dtype)
+        result = torch.nn.functional.linear(A, B_dq.t(), bias)
+        if out is not None:
+            out.copy_(result)
+            return out
+        return result
+
+    needs_grad = torch.is_grad_enabled() and (A.requires_grad or (bias is not None and bias.requires_grad))
+    if not needs_grad:
+        A_numel = A.numel()
+        if A_numel == 0:
+            if out is not None:
+                return out
+            return torch.empty((*A.shape[:-1], quant_state.shape[0]), dtype=A.dtype, device=A.device)
+
+        if not quant_state.nested:
+            result = torch.ops.bitsandbytes.gemm_4bit.default(
+                A,
+                B,
+                quant_state.shape,
+                quant_state.absmax,
+                quant_state.blocksize,
+                quant_state.quant_type,
+                bias=bias,
+            )
+        elif quant_state.state2.blocksize == 256:
+            result = torch.ops.bitsandbytes.gemm_4bit.default(
+                A,
+                B,
+                quant_state.shape,
+                quant_state.state2.absmax,
+                quant_state.blocksize,
+                quant_state.quant_type,
+                bias=bias,
+                absmax_8bit=quant_state.absmax,
+                absmax_code=quant_state.state2.code,
+                absmax_offset=quant_state.offset,
+            )
+        else:
+            raise NotImplementedError("nested quantization with state2.blocksize != 256 is not supported")
+        if out is not None:
+            out.copy_(result)
+            return out
+        return result
+
+    return MatMul4Bit.apply(A, B, out, bias, quant_state)