From f32444890916f2bf65153f960ea777333a571ddf Mon Sep 17 00:00:00 2001 From: thangdangjp Date: Wed, 9 Sep 2026 08:11:06 -0700 Subject: [PATCH] Warn once when Linear4bit computes in float32 on GPU `Linear4bit` has warnings for a float32 compute_dtype, but they live in `set_compute_type()`, which `forward()` only calls when `compute_type_is_set` is False -- i.e. when no compute_dtype was passed to the constructor. The dominant integration always passes one: `transformers.BitsAndBytesConfig` defaults `bnb_4bit_compute_dtype` to `torch.float32`, so a plain `BitsAndBytesConfig(load_in_4bit=True)` pins every layer to float32 compute and the existing warnings are unreachable. Users get the slow path with no runtime signal at all. That path is slow because float32 has no MMA 4-bit GEMM kernel, as the CUDA dispatch heuristic itself notes -- float32 only takes the custom kernel for M < 8 and otherwise falls back to unfused dequantize + F.linear. Measured on an A100 80GB (bnb 0.50, torch 2.13+cu130) with Qwen3-4B projection shapes, float32 vs bfloat16 compute_dtype: shape M=1 M=8 M=512 M=2048 q_proj (K=2560, N=4096) 0.68x 1.45x 5.64x 8.25x o_proj (K=4096, N=2560) 0.68x 0.84x 6.24x 7.73x gate_proj (K=2560, N=9728) 0.79x 2.97x 6.58x 9.30x down_proj (K=9728, N=2560) 0.81x 1.84x 7.91x 9.90x (>1 means float32 is slower). So single-token decoding is unaffected, but prefill and training run several times slower than they need to. Log a hint once per process per device type when a Linear4bit actually computes in float32 on a CUDA device. The check sits in `forward()` so it covers every way a module can reach that state, including the inferred one where `set_compute_type()` adopts float32 from float32 inputs. Two guards keep it cheap and safe: a `_fp32_compute_warned` latch reduces the steady state to one attribute load, and `torch.compiler.is_compiling()` keeps the branch out of Dynamo graphs so `fullgraph=True` compilation is unaffected. Also fix a stale comment claiming float32 is adopted "for speed", and a "compoute" typo. Author: Thang Dang - Fujitsu --- bitsandbytes/nn/modules.py | 47 +++++++++++++++++++++++++++++++++++--- tests/test_modules.py | 44 +++++++++++++++++++++++++++++++++++ 2 files changed, 88 insertions(+), 3 deletions(-) diff --git a/bitsandbytes/nn/modules.py b/bitsandbytes/nn/modules.py index ebc0b0943..d97428820 100644 --- a/bitsandbytes/nn/modules.py +++ b/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) diff --git a/tests/test_modules.py b/tests/test_modules.py index 95f78b6d3..c9601290c 100644 --- a/tests/test_modules.py +++ b/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