[quant][pt2] Preserve source_fn_stack after QAT fusion #110899

andrewor14 · 2023-10-09T22:03:33Z

Stack from ghstack (oldest at bottom):

-> [quant][pt2] Preserve source_fn_stack after QAT fusion #110899

Test Plan:
python test/test_quantization.py TestQuantizePT2EQAT.test_qat_preserve_source_fn_stack

Reviewers: jerryzh168, kimishpatel

Subscribers: jerryzh168, kimishpatel, supriyar

Differential Revision: D50101253

Test Plan: python test/test_quantization.py TestQuantizePT2EQAT.test_qat_preserve_source_fn_stack Reviewers: jerryzh168, kimishpatel Subscribers: jerryzh168, kimishpatel, supriyar [ghstack-poisoned]

pytorch-bot · 2023-10-09T22:03:36Z

🔗 Helpful Links

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

📄 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.

✅ No Failures

As of commit 3341aef with merge base 201d02e ():
💚 Looks good so far! There are no failures yet. 💚

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

Test Plan: python test/test_quantization.py TestQuantizePT2EQAT.test_qat_preserve_source_fn_stack Reviewers: jerryzh168, kimishpatel Subscribers: jerryzh168, kimishpatel, supriyar [ghstack-poisoned]

Test Plan: python test/test_quantization.py TestQuantizePT2EQAT.test_qat_preserve_source_fn_stack Reviewers: jerryzh168, kimishpatel Subscribers: jerryzh168, kimishpatel, supriyar ghstack-source-id: d3c2b02 Pull Request resolved: #110899

andrewor14 · 2023-10-09T22:15:03Z

@andrewor14 has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

test/quantization/pt2e/test_quantize_pt2e_qat.py

jerryzh168 · 2023-10-09T22:31:46Z

torch/ao/quantization/pt2e/qat_utils.py

+        (
+            replacement_conv_node,
+            replacement_bn_node,
+            replacement_getitem_node,
+            replacement_relu_node,
+        ) = _get_conv_bn_getitem_nodes(r.replacements)


I guess this is for identifying the nodes, I feel #110743 should help

Ok, let's do that in the future

jerryzh168

maybe we can extend #110743 to work with subgraph rewriter and query node by name

