Add BF16 support for int8_vectorwise_quant / LLM.int8 activation quant#1985
Add BF16 support for int8_vectorwise_quant / LLM.int8 activation quant#1985kru2710shna wants to merge 4 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.
kru2710shna
force-pushed
the
feature/int8-bf16-vectorwise-quant
branch
from
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. |
2 participants
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.