Repository navigation
Add BF16 support for int8_vectorwise_quant / LLM.int8 activation quant - #1985
kru2710shna wants to merge 6 commits into
Conversation
|
I think it would be worthwhile to update some of the other relevant tests to exercise the new bf16 path. E.g. in |
|
Will do |
Templates int8VectorQuant on T and adds bf16 kernel instantiations plus a cint8_vector_quant_bf16 C ABI entry point, mirroring the existing gemm_4bit_inference_naive fp16/bf16/fp32 pattern. The blockwise absmax reduction now accumulates in float rather than T: required for bf16 to compile cleanly and slightly improves fp16 accuracy (rowStats was already float, so downstream is unaffected). Removes the forced A.to(torch.float16) casts in MatMul8bitLt so bf16 activations quantize natively. Closes bitsandbytes-foundation#1868.
…rning - Restore explanatory comments in kInt8VectorQuant that were dropped during the float-accumulation rewrite (striped-load pattern, sparse-decomp absmax, outlier zeroing). - Remove the 'inputs will be cast to float16' warning in MatMul8bitLt, which no longer applies now that the forced fp16 cast is gone.
e30d08c to
cdd9e23
Compare
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
Extends existing tests to exercise the new bf16 activation-quant path: - test_matmullt: drop the forced fp16 cast so the existing bf16 param actually exercises the bf16 kernel end-to-end. - test_int8_linear_matmul_half: parametrize over fp16/bf16. - test_linear8bitlt_inference and new test_linear8bitlt_forward_dtypes: exercise the Linear8bitLt module forward in both dtypes. Addresses review feedback on bitsandbytes-foundation#1985.
|
Done in #1985 - e888a86. Added bf16 coverage across the relevant tests: test_matmullt: dropped the forced fp16 cast so the existing bf16 param now exercises the real path end-to-end. Note: test_linear8bit.py didn't exist, so I covered test_linear8bitlt.py instead. Also rebased onto latest main. Happy to adjust anything. |
|
There's still a test failure and a couple build issues to resolve.
|
Templates kdequant_mm_int32_fp16 and its launcher on output dtype, adds a cdequant_mm_int32_bf16 C ABI entry point, and dispatches on dtype in the Python backend. Fused bias now works for both fp16 and bf16 when the bias dtype matches the output dtype. Exposes a dtype kwarg on bitsandbytes.functional.int8_mm_dequant so callers can request bf16 output directly instead of relying on a cast afterward. Reverts the earlier test workaround now that int8_mm_dequant genuinely returns bf16 when requested. Addresses review feedback on bitsandbytes-foundation#1985 (previously deferred as a follow-up).
|
Went ahead and implemented this properly rather than just patching the test. int8_mm_dequant now natively supports bfloat16 output end-to-end: Templated kdequant_mm_int32_fp16 and its launcher on output dtype, added a bnb_bfloat16 instantiation and a cdequant_mm_int32_bf16 C ABI entry. This should close the last known gap for bf16 support in this PR. |
| template void dequant_mm_int32_fp16<half>( | ||
| int* A, float* rowStats, float* colStats, half* out, half* bias, int numRows, int numCols, bnb_stream_t stream | ||
| ); | ||
| template void dequant_mm_int32_fp16<bnb_bfloat16>( | ||
| int* A, float* rowStats, float* colStats, bnb_bfloat16* out, bnb_bfloat16* bias, int numRows, int numCols, | ||
| bnb_stream_t stream | ||
| ); |
There was a problem hiding this comment.
Now that it supports bf16 and fp16, the naming on this with the _fp16 suffix is a little confusing. It may make more sense to just rename it to dequant_mm_int32<T>.
| torch.testing.assert_close(out_accelerator_2, out_accelerator, rtol=1e-8, atol=1e-8) | ||
|
|
||
|
|
||
| @pytest.mark.parametrize("device", get_available_devices(no_cpu=True)) |
|
Hi @kru2710shna, There still remains build failures on CUDA < 12.4 and on ROCm. Additionally there are test failures where the build succeeded. |
- kInt8VectorQuant: CUDA < 12.4 only declares the bf16 __ldcs overload for sm_80+. Issue the streaming load through the half overload and reinterpret the bits. - pythonInterface.cpp: use bnb_bfloat16 instead of __nv_bfloat16 so ROCm builds. - Rename kdequant_mm_int32_fp16/dequant_mm_int32_fp16 to kdequant_mm_int32/dequant_mm_int32 now that they are templated on the output type (review feedback). - int8_mm_dequant (CUDA): accept any output dtype again. fp16/bf16 are written natively; other dtypes (fp32) are computed in fp16 and cast, as before. - MatMul8bitLt/MatMul8bitFp: only cast to fp16 when the input is not fp16/bf16, so fp32 inputs keep working and keep their fp32 output dtype. Backward keeps bf16 grads in bf16. - Int8Params: quantize bf16 weights natively instead of via fp16. - Add test_int8_mm_dequant_dtypes covering fp16/bf16/fp32 output and fused/unfused bias.
|
Fixed the remaining failures in 42b08e1: CUDA < 12.4: __ldcs(const __nv_bfloat16*) is only declared for sm_80+ before 12.4 (I reproduced this against the 12.1 headers on sm_75). The kernel now does the streaming load through the half overload, which works on every arch, and reinterprets the bits as bf16. ROCm: the bf16 C ABI entry points now use bnb_bfloat16 instead of __nv_bfloat16. Test failures: fp32 inputs and int8_mm_dequant(dtype=torch.float32) work again. They go through fp16 for quantization and dequant as before, so test_matmullt[fp32] and test_int8_scaled_mm[fp32] pass. Renamed the kernel and launcher to kdequant_mm_int32 / dequant_mm_int32 as you suggested. Added test_int8_mm_dequant_dtypes (fp16/bf16/fp32 output, with the bias either fused or added afterwards). |
Templates int8VectorQuant on T and adds bf16 kernel instantiations plus a cint8_vector_quant_bf16 C ABI entry point, mirroring the existing gemm_4bit_inference_naive fp16/bf16/fp32 pattern. The blockwise absmax reduction now accumulates in float rather than T: required for bf16 to compile cleanly and slightly improves fp16 accuracy (rowStats was already float, so downstream is unaffected). Removes the forced A.to(torch.float16) casts in MatMul8bitLt so bf16 activations quantize natively.
Closes #1868.