Open vayuda opened 3 months ago
This is quite cool and I've been thinking along similar lines
I think what we could to do to ship this is in quantization/
merge the pack and unpack functions and then have tests to ensure the the codegen is efficient. In practice you can test that a single kernel is launched by in your tests doing torch.compile(..., fullgraph=True)
- I'm not sure how we can validate that single buffer is used but perhaps @eellison does
And this can be a baseline for smaller dtypes. I'd be specific somewhere in the function names or docs that this is padding-based? Cause conceptually I can imagine another alternative where instead of wasting space you could pack 8 uint3 into 3 unint8 as a more general algorithm but that's finicky enough that we don't have to worry about it right now
Also @mobicham had been asking us for standardizing bitpacking logic so curious on his thoughts too
Thanks @vayuda , very interesting, thanks of sharing!
Normally, bit-unpacking is almost never used in isolation, it's either fused in a dequant kernel or a low-bit matmul kernel. There are two main things to consider while designing a bitpacking logic:
The axis along which quantization is performed: if you quantize along axis=0, and you bitpack along the same axis, the scale/zero can be accessed only once per group. However, if you quantize along axis=1 and you bitpack along axis=0, you'll have to access the scale/zero more than once and it makes dequantization slower. Here are two Triton dequant kernels for both cases I wrote, you can see in the second one, I had to access the zero/scale twice for 4-bit for axis=1, it would be even worse for lower bits: axis=0: https://github.com/mobiusml/hqq/blob/triton/hqq/kernels/triton/dequant.py#L33-L39 axis=1: https://github.com/mobiusml/hqq/blob/triton/hqq/kernels/triton/dequant.py#L65-L71 Since most of the methods quantize along axis=1, it would make sense to have a bitpacking logic that is optimized for that case.
The memory access pattern should be taken into account: if someone writes a Cuda or Triton optimized fused kernel, bitpacking should be structured in a way that can fully take advantage of tensor cores. @jeromeku suggested using interleaved access. Here's a 4-bit bitpacking example using that logic: https://github.com/mobiusml/hqq/blob/triton/hqq/kernels/triton/benchmark.py#L28-L35
@msaroufim do you know by any chance what kind of bitpacking logic is used in tiny_gemm?
@mobicham Thanks for the input. The interleaved accessing is interesting though I'm not really sure what it means to fully take advantage of tensor cores. I think this is something we can iterate on. For now I can create a version that can do row-wise pack/unpack.
As per @msaroufim suggestions, I will place these functions in the api file and write appropriate tests.
Even in relative isolation (without op support) bit packing/unpacking, is still useful for saving memory footprint when storing bool tensors / masks / bitsets:
But of course, more op support is needed for compressed bool tensors / bittensors / bitsets as well...
(Similarly, for some other usecases, it is still useful even when packing/unpacking is not fused into ops where the bottleneck is actually memory efficiency and speed overhead can be tolerated)
In order to support sub-byte dtypes for quantization, I (and many others) believe that it is better to pack these smaller dtypes into existing pytorch dtypes in order to reduce memory bandwidth contention for a bit of increased computation. Here is a preliminary algorithm in pytorch for doing this. It supports many types of conversions as seen in the tests.
Inspecting the compiled Triton code seems promising because it only launches one kernel and one buffer. Here is a snippit