Add Int4PlainInt32 Tensor #2845

liangan1 · 2025-08-22T01:56:30Z

This PR is used to enable the Int4XPUTensorIntZP. The pacing format name is "int4_xpu_int_zp"
Testcase:
bash python test/quantization/quantize_/workflows/int4/test_int4_xpu.py

pytorch-bot · 2025-08-22T01:56:34Z

🔗 Helpful Links

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

📄 Preview Python docs built from this PR

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

❌ 4 New Failures

As of commit 282f1a8 with merge base 3bf21d0 ():

NEW FAILURES - The following jobs have failed:

Run Regression Tests / test (CPU 2.8, linux.4xlarge, torch==2.8.0 --index-url https://download.pytorch.org/whl/cpu, cpu) / linux-job (gh)
test/quantization/quantize_/workflows/int4/test_int4_plain_int32_tensor.py::Int4PlainInt32Tensor::test_module_path_float16
Run Regression Tests / test (CUDA 2.8, linux.g5.12xlarge.nvidia.gpu, torch==2.8.0, cuda, 12.6) / linux-job (gh)
test/quantization/quantize_/workflows/int4/test_int4_plain_int32_tensor.py::Int4PlainInt32Tensor::test_module_path_float16
Run Regression Tests / test-nightly (CPU Nightly, linux.4xlarge, --pre torch --index-url https://download.pytorch.org/wh... / linux-job (gh)
test/quantization/quantize_/workflows/int4/test_int4_plain_int32_tensor.py::Int4PlainInt32Tensor::test_module_path_float16
Run Regression Tests / test-nightly (CUDA Nightly, linux.g5.12xlarge.nvidia.gpu, --pre torch --index-url https://downloa... / linux-job (gh)
test/quantization/quantize_/workflows/int4/test_int4_plain_int32_tensor.py::Int4PlainInt32Tensor::test_module_path_float16

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

jerryzh168 · 2025-08-26T00:11:03Z

torchao/quantization/quantize_/common/packing_format.py

+    "int4_xpu_int_zp is referring to the format used by int4 weight-only quantization on XPU with int zero point, which is a groupwise quantization format."
+    INT4_XPU_INT_ZP = "int4_xpu_int_zp"


please don't include int4 and xpu in the name, can you name this in terms of of how the quantized data is packed?

The int4 weight xpu is a plain format tensor according to this doc, it just pack 2 int4 weight elements in a byte and then store the 4*int4 as int32. So I change it to the plain.

I see, we have plain that stores 2*int4 as int8, can you reuse it or would need a new one? https://github.com/pytorch/ao/blob/main/torchao/quantization/quantize_/workflows/int4/int4_tensor.py

@liangan1 can you use PLAIN_INT32 for packing_format, and rename things accordingly (tensor subclass, files etc.)

Thanks @jerryzh168. I have added PLAIN_INT32 to be used by the xpu int4. Per my understanding, the packing format should be a dispatch policy to select the right tensor subclassing and a tensor subclass should cover a specific quantization recipe. So I suppose I should keep the current tensor name for int4 xpu.
In this PR, we just want to enable the int xpu with int zp domain. The current oneDNN backend can not support the float zp as CUDA/CPU backend and the feature is WIP. I plain to reuse this packing format in the future and dispatch the tensor with the zero point domain information.

you can reuse the packing format and the tensor for float32 zero_point as well in the future I think, but today we structure tensor subclass by: dtype + packing_format, so Int4PlainInt32 might be better

Done. change it to Int4PlainInt32. pls help to review again.

torchao/quantization/quantize_/workflows/int4/int4_xpu_tensor.py

jerryzh168 · 2025-08-28T00:03:46Z

torchao/quantization/quantize_/workflows/__init__.py

@@ -14,6 +14,9 @@
 from .int4.int4_tensor import (
    Int4Tensor,
 )
+from .int4.int4_xpu_tensor import (
+    Int4XPUTensorIntZP,


nit: we should rename the tensor names etc. as well

Per my understanding, the packing format should show the weight packing format, but I also want add the zp domian information in this tensor, so keep the tensor name as Int4XPUTensorIntZP.

the packing format should contain a description of how the quantized weights are laid out in memory, it should not mention XPU, and also not INTZP since these are not relevant to how the quantized weights laid out.

can you change it to something like Int4PlainInt32 since PLAIN_INT32 is the packing_format?

Int4PlainInt32 is good for me. But I still have some concern about the relation between packing_format and subclassing tensor? it must be 1 x 1?

I think so, it's 1x1, we are also going to do a follow up refactor to have Int4PackingFormat instead having a global PackingFormat, since there is not much reuse

theoretically, 1 packing format can correspond to multiple tensors with different dtypes, but we found that there is not much reuse of packing_format, that's why we want to move it to local to each dtype (int4, float8, intx)

what's the concern for this?

Thanks. Your explanation is clear enough. No more concern now.

torchao/quantization/quantize_/workflows/int4/int4_xpu_tensor.py

jerryzh168 · 2025-08-29T02:02:27Z

torchao/quantization/quantize_/workflows/int4/int4_plain_int32_tensor.py

+        eps = 1e-6
+        scale_dtype = None
+        zero_point_dtype = torch.int32
+        scale, zero_point = _choose_qparams_affine(


nit: we can use the official op?

ao/torchao/quantization/quant_primitives.py

Line 1194 in f02354d

def choose_qparams_affine(

same for quantize_affine

jerryzh168 · 2025-08-29T02:03:21Z

test/quantization/quantize_/workflows/int4/test_int4_plain_int32_tensor.py

+    return Int4WeightOnlyConfig(
+        group_size=group_size,
+        packing_format="plain_int32",
+        zero_point_domain=ZeroPointDomain.INT,


nit: we don't need this anymore I think, also we want to remove ZeroPointDomain in the future

Removed. But I have a question that how to selelct the int zp domain for user if there is no this param?

jerryzh168 · 2025-08-29T02:03:50Z

torchao/quantization/__init__.py

    "Int4PreshuffledTensor",
    "Int4MarlinSparseTensor",
-    "IntxOpaqueTensor",
+    "IntxOpaqueTensor",	   


nit: extra spaces?

jerryzh168 · 2025-08-29T02:04:19Z

torchao/quantization/quantize_/workflows/int4/int4_plain_int32_tensor.py

+    int4 weight-only quantization on XPU with oneDNN as backend (groupwise quantization only)
+
+    Tensor Attributes:
+        qdata: packed int4 weigh, always viewed as a 2D (N, K/2) tensor


nit: would be good to callout the dtype here, since it's important in this case

jerryzh168 · 2025-08-29T02:05:18Z

torchao/quantization/quantize_/workflows/int4/int4_plain_int32_tensor.py

+
+    Non-Tensor Attributes:
+        block_size: the block size for quantization, representing the granularity, for groupwise quantization, will have block_size (1, group_size).
+                    we only support group_size = 32/64/128.


add an assert in from_hp?

Corrected. There is no block size limitation for XPU backend.

jerryzh168 · 2025-08-29T03:23:16Z

torchao/quantization/quantize_/workflows/int4/int4_plain_int32_tensor.py

+        scale: (K/group_size, N), dtype is the same as the original Tensor dtype
+        zero_point: (K/group_size, N)


also dtypes for these for more clarity

Add the dtype information for zero_point.

meta-cla bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Aug 22, 2025

liangan1 added 3 commits August 22, 2025 09:55

Add Int4XPUTensorIntZP

ec3e065

Add int4_xpu_tensor

1dc5b2c

Update int4_xpu_tensor.py

e63b100

liangan1 requested review from jerryzh168 and andrewor14 August 25, 2025 07:55

liangan1 added the topic: new feature Use this tag if this PR adds a new feature label Aug 25, 2025

liangan1 changed the title ~~[WIP]Add Int4XPUTensorIntZP~~ Add Int4XPUTensorIntZP Aug 25, 2025

liangan1 added 4 commits August 25, 2025 15:55

Fix typo

5ef1ca2

Fix code format issue

a28dd89

fix bug

8a0f124

Fix code format

a0ff36f

jerryzh168 reviewed Aug 26, 2025

View reviewed changes

torchao/quantization/quantize_/workflows/int4/int4_xpu_tensor.py Outdated Show resolved Hide resolved

liangan1 added 5 commits August 26, 2025 09:46

Merge branch 'main' into liangan1/int4_xpu_int_zp

5e9c476

Update int4_xpu_tensor.py

2c4c2ce

change the pack format to plain

e48ea0b

fix typo

c4e5b9d

Update quant_api.py

7063e56

jerryzh168 reviewed Aug 28, 2025

View reviewed changes

torchao/quantization/quantize_/workflows/int4/int4_xpu_tensor.py Outdated Show resolved Hide resolved

liangan1 added 4 commits August 28, 2025 07:31

merge main branch

5b87d8b

Merge branch 'main' into liangan1/int4_xpu_int_zp

6076877

Update __init__.py

8d2acd2

Update __init__.py

43acd66

jerryzh168 reviewed Aug 29, 2025

View reviewed changes

liangan1 requested a review from jerryzh168 August 29, 2025 02:05

liangan1 changed the title ~~Add Int4XPUTensorIntZP~~ Add Int4PlainInt32 Tensor Aug 29, 2025

jerryzh168 reviewed Aug 29, 2025

View reviewed changes

liangan1 requested a review from jerryzh168 August 29, 2025 06:09

liangan1 added 4 commits August 29, 2025 09:45

change Int4XPUTensorIntZP to Int4PlainInt32

a047c00

Update __init__.py

3f70b2b

Refine code

402dd72

Refine code

282f1a8

		"int4_xpu_int_zp is referring to the format used by int4 weight-only quantization on XPU with int zero point, which is a groupwise quantization format."
		INT4_XPU_INT_ZP = "int4_xpu_int_zp"

		scale: (K/group_size, N), dtype is the same as the original Tensor dtype
		zero_point: (K/group_size, N)

Add Int4PlainInt32 Tensor #2845

Are you sure you want to change the base?

Add Int4PlainInt32 Tensor #2845

Uh oh!

Conversation

liangan1 commented Aug 22, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-bot bot commented Aug 22, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful Links

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

❌ 4 New Failures

Uh oh!

Choose a reason for hiding this comment

Uh oh!

liangan1 Aug 26, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

jerryzh168 Aug 29, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

liangan1 commented Aug 22, 2025 •

edited

Loading

pytorch-bot bot commented Aug 22, 2025 •

edited

Loading

liangan1 Aug 26, 2025 •

edited

Loading

jerryzh168 Aug 29, 2025 •

edited

Loading