Warn once, and only when it matters, for unaligned gemm_4bit - #2031
Open
guptaishaan wants to merge 1 commit into
Open
Warn once, and only when it matters, for unaligned gemm_4bit#2031guptaishaan wants to merge 1 commit into
guptaishaan wants to merge 1 commit into
Conversation
The CUDA gemm_4bit dispatch checked K % blocksize in an elif ahead of the kernel-selection heuristic, so any call with M <= 1536 and a misaligned K warned. The heuristic caps the fused kernel at M <= 512, so for larger M the dequant + F.linear fallback was the intended path all along and the warning was noise. This is what made it show up during training, where M is batch * seq_len. Run the heuristic first and warn only when misalignment is what actually cost us the fused kernel. Move the warning into a functools.cache'd _warn_gemm_4bit_unaligned(K, blocksize) in backends/utils.py so it fires at most once per shape per process, which matters for architectures whose K cannot be aligned (Qwen2.5-VL vision tower has K=3420). The XPU backend had a verbatim copy of the warning and now calls the same helper. No numerical change: every call that used to warn took the fallback then and takes it now.
Comment on lines
+337
to
+358
| @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 | ||
|
|
Member
There was a problem hiding this comment.
I don't think this test is entirely necessary. To me it's noise. I understand we're testing this changed logging behavior, but I still don't think the value is really there. Especially since this is written to run on all devices too. I would prefer to just remove it.
Member
|
Thanks for the PR. On the issue I had commented that I would like the warning to be emitted per Additionally, I was leaning on going to use |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Fixes #2027
In the CUDA
gemm_4bitdispatch theK % blocksize != 0check sat in anelifahead of the kernel-selection heuristic, so any call withM <= _gemm_4bit_custom_max_m(1536) and a misalignedKwarned, even though the heuristic caps the fused kernel atM <= 512and would have taken thedequant + F.linearfallback anyway. That is why the warning shows up during training, whereMisbatch * seq_len. On top of that it was emitted per call, so a model whoseKis inherently misaligned (Qwen2.5-VL vision tower,K = 3420) got one warning per forward.Changes:
backends/cuda/ops.py: run the heuristic first and warn only when misalignment is what actually cost us the fused kernel.backends/utils.py: the warning body moves into afunctools.cached_warn_gemm_4bit_unaligned(K, blocksize), so it fires at most once per shape per process.backends/xpu/ops.py: calls the same helper. The warning text there was a verbatim copy.No behaviour change beyond the warning. Every call that used to warn took the fallback then and takes it now, so numerics are identical.
Verified on an A40 (sm_86), CUDA 12.6, torch 2.13.0+cu126. With the issue's repro (
K = 3420,blocksize = 64, 10 calls) the counts go from 10/10/10 to 1/0/0 for M=1 inference, batched, and aLinear4bittraining forward+backward. New testtests/test_ops.py::Test4bitBlockwiseQuantOps::test_gemm_4bit_unaligned_warningfails on CUDA before the patch and passes after.tests/test_ops.pyand the 4-bit subset oftests/test_functional.pypass, and the pre-commit hooks pass on the changed files.Not verified on hardware I do not have: the XPU change (mechanical substitution, no control-flow change) and ROCm, which shares the CUDA dispatch and can only see fewer warnings after this. Arch coverage is sm_86 only. The fix keys on the heuristic's return value rather than any threshold, so it is arch-independent, but the exact
Mat which warnings stop is not.Thanks to @albertvillanova for the report and the precise diagnosis.