Test Plan: python test/test_quantization.py TestQuantizePT2EQAT.test_qat_preserve_source_fn_stack Reviewers: jerryzh168, kimishpatel Subscribers: jerryzh168, kimishpatel, supriyar Differential Revision: [D50101253](https://our.internmc.facebook.com/intern/diff/D50101253) [ghstack-poisoned]

andrewor14 · 2023-10-10T15:20:07Z

@andrewor14 has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

Test Plan: python test/test_quantization.py TestQuantizePT2EQAT.test_qat_preserve_source_fn_stack Reviewers: jerryzh168, kimishpatel Subscribers: jerryzh168, kimishpatel, supriyar Differential Revision: [D50101253](https://our.internmc.facebook.com/intern/diff/D50101253) [ghstack-poisoned]

Test Plan: python test/test_quantization.py TestQuantizePT2EQAT.test_qat_preserve_source_fn_stack Reviewers: jerryzh168, kimishpatel Subscribers: jerryzh168, kimishpatel, supriyar ghstack-source-id: fa25b14 Pull Request resolved: #110899

andrewor14 · 2023-10-10T17:11:30Z

@andrewor14 has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

Test Plan: python test/test_quantization.py TestQuantizePT2EQAT.test_qat_preserve_source_fn_stack Reviewers: jerryzh168, kimishpatel Subscribers: jerryzh168, kimishpatel, supriyar Differential Revision: [D50101253](https://our.internmc.facebook.com/intern/diff/D50101253) [ghstack-poisoned]

Test Plan: python test/test_quantization.py TestQuantizePT2EQAT.test_qat_preserve_source_fn_stack Reviewers: jerryzh168, kimishpatel Subscribers: jerryzh168, kimishpatel, supriyar ghstack-source-id: 163b993 Pull Request resolved: #110899

andrewor14 · 2023-10-10T19:34:51Z

@andrewor14 has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

andrewor14 · 2023-10-11T00:21:55Z

@pytorchbot merge

pytorchmergebot · 2023-10-11T00:23:53Z

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

kimishpatel · 2023-10-12T13:41:53Z

torch/ao/quantization/pt2e/qat_utils.py

    return (q_node, dq_node)

-def _get_conv_bn_getitem_nodes(nodes: List[Node]) -> Tuple[Node, Node, Node]:
+def _get_conv_bn_getitem_nodes(nodes: List[Node]) -> Tuple[Node, Node, Node, Optional[Node]]:


should also be renamed to _get_conv_bn_getitem_relu_nodes

kimishpatel · 2023-10-12T13:45:45Z

torch/ao/quantization/pt2e/qat_utils.py

+            replacement_relu_node,
+        ) = _get_conv_bn_getitem_nodes(r.replacements)
+
+        for original_node in _filter_nodes_map(r.nodes_map).values():


Andrew, why do you need to copy the metadata?

This is for lowering models trained by QAT. We found that the source_fn was not preserved and all convs in the model ended up being in the same partition conv2d, but they should be in different partitions

[quant][pt2] Preserve source_fn_stack after QAT fusion

0224b4e

Test Plan: python test/test_quantization.py TestQuantizePT2EQAT.test_qat_preserve_source_fn_stack Reviewers: jerryzh168, kimishpatel Subscribers: jerryzh168, kimishpatel, supriyar [ghstack-poisoned]

andrewor14 requested a review from jerryzh168 as a code owner October 9, 2023 22:03

pytorch-bot bot added the release notes: quantization release notes category label Oct 9, 2023

Update on "[quant][pt2] Preserve source_fn_stack after QAT fusion"

1e725e7

Test Plan: python test/test_quantization.py TestQuantizePT2EQAT.test_qat_preserve_source_fn_stack Reviewers: jerryzh168, kimishpatel Subscribers: jerryzh168, kimishpatel, supriyar [ghstack-poisoned]

Update on "[quant][pt2] Preserve source_fn_stack after QAT fusion"

aabb54f

Test Plan: python test/test_quantization.py TestQuantizePT2EQAT.test_qat_preserve_source_fn_stack Reviewers: jerryzh168, kimishpatel Subscribers: jerryzh168, kimishpatel, supriyar [ghstack-poisoned]

Update on "[quant][pt2] Preserve source_fn_stack after QAT fusion"

035335c

Test Plan: python test/test_quantization.py TestQuantizePT2EQAT.test_qat_preserve_source_fn_stack Reviewers: jerryzh168, kimishpatel Subscribers: jerryzh168, kimishpatel, supriyar [ghstack-poisoned]

jerryzh168 reviewed Oct 9, 2023

View reviewed changes

test/quantization/pt2e/test_quantize_pt2e_qat.py Outdated Show resolved Hide resolved

jerryzh168 reviewed Oct 9, 2023

View reviewed changes

test/quantization/pt2e/test_quantize_pt2e_qat.py Outdated Show resolved Hide resolved

jerryzh168 reviewed Oct 9, 2023

View reviewed changes

jerryzh168 approved these changes Oct 9, 2023

View reviewed changes

pytorch-bot bot added the ciflow/trunk Trigger trunk jobs on your pull request label Oct 11, 2023

pytorchmergebot added the merging label Oct 11, 2023

pytorchmergebot added Merged and removed merging labels Oct 11, 2023

pytorchmergebot closed this in 0e551bb Oct 11, 2023

kimishpatel reviewed Oct 12, 2023

View reviewed changes

facebook-github-bot deleted the gh/andrewor14/34/head branch October 14, 2023 14:24

[quant][pt2] Preserve source_fn_stack after QAT fusion #110899

[quant][pt2] Preserve source_fn_stack after QAT fusion #110899

Uh oh!

Conversation

andrewor14 commented Oct 9, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-bot bot commented Oct 9, 2023 • 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/110899

✅ No Failures

Uh oh!

andrewor14 commented Oct 9, 2023

Uh oh!

Uh oh!

Uh oh!

jerryzh168 Oct 9, 2023

Choose a reason for hiding this comment

Uh oh!

andrewor14 Oct 10, 2023

Choose a reason for hiding this comment

Uh oh!

jerryzh168 left a comment

Choose a reason for hiding this comment

Uh oh!

andrewor14 commented Oct 10, 2023

Uh oh!

andrewor14 commented Oct 10, 2023

Uh oh!

andrewor14 commented Oct 10, 2023

Uh oh!

andrewor14 commented Oct 11, 2023

Uh oh!

pytorchmergebot commented Oct 11, 2023

Merge started

Uh oh!

kimishpatel Oct 12, 2023

Choose a reason for hiding this comment

Uh oh!

kimishpatel Oct 12, 2023

Choose a reason for hiding this comment

Uh oh!

andrewor14 Oct 12, 2023

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

andrewor14 commented Oct 9, 2023 •

edited

Loading

pytorch-bot bot commented Oct 9, 2023 •

edited

Loading