Feat/blockwise fp8 quant #1668
Feat/blockwise fp8 quant #1668
Conversation
- first implementation of the DeepSeek blockwise quantizer (not fully fonctionnal) - amax has been unpdated - 2 more quantisation recipes has been added - a couple of things here and there to make it consistent
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/1668
Note: Links to docs will display an error until the docs builds have been completed. ❗ 1 Active SEVsThere are 1 currently active SEVs. If your PR is affected, please view them below: This comment was automatically generated by Dr. CI and updates every 15 minutes. |
facebook-github-bot
added
the
CLA Signed
|
could you share what gemm kernel you plan to use in this PR? I think a good first step here is to have a fast gemm. we have an issue tracking this here: #1594 |
|
this might be a good place to start: https://huggingface.co/deepseek-ai/DeepSeek-V3/blob/1d044fd82b15f1cedb197a288e50cc96a2c27205/inference/kernel.py#L63 |
|
Overall it would be great to be able to support this recipe in torchao. I think having a gemm with compelling performance that supports 128x1 and 128x128 scaling is something we need first, with benchmarks comparing to other recipes such as rowwise scaled, etc. |
|
Relevant PR in SGLang that adds the triton kernels - sgl-project/sglang#2575 (thanks to @HandH1998). I think it makes sense to add this as a starting point to torchao. |
Feat: Implementation of the DeepSeek blockwise quantization for fp8 tensors
WARNING: The code has been tested on the following files:
pytest test/float8/test_base.py
pytest test/float8/test_compile.py
pytest test/float8/test_numerics_integration.py
However, tests have not been performed on the following files due to limitations (Triton is unavailable on Windows and I don't own an NVIDIA GPU):
./test/float8/test_fsdp.sh
./test/float8/test_dtensor.sh
python test/float8/test_fsdp2/test_fsdp2.py