skip quant/dequant decomposed #2299

shiyang-weng · 2025-06-04T08:27:12Z

What we want to do now is to enable FP8 quantization in PyTorch. And similar as INT8 quantization, we need to insert quantize and dequantize ops into the graph.
However we met problems with these q/dq ops both in the PyTorch core and Torchao.

PyTorch core:

The quantize_per_tensor op does not support FP8. We want to fix it via pytorch/pytorch#153601. And as you commented, the op is deprecated.
Torchao:

In the fusion pass in Inductor, we want to match the pattern fp8_weight -> torchao.dequantize_affine_float8 -> fp32_op and fuse it as fp8_weight -> weight_pack -> fp8_op. We have done so for INT8 PT2E quantization. However, the pattern matching pass is applied after a constant folding pass in Inductor:
https://github.com/pytorch/pytorch/blob/100ec0b34aeff2b948dae33937857d0c86cf1646/torch/_inductor/fx_passes/freezing_patterns.py#L69C1-L74C1
After constant_fold(gm), the pattern will be folded as fp32_weight -> fp32_op. Then the original pattern cannot be found any more and the FP8 semantics is lost since the pattern is entirely in fp32 now.
For INT8, the int8_weight -> quantized_decomposed.dequantize_per_channel -> fp32_op pattern won't be folded because we mark quantized_decomposed.dequantize_per_channel impure so that it won't be folded: https://github.com/pytorch/pytorch/blob/100ec0b34aeff2b948dae33937857d0c86cf1646/torch/_inductor/constant_folding.py#L139C1-L149C1 . But for the torchao.dequantize_affine_float8, we cannot do this because
It is an op from Torchao, which is unknown to the constant folder
It is decomposed to smaller ops, so we cannot put it in the list as a single op.

Aligned dispatch_key with pt. quantize_affine and dequantize_affine will not be decomposed
Register meta func for q/dq
Register dq for uintx. Previously, dequant would be decomposed into small ops without the need for registration.

pytorch-bot · 2025-06-04T08:27:16Z

🔗 Helpful Links

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

📄 Preview Python docs built from this PR

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

❌ 3 New Failures, 1 Cancelled Job

As of commit 0eedbf3 with merge base 1239842 ():

NEW FAILURES - The following jobs have failed:

