int8 dynamic quant + bsr support

[jcaip]

This PR, based on #891 adds in int8 dynamicquant + bsr support. In this PR is a pretty large refactor of the superblock code, mainly removing duplicated code. There's still some work to be done but this PR is already getting pretty big, so want to check it in.

Changes:

• Use i8i8 -> bf16 matmul to maintain accuracy
• Added a block sparse layout type to AffineQuantizedTensor + check/impl.
• Cleaned up benchmark.py script and add a single line benchmark.sh file for acceleration numbers
• Updated eval.py and added a single line evaluate.sh file for accuracy numbers
• Lots of lint formatting and README updates
• torch.compile now working

TODO:

• move sparse layouts to their own folder
• promote blocksparse.py from prototype and add some tests
• move evaluate.py/benchmark.py out of superblock and into the top level folders?
• move superblock.py into training folder?

---

This PR is a prototype of adding int8 dynamic quant with BSR support.

Currently it works, but is not performant. (1.6 img/s vs 2.0 baseline) on Vit-H

This is because the block sparse mm is added as a custom op, since we cannot trace through the creation of a sparse bsr tensor in torch.compile:
Instead, we created a custom op that takes in the data tensors of a bsr tensor (crow_indicies, col_indicies, values) and wrapper subclass that holds these values. See here

This causes an issue for the quantization workflow, since we cannot fuse across the custom op. It's pretty much the same issue we face with 2:4 sparsity, where we materialize the full int32 intermediary matrix, which eats into our composed quant + sparse speedups.

Normally we avoid materializing the int32 intermediary matrix by fusing the dequant into the matmul, but we can't do so here because we use a custom op. I think the long term solution here is to just rip out the BSR triton code in core and have it in AO, in that case we can change the function to take in the data tensors themselves instead of a sparse bsr tensor, avoiding this need for a custom op.

Good news though, accuracy is not degraded at all when composing int8 and block sparsity

cc @cpuhrsch

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/821

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit a2519d8 with merge base 7dff17a ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[cpuhrsch]

cc @pearu

We could do epilogue fusion in core if we're adding _int_bsr_dense_addmm as a Template to https://github.com/pytorch/pytorch/tree/196748d49193eed2b000aa99b2c7bb1bff576d35/torch/_inductor/kernel .

We can do that similar to how we did it for _int_mm: https://github.com/pytorch/pytorch/pull/111125/files#diff-e1a7b1be52a02ad5595c2f38c1e5c1ac3086fd49b6638482ae192249f070f4dc

[jerryzh168]

@jcaip do you want to move these to sparsity folder? now we should be able to do everything outside of affine_quantized_tensor after #783

