diff --git a/bitsandbytes/backends/cuda/ops.py b/bitsandbytes/backends/cuda/ops.py index 1f3ce82e0..45ca37417 100644 --- a/bitsandbytes/backends/cuda/ops.py +++ b/bitsandbytes/backends/cuda/ops.py @@ -3,7 +3,6 @@ import functools from math import prod from typing import Optional -from warnings import warn import torch @@ -11,6 +10,7 @@ from ..._ops import register_kernel from ...cextension import lib +from ..utils import _warn_gemm_4bit_unaligned def _setup_ctypes(names, argtypes, restype=None): @@ -943,15 +943,13 @@ def _( # fallback. if M > _gemm_4bit_custom_max_m: use_custom = False - elif K % blocksize != 0: - warn( - f"inner dimension ({K}) is not aligned for fast kernel " - f"with blocksize={blocksize}, falling back to slower implementation.", - UserWarning, - ) - use_custom = False else: use_custom = _gemm_4bit_use_custom_fn(A.device.index, A.dtype, M, N, K) + if use_custom and K % blocksize != 0: + # Only warn when misalignment is what costs us the fused kernel; when the + # heuristic picks the fallback anyway (larger M, e.g. training), it doesn't. + _warn_gemm_4bit_unaligned(K, blocksize) + use_custom = False if not use_custom: return _dequant_linear_fallback( diff --git a/bitsandbytes/backends/utils.py b/bitsandbytes/backends/utils.py index a63e59a99..4a52b37d1 100644 --- a/bitsandbytes/backends/utils.py +++ b/bitsandbytes/backends/utils.py @@ -1,4 +1,6 @@ +import functools from importlib.metadata import metadata +from warnings import warn from packaging import version import torch @@ -75,6 +77,17 @@ def _get_4bit_code(quant_type: str, device: torch.device) -> torch.Tensor: return _code_4bit_cache[key] +@functools.cache +def _warn_gemm_4bit_unaligned(K: int, blocksize: int) -> None: + """Warn at most once per (K, blocksize): the misalignment is a fixed property of the + model, so repeating it on every gemm_4bit call is pure noise.""" + warn( + f"inner dimension ({K}) is not aligned for fast kernel " + f"with blocksize={blocksize}, falling back to slower implementation.", + UserWarning, + ) + + def get_gaudi_sw_version(): """ Returns the installed version of Gaudi SW. diff --git a/bitsandbytes/backends/xpu/ops.py b/bitsandbytes/backends/xpu/ops.py index 731200c53..a8f0fe734 100644 --- a/bitsandbytes/backends/xpu/ops.py +++ b/bitsandbytes/backends/xpu/ops.py @@ -2,7 +2,6 @@ import ctypes as ct import logging from typing import Optional -from warnings import warn from packaging import version import torch @@ -12,7 +11,7 @@ from ..._ops import register_kernel from ...cextension import ErrorHandlerMockBNBNativeLibrary, lib from ..default.ops import _gemm_4bit_default_impl -from ..utils import _get_4bit_code, triton_available +from ..utils import _get_4bit_code, _warn_gemm_4bit_unaligned, triton_available logger = logging.getLogger(__name__) @@ -186,11 +185,7 @@ def _( out = out + bias return out - warn( - f"inner dimension ({K}) is not aligned for fast kernel " - f"with blocksize={blocksize}, falling back to slower implementation.", - UserWarning, - ) + _warn_gemm_4bit_unaligned(K, blocksize) return _gemm_4bit_default_impl( A, diff --git a/tests/test_ops.py b/tests/test_ops.py index 4ca60f845..4fc57cbed 100644 --- a/tests/test_ops.py +++ b/tests/test_ops.py @@ -1,4 +1,5 @@ from math import prod +import warnings import pytest import torch @@ -333,6 +334,28 @@ def test_gemm_4bit(self, device, dtype, quant_type, compress_statistics, has_bia kwargs={"bias": bias}, ) + @pytest.mark.parametrize("device", get_available_devices()) + def test_gemm_4bit_unaligned_warning(self, device): + """Regression test for #2027: the blocksize-alignment warning must not be emitted + on every call, nor at all when the fused kernel was not going to be used anyway.""" + N, K, blocksize = 128, 3420, 64 # 3420 % 64 != 0 (Qwen2.5-VL vision tower) + B = torch.randn(N, K, dtype=torch.float16, device=device) + B_q, qs = bitsandbytes.functional.quantize_4bit(B, blocksize=blocksize, quant_type="nf4") + + def run(M): + A = torch.randn(M, K, dtype=torch.float16, device=device) + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + for _ in range(2): + torch.ops.bitsandbytes.gemm_4bit.default(A, B_q, list(B.shape), qs.absmax, blocksize, "nf4") + return [w for w in caught if "not aligned" in str(w.message)] + + # Large M always takes the dequant+F.linear path, aligned or not. + assert run(1024) == [] + + # When alignment does decide it, warn at most once per (K, blocksize). + assert len(run(1)) <= 1 + @pytest.mark.parametrize("device", get_available_devices()) @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16], ids=describe_dtype) @pytest.mark.parametrize("offset_dtype", [torch.float16, torch.bfloat16], ids=describe_dtype)