PR Label Check / Check PR Labels (gh)
Process completed 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)
test/quantization/pt2e/test_x86inductor_fusion.py::TestPatternMatcher::test_smooth_quant_with_int_mm_has_bias_True_float32_per_channel_quant_True_dynamic_True
Run Regression Tests / test-nightly (CUDA Nightly, linux.g5.12xlarge.nvidia.gpu, --pre torch --index-url https://downloa... / linux-job (gh)
test/sparsity/test_sparse_api.py::TestQuantSemiSparse::test_quant_semi_sparse_compile_False

CANCELLED JOB - The following job was cancelled. Please retry:

Run Float8 Tests / test (H100, linux.aws.h100, --pre torch torchvision torchaudio --index-url https://download.pytor... / linux-job (gh)
##[error]The operation was canceled.

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

facebook-github-bot · 2025-06-04T08:27:18Z

Hi @shiyang-weng!

Thank you for your pull request and welcome to our community.

Action Required

In order to merge any pull request (code, docs, etc.), we require contributors to sign our Contributor License Agreement, and we don't seem to have one on file for you.

Process

In order for us to review and merge your suggested changes, please sign at https://code.facebook.com/cla. If you are contributing on behalf of someone else (eg your employer), the individual CLA may not be sufficient and your employer may need to sign the corporate CLA.

Once the CLA is signed, our tooling will perform checks and validations. Afterwards, the pull request will be tagged with CLA signed. The tagging process may take up to 1 hour after signing. Please give it that time before contacting us about it.

If you have received this in error or have any questions, please contact us at cla@meta.com. Thanks!

Xia-Weiwen · 2025-06-04T08:35:18Z

torchao/quantization/quant_primitives.py

+
+
+_quantize_affine, _quantize_affine_meta = register_custom_op_with_meta(
+    _quantize_affine, _quantize_affine_meta)


Can we separate the registrations for meta and non-meta device and also use them as decorators instead calling it? Then it would be better aligned with the rest of the file.

Can we separate the registrations for meta and non-meta device and also use them as decorators instead calling it? Then it would be better aligned with the rest of the file.

Done

…key in pytorch

leslie-fang-intel · 2025-06-09T23:37:57Z

torchao/utils.py


            lib_namespace = lib.ns
            op = getattr(getattr(torch.ops, lib_namespace), op_name)
-            register_decomposition([op])(fn)
+            if dispatch_key == "CompositeImplicitAutograd":
+                register_decomposition([op])(fn)


Feeling adding a flag in _register_custom_op to indicate whether this op should be decomposed will be clearer, cc @jerryzh168 what's your suggestions?

yeah I feel adding a flag might be easier here

…ay is to register torchao.dequantize_affine.default but fail

shiyang-weng · 2025-06-12T02:03:19Z

torchao/dtypes/uintx/uintx_layout.py

+@implements(torch.ops.torchao.dequantize_affine)
+def _(func, types, args, kwargs):
+    return return_and_correct_aliasing(
+        func.default,  # work around


Better way may be

@implements(torch.ops.torchao.dequantize_affine.default) def _(func, types, args, kwargs): return return_and_correct_aliasing( func, args, kwargs, _dequantize_affine_impl(*args), )

But there will be following issue. NotImplementedError: UintxTensor dispatch: attempting to run unimplemented operator/function: func=<OpOverload(op='aten.to', overload='dtype')>, types=(<class 'torchao.dtypes.uintx.uintx_layout.UintxTensor'>,)

are you using uintx layout for float8? why is this changed?

are you using uintx layout for float8? why is this changed?

dequantize_affine/quantize_affine also used in uintx.
Without this patch, dequant would be decomposed, so it's ok for uintx, but now it is not decomposed. It is necessary to be found on dispatch__torch_function_.
So I use _implements to add it.

But there will be following issue. NotImplementedError: UintxTensor dispatch: attempting to run unimplemented operator/function: func=<OpOverload(op='aten.to', overload='dtype')>, types=(<class 'torchao.dtypes.uintx.uintx_layout.UintxTensor'>,)

for this, you could implement to in the UintxTensor as well I think

shiyang-weng · 2025-06-12T08:24:06Z

Remove this pr still has this accuracy error.
So this error may not be caused by this patch

$ conda list |grep torch
pytorch-triton            3.3.1+gitc8757738          pypi_0    pypi
torch                     2.8.0.dev20250608+cu126          pypi_0    pypi
torchao                   0.12.0+git3aa9361           dev_0    <develop>
torchaudio                2.8.0.dev20250609+cu126          pypi_0    pypi
torchvision               0.23.0.dev20250609+cu126          pypi_0    pypi
$ git log -1
commit 3aa93619466739c9d9845e1db3bfb2ff0f464857 (HEAD, origin/main, origin/HEAD, main)
$ rm -rf /tmp/torchinductor_pt-gpu/ && /data/wengshiy/conda/envs/ao/bin/pytest -s test/integration/test_integration.py::TestSubclass
       if q.is_Number and p.is_Number:
>           assert p >= 0, p
E           torch._inductor.exc.InductorError: AssertionError: -420769046757011/1000000000000000
E           
../conda/envs/ao/lib/python3.10/site-packages/torch/utils/_sympy/functions.py:488: InductorError

shiyang-weng · 2025-06-13T10:37:56Z

Could you help reivew this pr? @jerryzh168
Looks CI errors are unrelated.

And I need to fill out CLA, do you know who I can fill out on the "point of contact"?

…tration failed

jerryzh168 · 2025-06-13T15:39:20Z

test/quantization/test_quant_primitives.py

+                quant_min=torch.finfo(float8_dtype).min,
+                quant_max=torch.finfo(float8_dtype).max,


I remember now we have quantize_affine_float8?

Without this patch, dequantize_affine/quantize_affine will be decomposed both on int8 and fp8.
If you think it's better to handle fp8 operators separately, I will create another pr to register fp8 operators separately

we should use quantize_affine_float8 / dequantize_affine_float8 for float8 I think

jerryzh168 · 2025-06-13T16:08:55Z

for CLA I think you can just follow the process in the CI:

In order for us to review and merge your suggested changes, please sign at https://code.facebook.com/cla. If you are contributing on behalf of someone else (eg your employer), the individual CLA may not be sufficient and your employer may need to sign the corporate CLA.

Once the CLA is signed, our tooling will perform checks and validations. Afterwards, the pull request will be tagged with CLA signed. The tagging process may take up to 1 hour after signing. Please give it that time before contacting us about it.

If you have received this in error or have any questions, please contact us at [cla@meta.com](mailto:cla@meta.com?subject=CLA%20for%20pytorch%2Fao%20%232299). Thanks!

jerryzh168 · 2025-06-17T03:17:23Z

torchao/utils.py

@@ -179,7 +179,7 @@ def find_multiple(n: int, *args: int) -> int:
    return n + k - (n % k)


-def _register_custom_op(lib):
+def _register_custom_op(lib, dispatch_key="CompositeImplicitAutograd"):


are you planning to change dispatch key to a flag of whether to decompose or not

shiyang-weng · 2025-07-21T07:17:30Z

Fixed on #2379

shiyang-weng marked this pull request as draft June 4, 2025 08:27

Xia-Weiwen reviewed Jun 4, 2025

View reviewed changes

shiyang-weng added 4 commits June 4, 2025 16:07

support not constant folding on inductor

cae0123

fix linting

fb2006c

refine code

01a3587

add test

d9d3afb

shiyang-weng changed the title ~~[WIP] skip quant/dequant decomposed~~ skip quant/dequant decomposed Jun 9, 2025

Xia-Weiwen approved these changes Jun 9, 2025

View reviewed changes

shiyang-weng marked this pull request as ready for review June 9, 2025 08:14

shiyang-weng added 9 commits June 9, 2025 09:42

dont_constant_fold is not required after alignment with the dispatch_…

cdabf2d

…key in pytorch

finx linting

ad96296

fix linting

88ea7f2

register meta ops; refine code

70bc8a4

remove test temp to rebase

9b71d19

Merge remote-tracking branch 'origin/main' into wengshiy/quant

88cbde8

re-add test

c85da2c

fix merge issue

14f5658

fix linting

0532521

leslie-fang-intel reviewed Jun 9, 2025

View reviewed changes

shiyang-weng added 4 commits June 11, 2025 13:10

fix register_meta bugs but CI still has issues

49a8efb

fix meta bug

03a9e55

work around to register torchao.dequantize_affine for uintx; better w…

0edde91

…ay is to register torchao.dequantize_affine.default but fail

enable some unitx tests on cpu; other unitx tests must run on cuda

e087fb5

shiyang-weng commented Jun 12, 2025

View reviewed changes

shiyang-weng requested review from Xia-Weiwen and leslie-fang-intel June 12, 2025 03:06

Xia-Weiwen approved these changes Jun 12, 2025

View reviewed changes

leslie-fang-intel approved these changes Jun 13, 2025

View reviewed changes

shiyang-weng added 2 commits June 13, 2025 13:16

enable kwargs for uintx quant

29daf70

Change to use torch.library.impl register meta because previous regis…

a1fbb9c

…tration failed

jerryzh168 reviewed Jun 13, 2025

View reviewed changes

fix stride bugs

0eedbf3

shiyang-weng mentioned this pull request Jun 16, 2025

[float8] Prevent quantize_affine_float8/dequantize_affine_float8 decomposed on inductor #2379

Merged

jerryzh168 reviewed Jun 17, 2025

View reviewed changes

facebook-github-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 Jun 19, 2025

shiyang-weng closed this Jul 21, 2025



		_quantize_affine, _quantize_affine_meta = register_custom_op_with_meta(
		_quantize_affine, _quantize_affine_meta)

		quant_min=torch.finfo(float8_dtype).min,
		quant_max=torch.finfo(float8_dtype).max,

skip quant/dequant decomposed #2299

skip quant/dequant decomposed #2299

Uh oh!

Conversation

shiyang-weng commented Jun 4, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-bot bot commented Jun 4, 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/2299

❌ 3 New Failures, 1 Cancelled Job

Uh oh!

facebook-github-bot commented Jun 4, 2025

Action Required

Process

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!

shiyang-weng Jun 12, 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!

shiyang-weng commented Jun 12, 2025

Uh oh!

shiyang-weng commented Jun 13, 2025

Uh oh!

jerryzh168 Jun 13, 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!

jerryzh168 commented Jun 13, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

shiyang-weng commented Jul 21, 2025

Uh oh!

Uh oh!

shiyang-weng commented Jun 4, 2025 •

edited

Loading

pytorch-bot bot commented Jun 4, 2025 •

edited

Loading

shiyang-weng Jun 12, 2025 •

edited

Loading

jerryzh168 Jun 13, 2025 •

edited

Loading