[WIP] Codebook quantization flow

[DerekLiu35]

This PR adds codebook quantization flow per #1195

Usage

import torch
from torchao.prototype.quantization.codebook.codebook_quantized_tensor import CodebookQuantizedTensor

input_tensor = torch.randn(1024, 1024,  device='cuda')

block_size = (1, 1)
code_dtype = torch.uint4

quantized_tensor = CodebookQuantizedTensor.from_float(input_tensor, block_size, code_dtype)

dequantized_tensor = quantized_tensor.dequantize()

ToDo

• make fit_kmeans faster. Right now it takes >1 hour if you try to quantize a 1B model.

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

✅ You can merge normally! (1 Unrelated Failure)

As of commit b938a7c with merge base 46b8796 ():

BROKEN TRUNK - The following job failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• Run Regression Tests / test-nightly (CUDA Nightly, linux.g5.12xlarge.nvidia.gpu, --pre torch --index-url https://downloa... / linux-job (gh) (trunk failure)
  test/integration/test_integration.py::TestSubclass::test_int8_dynamic_quant_subclass_api_5_cuda

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

[jerryzh168]

thanks for the contribution! yeah "> 1hour" seems a bit too slow, any ideas to speedup?

[jerryzh168]

also after this is done, it would useful if you can add codebookquant to generate.py (

ao/torchao/_models/llama/generate.py

Line 209 in b714026

if quantization:

) and eval (

ao/torchao/_models/llama/eval.py

Line 71 in b714026

if quantization:

) to test the e2e model performance and accuracy

[DerekLiu35]

thanks for the contribution! yeah "> 1hour" seems a bit too slow, any ideas to speedup?

I think

• For block_size = (1, 1), It's similar to nf4tensor, so we can use absolute distance for scalars instead of euclidean distance
• We could also decrease max_iter from 1000 to 200 for fit_kmeans but this would increase quantization error.

[DerekLiu35]

changed max_iter from 1000 to 200 added codebook to eval and generate
I also added scales with group size 64 since when I ran eval codebook with torch.uint4 it was getting high perplexity, the perplexity is still ~200 even with scales with group size 64

[jerryzh168]

changed max_iter from 1000 to 200 added codebook to eval and generate I also added scales with group size 64 since when I ran eval codebook with torch.uint4 it was getting high perplexity, the perplexity is still ~200 even with scales with group size 64

thanks, why the perplexity is so high? we get around 12/13 when using int4wo-64 on llama2: https://github.com/pytorch/ao/tree/main/torchao/quantization#cuda-backend

[DerekLiu35]

thanks, why the perplexity is so high? we get around 12/13 when using int4wo-64 on llama2: https://github.com/pytorch/ao/tree/main/torchao/quantization#cuda-backend

I found that it was because I was setting scales to the norm of each scale group instead of max of each scale group. If you set scales to be max of each scale group wikitext perplexity is ~11.6 for
llama-3.2-3B.
Though, AQLM initializes their scales as the norm of each scale group instead of max, but I guess they only do vector quantization and no scalar quantization.

[jerryzh168]

Thanks if both performance and accuracy are reasonable, I think the main thing is to add a section for codebook quant: https://github.com/pytorch/ao/blob/main/torchao/quantization/README.md#uintx-quantization and show the perplexity and token/s result from eval.py and generate.py

[jerryzh168]

block_size is a general arg that allows people to do all kinds of granularities:

ao/torchao/quantization/quant_primitives.py

Lines 277 to 287 in 53d2486

	Note:
	How can block_size represent different granularities?
	let's say we have a Tensor of size: (3, 3, 10, 10), here is the table showing how block_size represents different
	granularities:
	granularity type | block_size
	per_tensor | (3, 3, 10, 10)
	per_axis (axis=0) | (1, 3, 10, 10)
	per_axis (axis=1) | (3, 1, 10, 10)
	per_group (groupsize=2) | (3, 3, 10, 2)
	per_group (groupsize=2) for axis = 3 | (3, 3, 2, 10)

, you can also refer to

ao/torchao/quantization/quant_primitives.py

Lines 367 to 375 in 53d2486

	shape_for_reduction, reduction_dims = _get_reduction_params(
	block_size, input.size()
	)
	original_shape = input.shape
	input = input.view(shape_for_reduction)
	shape_after_reduction = shape_for_reduction
	for i in reduction_dims:
	shape_after_reduction[i] = 1
	scale = scale.view(shape_after_reduction)

for some helper functions that helps to make the shape correct

is there a reason why it's assumed to be 2d here?

[jerryzh168]

I guess it's fine we start with 2d for now as well, but would be good to add an assert and create an issue for further development

[DerekLiu35]

I made it 2d because that is how it was implemented in AQLM code. The code will be a little more complicated, but it should be possible to generalize to more than 2d.

[DerekLiu35]

benchmarks were run on a single NVIDIA-A6000 GPU.

Model	Technique	wikitext-perplexity	Tokens/Second	Memory Bandwidth (GB/s)	Peak Memory (GB)	Model Size (GB)
Llama-3-8B	Base (bfloat16)	7.590	32.36	485.71	16.19	15.01
	codebook-4-64	9.533	1.73	8.62	23.11	4.98
Llama-3.1-8B	Base (bfloat16)	7.713	32.16	482.70	16.35	15.01
	codebook-4-64	10.095	1.73	8.63	23.11	4.98

Seems slow, might need custom kernels

[jerryzh168]

is this needed? this might be the reason why it's slow, what do you mean by index?

[DerekLiu35]

I think so. I do indexing in dequant = codebook[codes] in dequantize_codebook. I got an error when I tried doing codebook[codes] when codes was uint8.

[jerryzh168]

what is the error? might be easy to support this in pytorch I feel

[jerryzh168]

only int and long are supported right now: https://github.com/pytorch/pytorch/blob/6e203ae6deaceb370e497bd50f2d02e894f5e9cc/aten/src/ATen/Dispatch.h#L801

[DerekLiu35]

UserWarning: indexing with dtype torch.uint8 is now deprecated, please use a dtype torch.bool instead.
IndexError: The shape of the mask [2] at index 0 does not match the shape of the indexed tensor [3, 3] at index 0.

[jerryzh168]

OK, let's add a TODO here to follow up and land this for now, is perplexity number expected here? it looks like it's slightly worse than int4wo: https://github.com/pytorch/ao/blob/main/torchao/quantization/README.md#cuda-backend

[DerekLiu35]

I think perplexity is expected to be slightly worse than asymmetric int4wo (with scales and zero points) because I only add scales. But it should usually be better than symmetric int4wo (only scales, no zero points), though I initialize the centroids in kmeans randomly so it could be worse if centroids have bad initialization.

[jerryzh168]

I think model level benchmark is fine for now, maybe add some unittest to test some basic functionality like getting codebook, quantize, dequantize before landing?

[jerryzh168]

also for model level, is it possible to repro some accuracy result from AQLM: https://github.com/Vahe1994/AQLM/tree/main

[DerekLiu35]

I tested if I could convert an AQLM quantize model to my implementation as a sanity check and it seems like it works https://gist.github.com/DerekLiu35/c1cb9594c515e92c64762cbc8d087f7a.
Though, I get ~300 perplexity for llama-3-8b when I do similar setting as AQLM (block_size=(1, 8), code_dtype=torch.int32 (which makes codebook_size = 2**16), and scale_block_size=input.shape[1]). Maybe the extra tuning is necessary for lower perplexity?

[jerryzh168]

Thanks! So the representation can be verified, but we are not sure how to repro the accuracy, and in original AQLM they'd need full model finetuning to get a good accuracy, seems like a good next step, if you are interested in improving this further

[jerryzh168]

looks like a large drop, we could use a larger dtype, e.g. torch.uint8, for test I think, so we can get something closer

[jerryzh168]

we typically use SQNR:

ao/torchao/quantization/utils.py

Line 50 in 46b8796

def compute_error(x, y):

, numbers greater than 20 or 30 are reasonable

[DerekLiu35]

Thanks! So the representation can be verified, but we are not sure how to repro the accuracy, and in original AQLM they'd need full model finetuning to get a good accuracy, seems like a good next step, if you are interested in improving this further

Yeah, I'd definitely be interested in trying to implement AQLM to improve accuracy.

[jerryzh168]

nit: it's fine to just rely on sqnr for error checking btw.

[jerryzh168]

Thanks for the contribution and addressing all the comments!

[jcaip]

@mostafaelhoushi

To overload linear:

@implements_torch_function(nn.functional.linear)
def function_linear(tensor, *args, **kwargs):
    breakpoint()
    // torch.ops.torchao.tinygemm(args)
    return tensor.detach()