(meaning register layout_type/layout tensor and dispatch/impl for quantized linear

[jcaip]

No, not for this PR at this time. I think we should resolve the perf issues before we merge.

[cpuhrsch]

I wonder if a sufficiently large batch size would eventually show a gain even without epilogue fusion.

[jcaip]

cc @cpuhrsch It's possible, the ImageNet benchmark is not the kindest benchmark either since we can't quantize the attention layers as well. So we may be able to see perf gains with SAM as well. I added this as a hackathon project for CUDA-MODE so maybe we can pick up a contributor there.

[pearu]

When running

python benchmark.py --model vit_h_14   --batch-size 256   --sparsity-linear 0.8   --sp-linear-tile-size 64 --bsr 64 --sparsity bsr --quantization

notice the warning messages:

torch/sparse/_triton_ops.py:795: UserWarning: bsr_dense_addmm uses non-optimal triton kernel parameters for M=5120 K=1280 N=65792 Ms=64, Ks=64 beta=0 alpha=1
  warn_once(
torch/sparse/_triton_ops.py:795: UserWarning: bsr_dense_addmm uses non-optimal triton kernel parameters for M=1280 K=5120 N=65792 Ms=64, Ks=64 beta=0 alpha=1

When using the tuned parameters, the performance increase is about 0.1 img/s.

I'll create a torch PR with the corresponding tuned parameters shortly.

[pearu]

Suggested change

	return bsr_dense_mm(weight_bsr, A).t().contiguous()
	return bsr_dense_mm(weight_bsr, A).t()

That should increase the performance of

benchmark.py --model vit_h_14   --batch-size 256   --sparsity-linear 0.8   --sp-linear-tile-size 64 --bsr 64 --sparsity bsr --quantization

by ~0.9 img/s.

[pearu]

Suggested change

	tmp_t = tmp.t().contiguous()
	tmp_t = tmp.t()

to increase the performance about 0.2 img/s.

[pearu]

TODO: move scaling operations to bsr_dense_addmm kernel.

[pearu]

pytorch/pytorch#136104

[jcaip]

cc @pearu @cpuhrsch

I've pulled in your changes to this PR here, seeing some nice speedups! Unfortunately I think I'm running into some saturation issues when doing the int8 int8 -> int8 BSR matmul.

This is causing poor accuracy when running BSR with quantization, evaluate.sh should print out the results. I'm seeing 0 accuracy when running with BSR+ quant. I can reproduce this is the test too, when I bump up the matrix sizes, the results I'm seeing for the quantized mm is quite different from the reference result:

We ran into something similar with cuSPASRELt and we solved that by adding support for i8i8->bf16 matmul. Would it be possible to add something similar to the BSR kernels? Something like aout_dtype option would work.

[pearu]

Would it be possible to add something similar to the BSR kernels? Something like aout_dtype option would work.

We have it in terms of out argument. In blocksparse_int_addmm, try changing

out = A.new_empty(original_batch_dims_broadcasted + (M, N))

to

out = A.new_empty(original_batch_dims_broadcasted + (M, N), dtype=torch.int32)

that should resolve the saturation problem. I did't use bfloat16 as out dtype because triton dot accumulator uses int32.

Notice that we don't have tuned parameters for bsr_dense_addmm(<int8>)->int32, so the performance is expected to be lower. Using --tune-kernel-params will not be helpful either because optimize_bsr_dense_addmm assumes that out.dtype is int8 if dtype is int8.
There is an easy fix though: we'll need to add out_dtype parameter to optimize_bsr_dense_addmm. I'll cook up a PR for this..

[jcaip]

Nice, that makes sense.I still run into numerical issues with torch.int32, but passing in torch.bfloat16 seems to fix those :)

i8i8->bf16 eval:

Test:  Acc@1 77.087 Acc@5 105.867
vit_b_16,256,bfloat16,bsr,64,0.8,True,77.08733974358974

i8i8->i8/32 eval:

Test:  Acc@1 0.495 Acc@5 1.945                                                                                       
vit_b_16,256,bfloat16,bsr,64,0.8,True,0.4947916666666667 

Will spend some time benchmarking perf + closing out the PR tomorrow, but preliminary seeing 436.59 ms for bf16 vs 362.064 ms for int8.

[pearu]

There is an easy fix though: we'll need to add out_dtype parameter to optimize_bsr_dense_addmm. I'll cook up a PR for this..

Done in pytorch/pytorch#136626

[pearu]

I still run into numerical issues with torch.int32

Can you provide a simple reproducer?

What is the size of inputs to bsr_dense_addmm? In the worst case scenario (all entries have value 2**7), the maximal intermediate size for int8 @ int8 -> int32 is about 2**31 / 2**14 = 131072 when int32 accumulator would not overflow.

[jcaip]

@pearu

I don't think it is a saturation issue, but rather a numerical one. I think the outputs of the layer are small values to close 0. When we run with an int intermediately, these all values all get quantized to 0. But this leads to the model spitting out all zeros which ruins the model accuracy.

If you change the dtype to int32 in my branch, you can run sh evaluate.sh after setting IMAGENET_PATH to reproduce.