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. ❌ 1 New FailureAs of commit 10c9436 with merge base c9b9adc ( NEW FAILURE - The following job has failed:
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
the-tuning-machine
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 |
| 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. |
> the autotuner was optimizing based only on small M sizes at the beginning of the benchmark > added a `M_bucket` key to the autotuner to enable tuning based on similar M sizes > added `128` to the `BLOCK_SIZE_M` configuration, which improves performance for large M values > launcher now takes `block_size` into account (although using `block_size=128` is recommended for best performance)
…:Degnel/ao into feat/blockwise_fp8_quant_triton_gemm_ker
|
@danielvegamyhre @drisspg I finally got the time to dive into the problem. On an L4, the default DeepSeek parameters are already quite solid. I eventually found out that adding 128 to the BLOCK_SIZE_M configuration, leads to some performance gains for large M (before that change, computing the matmul in fp16 was actually faster than the optimized path for M = 512 even with the key fix). |
|
@Degnel Thanks for fixing the autotune config for M and verifying the perf. I want to run everything locally to double check but I'm getting an import error again trying to run the benchmark. This happened previously and you fixed it, did the fix get undone somehow? |
|
@danielvegamyhre I just checked, and I don’t see 'from torchao.prototype.blockwise_fp8 import blockwise_fp8_gemm' anywhere in the latest version of the code (https://github.com/pytorch/ao/pull/1763/files). Maybe you're using an older version? Let me know if that’s not the case. If needed, the correct import should be: 'from torchao.prototype.blockwise_fp8.blockwise_quantization import blockwise_fp8_gemm'. |
Re-checked out the PR and am getting a different import error when running the benchmark now: This particular line does appear in the latest PR. thanks for sticking with this, looks good to me once we get this fixed! |
Hi @danielvegamyhre, I think that you should run 'pip install -e .' at the root of your project. It looks like Python does not recognize the newly added module. |
|
Latest benchmarks, |
|
Benchmarks on H100, looks like for high values of M we get a slight slowdown of a few % at most but nothing like the 0.38 observed on A100 by @Degnel. We can include these in the prototype README as a next step. Thanks for adding this @Degnel! |
|
Thanks a lot @danielvegamyhre for your assistance and for running the benchmarks on the H100! Would you like me to go ahead and add the blockwise quantization to the TorchAO API now, or do you think it's too early given the current performance results? Another next step could be to try make a bench on a QAT task (this could be an occasion to update the README). Let me know what you think would be relevant. |
Sorry for the late reply, I just saw this - I think a good next step would be to make float8 blockwise linear differentiable via an autograd func so it can be used for training, and adding some training benchmarks. Here is an example with triton kernels that uses autograd functions to make it differentiable for training. Do you want to work on this? I would be excited to land this in torchao |
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.