Add dont constant fold flag by LevelDownRefine · Pull Request #154945 · pytorch/pytorch

LevelDownRefine · 2025-06-03T02:17:22Z

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 #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.
So, we think an easy and short-term solution is to modify the ops in PyTorch core via #153601.
However, if we want to resolve the issue with Torchao, we need to
Add a method in the constant folder in Inductor to allow registration of impure ops

Based on Jansel‘s reply, add dont constant fold flag on this patch

cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @ipiszy @chenyang78 @kadeng @muchulee8 @amjames @chauhang @aakhundov

pytorch-bot · 2025-06-03T02:17:26Z

🔗 Helpful Links

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

📄 Preview Python docs built from this PR
📄 Preview C++ docs built from this PR
❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

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 6437a0d with merge base be2ad70 ():

UNSTABLE - The following job is marked as unstable, possibly due to flakiness on trunk:

pull / linux-jammy-py3-clang12-executorch / test (executorch, 1, 1, linux.2xlarge, unstable) (gh)
exir/backend/test/test_to_backend_multi_method.py::TestToBackendMultiMethod::test_multi_method_end_to_end

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

Xia-Weiwen · 2025-06-03T02:24:26Z

Thanks for the PR. I would suggest elaborating more about the background, like we have discussed elsewhere.

LevelDownRefine · 2025-06-03T02:35:10Z

I want to add the flag dont_constant_fold on torch/_inductor/config.py, but get following warming.

W0529 15:14:03.946369 3148528 torch/_inductor/codecache.py:632] [0/0] Failed to pickle cache key
W0529 15:14:03.946369 3148528 torch/_inductor/codecache.py:632] [0/0] Traceback (most recent call last):
W0529 15:14:03.946369 3148528 torch/_inductor/codecache.py:632] [0/0]   File "/home/wengshiy/pytorch/torch/_inductor/codecache.py", line 628, in dumps
W0529 15:14:03.946369 3148528 torch/_inductor/codecache.py:632] [0/0]     self.dump(obj)         
W0529 15:14:03.946369 3148528 torch/_inductor/codecache.py:632] [0/0] TypeError: cannot pickle 'PyCapsule' object

So I add this flag on torch/_inductor/constant_folding.py firstly

LevelDownRefine · 2025-06-03T02:43:44Z

Thanks for the PR. I would suggest elaborating more about the background, like we have discussed elsewhere.

Done

…t_fold

leslie-fang-intel

Lint Failure?

LevelDownRefine · 2025-06-04T09:08:57Z

Lint Failure?

Fixed

LevelDownRefine · 2025-06-04T09:10:40Z

Hi @eellison @jansel , could you help review this patch?

…t_fold

Co-authored-by: Jason Ansel <jansel@jansel.net>

LevelDownRefine · 2025-06-05T02:52:43Z

@pytorchbot merge

pytorch-bot · 2025-06-05T02:52:47Z

This PR has pending changes requested. Please address the comments and update the PR before merging.

LevelDownRefine · 2025-06-05T08:35:19Z

@pytorchbot merge

pytorchmergebot · 2025-06-05T08:37:50Z

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here

…t_fold

LevelDownRefine · 2025-06-09T03:07:17Z

skip halide on my test to avoid CI issue.

In my test(test/inductor/test_torchinductor.py), I check the function name to see if mul is constant_fold.
But on test/inductor/test_halide.py, function name will always be halide_kernel_0.

jansel · 2025-06-09T03:13:04Z

@pytorchbot merge

pytorchmergebot · 2025-06-09T03:15:07Z

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here

pytorchmergebot · 2025-06-09T03:41:47Z

Merge failed

Reason: 1 jobs have failed, first few of them are: inductor / unit-test / linux-jammy-cpu-py3.12-gcc11-inductor-triton-cpu / test (inductor-triton-cpu, 1, 1, linux.12xlarge)

Details for Dev Infra team

Raised by workflow job

LevelDownRefine · 2025-06-09T05:39:02Z

@pytorchbot merge

pytorch-bot · 2025-06-09T05:39:06Z

Pull workflow has not been scheduled for the PR yet. It could be because author doesn't have permissions to run those or skip-checks keywords were added to PR/commits, aborting merge. Please get/give approval for the workflows and/or remove skip ci decorators before next merge attempt. If you think this is a mistake, please contact PyTorch Dev Infra.

…t_fold

jansel · 2025-06-10T14:44:24Z

@pytorchbot merge

pytorchmergebot · 2025-06-10T14:46:17Z

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here

pytorch-bot Bot added the module: inductor label Jun 3, 2025

