Feat: Implementation of the DeepSeek blockwise quantization for fp8 tensors #1763
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/1763
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit 0af2b9b with merge base 137b079 ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
facebook-github-bot
added
the
CLA Signed
|
Thanks for running the tests. I have two questions regarding the errors:
|
Can you clarify what you mean? Are tests failing in CI due to a missing triton installation? That shouldn't be happening, please share the link/logs if so.
We just use helpers which skip tests if GPU architecture is not at least SM 89: Line 619 in f478692 You can find examples in the float8 tests (example). |
Indeed, they are. It looks like only the CPU runs are failing. I presume that bitsandbytes might not install triton when no GPU is available (I might be missing something there). Here is an instance of a failing log: https://github.com/pytorch/ao/actions/runs/13484452669/job/37730985419?pr=1763#step:14:1276
Thank you for the hint, I've locally updated the code accordingly 👍 |
| W_q, W_s = fp8_blockwise_weight_quant(W, block_size, dtype) | ||
| output_blockwise = blockwise_fp8_gemm(A_q, A_s, W_q, W_s) | ||
|
|
||
| quantize_(lin, int8_dynamic_activation_int4_weight()) |
There was a problem hiding this comment.
why is int8_dynamic_activation_int4_weight being used here?
There was a problem hiding this comment.
Thank's for noticing it. I was aiming for a static W4A8 quantization and I overlooked that it was dynamic. I will try to address this within the week.
Also @Degnel you should skip tests requiring triton if CUDA is not available. |
|
@Degnel thanks for your work on this, i ran the tests and it looks like your blockwise fp8 gemm test is failing due to quantization error |
Thanks for pointing that out! I had also noticed the issue, and I think I was just a bit too harsh with the threshold. I'll increase it to make it more reasonable. That said, I'll still double-check the calculations manually to ensure everything is mathematically correct. |
|
@danielvegamyhre I believe that everything should be alright except for the PR Label Check (I'm not sure if I have the required rights to edit this). The test-mps-ops (macos-m1-stable) failed, but I think that the merge will fix it as it seems to be a newly introduced test. |
danielvegamyhre
added
the
topic: new feature
|
The test-mps-ops (macos-m1-stable) failed once again. I've seen other recent PRs both succeeding and failing this test (due to the same missing package 'importlib_metadata'). I don't think this is related to the code I wrote, but I might be missing something. Please, let me know if you have any insights. |
|
the test mps is unrleated, re-running tests |
|
It seems like the new PRs are not failing anymore due to the macOS tests. Maybe we should try to rerun it here :) @danielvegamyhre @drisspg |
|
Sorry, could you do 1 more rebase to kick back off ci |
- fp8 triton gemm - quant, dequant and linear utils - time & precision benchmarks - basic tests
Degnel
force-pushed
the
feat/blockwise_fp8_quant_triton_gemm_ker
branch
from
e8edea9 to
e41457c
Compare
No problem, it should be ok |
|
Thank you @drisspg I've made the linting |
|
cc @danielvegamyhre can you carry this across the finish line |
There was a problem hiding this comment.
Great progress @Degnel and sorry for the delay! Just did another pass and left a couple comments.
| W_q, W_s = fp8_blockwise_weight_quant(W, block_size, dtype) | ||
| output_blockwise = blockwise_fp8_gemm(A_q, A_s, W_q, W_s) | ||
|
|
||
| qact = _int8_symm_per_token_reduced_range_quant_cutlass(A) |
There was a problem hiding this comment.
Why are int4 and int8 quant being used here?
Note these no longer exist in torchao (_int4_symm_per_token_quant_cutlass and _int8_symm_per_token_reduced_range_quant_cutlass) so they should be removed/changed but I'm unclear why int4/int8 quant is being used in the first place, can you clarify?
There was a problem hiding this comment.
To be transparent, this code was written a few months ago, and I don't fully recall the rationale behind using those specific quantization modes. At the time, the goal was to establish a baseline for comparison across different quantization strategies, and int4/int8 quantization options seemed like a good starting point. Looking back, I agree it was probably a poor choice. I've removed the deprecated imports.
…:Degnel/ao into feat/blockwise_fp8_quant_triton_gemm_ker
|
@Degnel I get an import error when trying to run the benchmarks, can you take a look? |
|
Interesting, I see solid speedup vs fp16 for small M but it falls off as M increases, to the point where fp8 blockwise is actually 4x slower than fp16 for larger shapes, which is the opposite of what I would expect... this on H100s for context. Maybe block sizes in autotuner are not optimal for shapes of this size? What do you think @Degnel @drisspg ? I suppose we can merge this is an initial start but we should profile this and fix the perf for larger shapes. |
Thanks @danielvegamyhre. My first intuition was that it might be due to a hidden tensor copy, but the code looks fairly straightforward and I don’t immediately see anything problematic there. That said, the original paper does mention some know inefficiencies: “partial results will be copied from Tensor Cores to CUDA cores”. However, the sizes they benchmarked seem to be at least larger than the point where we start seeing the drop in efficiency. I’ll try to dig into this over the next few days and see if I can pinpoint what’s going on. |
This PR is the first step towards addressing issue #1594. It includes the following implementations:
If the code is validated, it would be great to bench it on H100.