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