Register codebook quant ops

[jerryzh168]

Summary:
Register the codebook quant / dequant ops as custom ops so they can be recongnized after export

Test Plan:
python test/prototype/test_codebook_quant.py -k test_export

Reviewers:

Subscribers:

Tasks:

Tags:

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 13 New Failures

As of commit ebaa5fc with merge base 9516764 ():

NEW FAILURES - The following jobs have failed:

• Code Analysis with Ruff / build (3.9) (gh)
  torchao/swizzle/swizzle_ops.py:51:5: F811 Redefinition of unused swizzle_mm from line 29
• Run Float8 Tests / test (H100, linux.aws.h100, --pre torch torchvision torchaudio --index-url https://download.pytor... / linux-job (gh)
  RuntimeError: Command docker exec -t e5a2167ff4ef8b2c774ed742b25164e727244eab5a650c62f699603312e838d3 /exec failed with exit code 1
• Run Float8 Tests / test (SM-89, linux.g6.4xlarge.experimental.nvidia.gpu, --pre torch --index-url https://download.p... / linux-job (gh)
  RuntimeError: Command docker exec -t 17a487ee9fb0aecf6c028899c313b820f0f938709b934a95652ac2fc005bfb0b /exec failed with exit code 1
• Run Regression Tests / test (CPU 2.3, linux.4xlarge, torch==2.3.0 --index-url https://download.pytorch.org/whl/cpu, cpu) / linux-job (gh)
  RuntimeError: Command docker exec -t 4f025e424e4ca6936137153235d223e3b381630159666dd98b054ef17e673cc0 /exec failed with exit code 1
• Run Regression Tests / test (CPU 2.4, linux.4xlarge, torch==2.4.0 --index-url https://download.pytorch.org/whl/cpu, cpu) / linux-job (gh)
  RuntimeError: Command docker exec -t a5327c0ccc05c7967d85363ac858dae7bef7f5fc4b622b9ea2872861e59ea45a /exec failed with exit code 1
• Run Regression Tests / test (CPU 2.5.1, linux.4xlarge, torch==2.5.1 --index-url https://download.pytorch.org/whl/cpu, cpu) / linux-job (gh)
  RuntimeError: Command docker exec -t 18a70960813622c56c156534ba54c782cd6d177e2915fdaeca7053aea1189155 /exec failed with exit code 1
• Run Regression Tests / test (CUDA 2.3, linux.g5.12xlarge.nvidia.gpu, torch==2.3.0, cuda, 12.1) / linux-job (gh)
  RuntimeError: Command docker exec -t 14fd4442f791721f5b74c124d30cfe1c07a7010ac2a119a276ef4d1db8e271c0 /exec failed with exit code 1
• Run Regression Tests / test (CUDA 2.4, linux.g5.12xlarge.nvidia.gpu, torch==2.4.0, cuda, 12.1) / linux-job (gh)
  RuntimeError: Command docker exec -t a4654e8b7f3b5e433924a2ce5567e597fe58b09eaabdfa54df943823c9d6689d /exec failed with exit code 1
• Run Regression Tests / test (CUDA 2.5.1, linux.g5.12xlarge.nvidia.gpu, torch==2.5.1 --index-url https://download.pytorch... / linux-job (gh)
  RuntimeError: Command docker exec -t 8b4c6575c9f51512dbfcc251184386f719b994c3c1c52ad288683558d6bb9e3a /exec failed with exit code 1
• Run Regression Tests / test-nightly (CPU Nightly, linux.4xlarge, --pre torch --index-url https://download.pytorch.org/wh... / linux-job (gh)
  RuntimeError: Command docker exec -t 4c71fccd69655be1d181c2f70e3573147b8ba7c3c7e60833f9e7bc6d32fbddc8 /exec failed with exit code 1
• Run Regression Tests / test-nightly (CUDA Nightly, linux.g5.12xlarge.nvidia.gpu, --pre torch --index-url https://downloa... / linux-job (gh)
  RuntimeError: Command docker exec -t 93dd081097342fd30f3962cfa5b761abb39b96d81f66ee32f2c0e2f9a525d853 /exec failed with exit code 1
• Run TorchAO Experimental Tests / test-cpu-ops (macos-14) (gh)
  ninja: build stopped: subcommand failed
• Run TorchAO Experimental Tests / test-mps-ops (macos-m1-stable) (gh)
  ninja: build stopped: subcommand failed

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

[metascroy]

nit: say what k is

[jerryzh168]

will update docs after I update the code to support block_size

[metascroy]

IIUC, this does not look like it supports granularity, which we will want.

From what I can tell, k is the idx range, e.g., for 4-bit quantization, k = 16. Each idx=i is mapped to the tensor codebook[i]. So we have 1 codebook/LUT for the tensor that maps indices to tensors.

This seems a bit complicated to me. For CoreML, the default is each idx maps to a scalar (but they also support mapping to a vector). I'm not sure if anyone will need tensor-valued look up values.

But we do want granularity in the sense that we can have one codebook per channel, grouped channel, tensor, etc.

Maybe this is what was originally intended for the block_size (based on https://github.com/pytorch/ao/pull/1299/files/53874a005cb174f764363a7c3a22f653ccf738df#r1870108715), but I understand the code correctly, that's not what got implemented.

[jerryzh168]

I think the scale_block_size in choose_qparams_codebook or the shape of scales in the dequant op is supposed to allow us to control the granularity, the block_sizearg seems to have a different meaning than the block_size in other ops, so we should probably rename it, may guess is the block_size of tensor values that share the same kmeans cluster value.

[jerryzh168]

oh wait, the granularity of codebook is separate, let me take a look again

[metascroy]

Stamping to unblock as discussed