Warn once when Linear4bit computes in float32 on GPU

bitsandbytes/nn/modules.py

@@ -3,6 +3,7 @@
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
 import copy
+import functools
 import logging
 from typing import Any, Optional, TypeVar, Union, overload
 
@@ -25,6 +26,26 @@
 T = TypeVar("T", bound="torch.nn.Module")
 
 
+@functools.cache
+def _warn_fp32_compute_dtype(device_type: str) -> None:
+    """Log a one-time hint that a `Linear4bit` is computing in float32 on a GPU device.
+
+    `functools.cache` keeps this to one message per process per device type: a 4-bit model
+    has hundreds of `Linear4bit` layers and each one would otherwise repeat the hint.
+    """
+    logger.warning(
+        "Linear4bit is computing in torch.float32 on a %s device, which is much slower than "
+        "16-bit compute: bitsandbytes has no tensor-core (MMA) 4-bit GEMM kernel for float32, "
+        "so these layers fall back to the SIMT kernel or to an unfused dequantize + matmul. "
+        "Pass bnb_4bit_compute_dtype=torch.bfloat16 (or torch.float16) to "
+        "transformers.BitsAndBytesConfig, or compute_dtype= to Linear4bit directly, unless you "
+        "specifically need float32 compute. Note that float32 is the default in "
+        "transformers.BitsAndBytesConfig. This is logged once; silence it with "
+        "logging.getLogger('bitsandbytes.nn.modules').setLevel(logging.ERROR).",
+        device_type,
+    )
+
+
 class StableEmbedding(torch.nn.Embedding):
     """
     Custom embedding layer designed to improve stability during training for NLP tasks by using 32-bit optimizer states. It is designed to reduce gradient variations that can result from quantization. This embedding layer is initialized with Xavier uniform initialization followed by layer normalization.
@@ -534,6 +555,11 @@ class Linear4bit(nn.Linear):
     ```
     """
 
+    # Class-level default so that modules restored from a checkpoint saved by an older
+    # version, which has no such instance attribute, still resolve it. Instances that emit
+    # the warning shadow it with True; it is deliberately not part of the state dict.
+    _fp32_compute_warned = False
+
     def __init__(
         self,
         input_features,
@@ -574,11 +600,11 @@ def __init__(
 
     def set_compute_type(self, x):
         if x.dtype in [torch.float32, torch.bfloat16]:
-            # the input is in a dtype that is safe to compute in, we switch
-            # to this type for speed and stability
+            # the input is in a dtype that is safe to compute in, so we adopt it. Note that
+            # float32 is safe but slow on GPU; forward() warns about that case separately.
             self.compute_dtype = x.dtype
         elif x.dtype == torch.float16:
-            # we take the compoute dtype passed into the layer
+            # we take the compute dtype passed into the layer
             if self.compute_dtype in [None, torch.float32] and (x.numel() == x.shape[-1]):
                 # single batch inference with input torch.float16 and compute_dtype float32 -> slow inference when it could be fast
                 # warn the user about this
@@ -623,6 +649,21 @@ def forward(self, x: torch.Tensor):
             self.set_compute_type(x)
             self.compute_type_is_set = True
 
+        # A float32 compute_dtype costs several times the runtime of a 16-bit one on GPU (see
+        # _warn_fp32_compute_dtype). It is easy to end up here without meaning to, because
+        # transformers.BitsAndBytesConfig defaults bnb_4bit_compute_dtype to float32 -- and in
+        # that case compute_type_is_set is already True, so set_compute_type()'s own warnings
+        # about float32 never run. The `_fp32_compute_warned` latch keeps the common path to a
+        # single attribute load, and is_compiling() keeps this branch out of Dynamo graphs.
+        if (
+            self.compute_dtype == torch.float32
+            and not self._fp32_compute_warned
+            and x.device.type == "cuda"
+            and not torch.compiler.is_compiling()
+        ):
+            self._fp32_compute_warned = True
+            _warn_fp32_compute_dtype(x.device.type)
+
         inp_dtype = x.dtype
         if self.compute_dtype is not None:
             x = x.to(self.compute_dtype)

tests/test_modules.py

@@ -433,6 +433,50 @@ def test_4bit_linear_warnings(device, caplog):
     assert any("inference." in msg for msg in caplog.messages)
 
 
+@pytest.mark.parametrize("device", get_available_devices())
+def test_4bit_linear_fp32_compute_dtype_warning(device, caplog):
+    """float32 compute is several times slower than 16-bit on GPU (no MMA 4-bit GEMM kernel
+    exists for it), and is easy to select unintentionally because
+    transformers.BitsAndBytesConfig defaults bnb_4bit_compute_dtype to float32. Check that we
+    say so once, and only where a faster 16-bit kernel actually exists.
+    """
+    dim = 64
+    marker = "4-bit GEMM kernel for float32"
+    expect_warning = device == "cuda"
+
+    def build(compute_dtype):
+        return nn.Sequential(
+            *[bnb.nn.Linear4bit(dim, dim, compute_dtype=compute_dtype, quant_type="nf4") for _ in range(4)]
+        ).to(device)
+
+    # An explicitly configured float32 compute_dtype: compute_type_is_set is already True, so
+    # set_compute_type() -- and its own float32 warnings -- never runs.
+    bnb.nn.modules._warn_fp32_compute_dtype.cache_clear()
+    with caplog_at_level(caplog, logging.WARNING, "bitsandbytes.nn.modules"):
+        net = build(torch.float32)
+        inp = torch.rand(8, dim, device=device, dtype=torch.float32)
+        for _ in range(3):
+            net(inp)
+    # 4 layers x 3 forwards: the hint is logged once per process per device type, not per call.
+    assert len([msg for msg in caplog.messages if marker in msg]) == (1 if expect_warning else 0)
+
+    # A float32 compute_dtype inferred from float32 inputs by set_compute_type() is just as slow.
+    caplog.clear()
+    bnb.nn.modules._warn_fp32_compute_dtype.cache_clear()
+    with caplog_at_level(caplog, logging.WARNING, "bitsandbytes.nn.modules"):
+        net = build(None)
+        net(torch.rand(8, dim, device=device, dtype=torch.float32))
+    assert len([msg for msg in caplog.messages if marker in msg]) == (1 if expect_warning else 0)
+
+    # A 16-bit compute_dtype is the recommended setting and must stay silent.
+    caplog.clear()
+    bnb.nn.modules._warn_fp32_compute_dtype.cache_clear()
+    with caplog_at_level(caplog, logging.WARNING, "bitsandbytes.nn.modules"):
+        net = build(torch.bfloat16)
+        net(torch.rand(8, dim, device=device, dtype=torch.bfloat16))
+    assert not [msg for msg in caplog.messages if marker in msg]
+
+
 @pytest.mark.parametrize("device", get_available_devices())
 def test_4bit_embedding_warnings(device, caplog):
     num_embeddings = 128