LevelDownRefine changed the title ~~Wengshiy/dont constant fold~~ Add dont constant fold flag Jun 3, 2025

Xia-Weiwen added the release notes: inductor label Jun 3, 2025

LevelDownRefine marked this pull request as draft June 3, 2025 02:27

pytorchbot added the open source label Jun 3, 2025

Xia-Weiwen requested a review from leslie-fang-intel June 3, 2025 02:50

LevelDownRefine added 2 commits June 3, 2025 09:34

add dont_constant_fold flag

2beccf5

Merge remote-tracking branch 'origin/main' into wengshiy/dont_constan…

f537184

…t_fold

leslie-fang-intel requested changes Jun 4, 2025

View reviewed changes

LevelDownRefine marked this pull request as ready for review June 4, 2025 09:08

LevelDownRefine added 2 commits June 4, 2025 13:33

fix lint

53cc0e4

Merge remote-tracking branch 'origin/main' into wengshiy/dont_constan…

d1b9a55

…t_fold

jansel reviewed Jun 4, 2025

View reviewed changes

Comment thread torch/_inductor/constant_folding.py Outdated

jansel previously approved these changes Jun 4, 2025

View reviewed changes

Update torch/_inductor/constant_folding.py

6c249ed

Co-authored-by: Jason Ansel <jansel@jansel.net>

LevelDownRefine requested a review from leslie-fang-intel June 5, 2025 03:02

leslie-fang-intel previously approved these changes Jun 5, 2025

View reviewed changes

pytorch-bot Bot added the ciflow/trunk Trigger trunk jobs on your pull request label Jun 5, 2025

pytorchmergebot added the merging label Jun 5, 2025

Merge remote-tracking branch 'origin/main' into wengshiy/dont_constan…

505017e

…t_fold

LevelDownRefine requested a review from leslie-fang-intel June 9, 2025 02:59

jansel approved these changes Jun 9, 2025

View reviewed changes

pytorch-bot Bot added the ciflow/trunk Trigger trunk jobs on your pull request label Jun 9, 2025

jansel added the ciflow/inductor label Jun 9, 2025

pytorchmergebot added the merging label Jun 9, 2025

pytorchmergebot removed the merging label Jun 9, 2025

pytorch-bot Bot removed ciflow/trunk Trigger trunk jobs on your pull request ciflow/inductor labels Jun 9, 2025

LevelDownRefine added 2 commits June 9, 2025 10:38

skip halide

ea29dec

Merge remote-tracking branch 'origin/main' into wengshiy/dont_constan…

61932be

…t_fold

jansel added ciflow/trunk Trigger trunk jobs on your pull request ciflow/inductor labels Jun 9, 2025

skip triton

6437a0d

pytorchmergebot added the merging label Jun 10, 2025

pytorchmergebot closed this in b44306d Jun 10, 2025

pytorchmergebot removed the merging label Jun 10, 2025

Conversation

LevelDownRefine commented Jun 3, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-bot Bot commented Jun 3, 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/154945

✅ You can merge normally! (1 Unrelated Failure)

Uh oh!

Xia-Weiwen commented Jun 3, 2025

Uh oh!

LevelDownRefine commented Jun 3, 2025

Uh oh!

LevelDownRefine commented Jun 3, 2025

Uh oh!

leslie-fang-intel left a comment

Choose a reason for hiding this comment

Uh oh!

LevelDownRefine commented Jun 4, 2025

Uh oh!

LevelDownRefine commented Jun 4, 2025

Uh oh!

Uh oh!

LevelDownRefine commented Jun 5, 2025

Uh oh!

pytorch-bot Bot commented Jun 5, 2025

Uh oh!

LevelDownRefine commented Jun 5, 2025

Uh oh!

pytorchmergebot commented Jun 5, 2025

Merge started

Uh oh!

LevelDownRefine commented Jun 9, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

jansel commented Jun 9, 2025

Uh oh!

pytorchmergebot commented Jun 9, 2025

Merge started

Uh oh!

pytorchmergebot commented Jun 9, 2025

Merge failed

Uh oh!

LevelDownRefine commented Jun 9, 2025

Uh oh!

pytorch-bot Bot commented Jun 9, 2025

Uh oh!

jansel commented Jun 10, 2025

Uh oh!

pytorchmergebot commented Jun 10, 2025

Merge started

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

8 participants

LevelDownRefine commented Jun 3, 2025 •

edited

Loading

pytorch-bot Bot commented Jun 3, 2025 •

edited

Loading

LevelDownRefine commented Jun 9, 2025 •

edited

Loading