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
Copy link
Copy Markdown
Sorry, something went wrong.
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 thread
tests/test_ops.py
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 | ||
|
|
Copy link
Copy Markdown
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.
Sorry, something went wrong.
Copy link
Copy Markdown
Member
|
Thanks for the PR. On the issue I had commented that I would like the warning to be emitted per (N, K, blocksize) combination. Can you please make this adjustment? Additionally, I was leaning on going to use logger instead of warnings but I'm not super particular on that point. |
Sorry, something went wrong.
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
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_4bit dispatch the K % blocksize != 0 check sat in an elif ahead of the kernel-selection heuristic, so any call with M <= _gemm_4bit_custom_max_m (1536) and a misaligned K warned, even though the heuristic caps the fused kernel at M <= 512 and would have taken the dequant + F.linear fallback anyway. That is why the warning shows up during training, where M is batch * seq_len. On top of that it was emitted per call, so a model whose K is inherently misaligned (Qwen2.5-VL vision tower, K = 3420) got one warning per forward.
Changes:
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 a Linear4bit training forward+backward. New test tests/test_ops.py::Test4bitBlockwiseQuantOps::test_gemm_4bit_unaligned_warning fails on CUDA before the patch and passes after. tests/test_ops.py and the 4-bit subset of tests/test_functional.py pass, 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 M at which warnings stop is not.
Thanks to @albertvillanova for the report and the precise diagnosis.