Warn once when Linear4bit computes in float32 on GPU

[thangdangjp]

Problem

Linear4bit already has warnings for a float32 compute_dtype, but they are unreachable in the case that matters most.

They live in set_compute_type(), which forward() only calls when compute_type_is_set is False — that is, only when no compute_dtype was passed to the constructor:

if not self.compute_type_is_set:
    self.set_compute_type(x)
    self.compute_type_is_set = True

But the dominant integration always passes one. transformers.BitsAndBytesConfig defaults bnb_4bit_compute_dtype to torch.float32, and integrations/bitsandbytes.py forwards it to Linear4bit(...). So a plain BitsAndBytesConfig(load_in_4bit=True) sets compute_type_is_set = True with compute_dtype=torch.float32, set_compute_type() never runs, and the user lands on the slow path with no runtime signal at all — even when the model itself was loaded in bfloat16.

Additionally, the branch that does run silently picks float32 itself when it sees float32 inputs, under a comment claiming that is done "for speed and stability". The float16 branch warns about exactly this situation; the float32 branch says nothing.

Why float32 compute is slow

Because there is no MMA 4-bit GEMM kernel for float32 — as the CUDA dispatch heuristic in backends/cuda/ops.py notes itself:

# fp32 has no MMA kernel; pre-sm75 has no MMA kernel; sm75 has fp16 MMA only.
# For all of these, custom only wins in the SIMT range (M<8).
if dtype == torch.float32 or major < 7:
    return M < 8

So float32 takes the custom kernel only for M < 8 and otherwise falls back to unfused dequantize + F.linear.

Measured on an A100 80GB PCIe (bitsandbytes 0.50, torch 2.13.0+cu130), bnb.nn.Linear4bit with nf4 + double quant at Qwen3-4B projection shapes, compute_dtype=float32 vs bfloat16, float32 inputs in both cases. Values are float32 time / bfloat16 time, so >1 means float32 is slower:

shape	M=1	M=8	M=512	M=2048
q_proj (K=2560, N=4096)	0.68x	1.45x	5.64x	8.25x
k_proj (K=2560, N=1024)	0.68x	1.35x	1.57x	4.96x
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

Single-token decoding is unaffected (slightly faster, in fact), which is part of why this is easy to miss. Prefill and training, however, run several times slower than they need to.

This is also easy to hit accidentally in a QLoRA setup: peft's prepare_model_for_kbit_training() upcasts norms and embeddings to float32, so Linear4bit layers receive float32 activations, and if bnb_4bit_compute_dtype was left at its default the whole model computes in float32.

Change

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() rather than in set_compute_type() or __init__(), because that is the only place that sees both the resolved compute_dtype and the real device, and it therefore covers every route into this state — the explicitly-configured one above, the inferred one where set_compute_type() adopts float32 from float32 inputs, and pre-quantized checkpoints loaded via device_map (which bypass Params4bit._quantize).

Two guards keep that safe:

• _fp32_compute_warned, a class-level default shadowed per instance, so the steady-state cost on the hot path is one attribute load and a negation. Deduplication across modules is handled by functools.cache on the log helper, which matters because a 4-bit model has hundreds of these layers (cf. gemm_4bit: blocksize-alignment warning emitted on every call, and now also during training (new in 0.50.0) #2027, where a per-call warning was itself the bug).
• torch.compiler.is_compiling(), so the branch is constant-folded away under Dynamo and never causes a graph break.

The message names both bnb_4bit_compute_dtype and Linear4bit(compute_dtype=...), notes that float32 is the Transformers default, allows for users who genuinely want float32, and says how to silence it.

Also fixes the stale "for speed and stability" comment and a compoute typo.

Testing

tests/test_modules.py::test_4bit_linear_fp32_compute_dtype_warning covers three cases across all available devices: explicitly-configured float32 warns exactly once over 4 layers × 3 forwards; float32 inferred by set_compute_type() also warns once; bfloat16 stays silent. It asserts the warning appears only on cuda.

Run on an A100 80GB, CUDA 13.0, torch 2.13.0+cu130, Python 3.12:

• tests/test_modules.py — 177 passed
• tests/test_linear4bit.py -k "not torch_compile" — 357 passed
• tests/test_linear4bit.py::test_linear4bit_torch_compile, all 32 fullgraph=True + fp32 + cuda variants — 32 passed, confirming no graph break
• pre-commit run --all-files — all 10 hooks pass

I also verified the deduplication directly outside pytest: 200 Linear4bit layers × 5 forwards produces exactly 1 log record.

---

Submitted by Thang Dang, Fujitsu Research of America