Add Triton kernels for fp8 blockwise quantization and GEMMs #2617
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/2617
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit 14420ef with merge base 0e00df3 ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
facebook-github-bot
added
the
CLA Signed
danielvegamyhre
force-pushed
the
danielvegamyhre/stack/17
branch
from
4c0250f to
89357e5
Compare
danielvegamyhre
added
the
topic: not user facing
danielvegamyhre
force-pushed
the
danielvegamyhre/stack/17
branch
2 times, most recently
from
b2e5d4d to
7f21a73
Compare
| ref_fp32 = ref_fp8.to(torch.float32) | ||
|
|
||
| # Check that the quantized tensors are close | ||
| assert torch.allclose(triton_fp32, ref_fp32, rtol=1e-3, atol=1e-3), ( |
There was a problem hiding this comment.
shouldn't this be bit exact? If it's not bit exact and there no exact reason why not, I wouldn't really trust the triton kernel.
IMO I would just go with torch native kernels for everything for now (since they are the easiest to verify numerical correctness for) except the gemms for now, and leave writing triton kernels for quant of weights/activations as a future thing
There was a problem hiding this comment.
Yes, updated the tests to use torch.equal instead of allclose to assert bitwise equivalence, let me know what you think.
danielvegamyhre
force-pushed
the
danielvegamyhre/stack/17
branch
from
7f21a73 to
5ce55f6
Compare
|
|
||
| sqnr = compute_error(C, C_q) | ||
| min_sqnr = 28.0 | ||
| print(f"blockwise_fp8_gemm_1x128_128x128 ({M},{N},{K}) SQNR: {sqnr}") |
There was a problem hiding this comment.
remove the prints before landing
| ] | ||
|
|
||
|
|
||
| @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") |
| ref_fp32 = ref_fp8.to(torch.float32) | ||
|
|
||
| # Check that the quantized tensors are close | ||
| assert torch.equal(triton_fp32, ref_fp32), ( |
There was a problem hiding this comment.
nit: torch.testing.assert_close(..., rtol=0, atol=0) everywhere
stack-info: PR: #2617, branch: danielvegamyhre/stack/17
danielvegamyhre
force-pushed
the
danielvegamyhre/stack/17
branch
from
5ce55f6 to
14420ef
Compare
|
Thanks for the review @vkuzo, I finished addressing your comments, this is ready for another look. |
Stacked PRs:
Add Triton kernels for fp8 blockwise quantization and GEMMs
GEMMS:
blockwise_fp8_gemm_1x128_128x128for:out = input @ weight.Tgrad_input = grad_output @ weightblockwise_fp8_gemm_1x128_128x1for:grad_weight = grad_output.T @ inputQuantization:
fp8_blockwise_act_quant_lhsfp8_blockwise_act_quant_rhsfp8_blockwise_act_quant_transposed_lhsfp8_blockwise_weight_quant_rhsfp8_blockwise_weight_quant_transposed_rhsTest plan
pytest test/prototype/blockwise_fp8_training/test_blockwise_kernels.pyAttempted usage of DeepGEMM cutlass kernels
Unfortunately the GEMM APIs in @vkuzo's PoC here no longer exist in DeepGEMM. I tried using the new GEMM APIs (
fp8_gemm_ntetc), and:Attempted usage of torch._scaled_mm
I also tried using torch._scaled_mm in torch nightly, and the error messages indicate it does not support groupwise scaled "A" tensor with blockwise scaled "B" tensor, so I went ahead and finished writing these triton GEMMs.
Today, however, I talked to Luca and it seems the error message is inaccurate and it is indeed supported, but will require some changes to scale strides and alignment to adhere to the API requirements.
If we want to make this prototype more performant, I can make these changes and swap out the GEMMs to use torch._scaled_mm (and update the PT core error message to be more accurate).