[pt2e] Avoid getting model device once per node #159901

andrewor14 · 2025-08-05T20:56:15Z

Summary: Previously, we call assert_and_get_unqiue_device once per node in both prepare and convert. This is expensive and unnecessary since the model device is the same across all nodes, so we should just call this once in the beginning and reuse the same model device across all the nodes.

Test Plan:
python test/test_quantization.py -k TestQuantizePT2E

cc @ezyang @SherlockNoMad @EikanWang @jgong5 @wenzhe-nrv

pytorch-bot · 2025-08-05T20:56:19Z

🔗 Helpful Links

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

📄 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

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

✅ No Failures

As of commit c4be20c with merge base 791eff9 ():
💚 Looks good so far! There are no failures yet. 💚

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

facebook-github-bot · 2025-08-05T20:56:41Z

@andrewor14 has imported this pull request. If you are a Meta employee, you can view this in D79674759.

andrewor14 · 2025-08-05T20:57:50Z

torch/ao/quantization/utils.py

-            "Both 'meta' and 'cpu' are present in the list of devices. Module can have one device. We Select 'cpu'."
-        )
-        devices = {torch.device("cpu")}
-    ""


@jerryzh168 is it safe to remove this?

let me check

facebook-github-bot · 2025-08-05T21:06:21Z

@andrewor14 has imported this pull request. If you are a Meta employee, you can view this in D79674759.

**Summary:** Previously, we call `assert_and_get_unqiue_device` once per node in convert. This is expensive and unnecessary since the model device is the same across all nodes, so we should just call this once in the beginning and reuse the same model device across all the nodes. torchao version of pytorch/pytorch#159901 **Test Plan:** ``` python test/quantization/pt2e/test_quantize_pt2e.py ```

facebook-github-bot · 2025-08-05T21:25:07Z

@andrewor14 has imported this pull request. If you are a Meta employee, you can view this in D79674759.

**Summary:** Previously, we call `assert_and_get_unqiue_device` once per node in both prepare and convert. This is expensive and unnecessary since the model device is the same across all nodes, so we should just call this once in the beginning and reuse the same model device across all the nodes. torchao version of pytorch/pytorch#159901 Note: The prepare path is not completely done yet, since we are blocked on the pytorch PR on being merged. It's different from convert since it still calls utility functions from `torch.ao.quantization.fx`. **Test Plan:** ``` python test/quantization/pt2e/test_quantize_pt2e.py ```

facebook-github-bot · 2025-08-05T21:46:56Z

@andrewor14 has imported this pull request. If you are a Meta employee, you can view this in D79674759.

facebook-github-bot · 2025-08-06T00:03:24Z

@andrewor14 has imported this pull request. If you are a Meta employee, you can view this in D79674759.

iremyux · 2025-08-06T08:49:00Z

Adding ciflow/win-arm64 label to trigger Windows Arm64 CI and its test purposes - nothing about this PR specifically. (It should not effect the acceptance of the PR even if it fails.)

**Summary:** Previously, we call `assert_and_get_unqiue_device` once per node in both prepare and convert. This is expensive and unnecessary since the model device is the same across all nodes, so we should just call this once in the beginning and reuse the same model device across all the nodes. torchao version of pytorch/pytorch#159901 Note: The prepare path is not completely done yet, since we are blocked on the pytorch PR on being merged. It's different from convert since it still calls utility functions from `torch.ao.quantization.fx`. **Test Plan:** ``` python test/quantization/pt2e/test_quantize_pt2e.py ```

facebook-github-bot · 2025-08-06T16:08:55Z

@andrewor14 has imported this pull request. If you are a Meta employee, you can view this in D79674759.

facebook-github-bot · 2025-09-02T21:37:56Z

@andrewor14 has imported this pull request. If you are a Meta employee, you can view this in D79674759.

**Summary:** Previously, we call `assert_and_get_unqiue_device` once per node in both prepare and convert. This is expensive and unnecessary since the model device is the same across all nodes, so we should just call this once in the beginning and reuse the same model device across all the nodes. torchao version of pytorch/pytorch#159901 Note: The prepare path is not completely done yet, since we are blocked on the pytorch PR on being merged. It's different from convert since it still calls utility functions from `torch.ao.quantization.fx`. **Test Plan:** ``` python test/quantization/pt2e/test_quantize_pt2e.py ```

jerryzh168 · 2025-09-03T00:25:25Z

torch/ao/quantization/observer.py

+                    model.graph,
+                    "_scale",
+                    scale,
+                    scale.device,


value could be a non-Tensor I think? according to create_getattr_from_value L274, so should this be the device from model?

Ok, seems like if we just get it from the model there will be no benefits, so how about here we get it from the scale if it's a tensor and get it from the model otherwise?

**Summary:** Previously, we call `assert_and_get_unqiue_device` once per node in both prepare and convert. This is expensive and unnecessary since the model device is the same across all nodes, so we should just call this once in the beginning and reuse the same model device across all the nodes. torchao version of pytorch/pytorch#159901 Note: The prepare path is not completely done yet, since we are blocked on the pytorch PR on being merged. It's different from convert since it still calls utility functions from `torch.ao.quantization.fx`. **Test Plan:** ``` python test/quantization/pt2e/test_quantize_pt2e.py ```

**Summary:** Previously, we call `assert_and_get_unqiue_device` once per node in both prepare and convert. This is expensive and unnecessary since the model device is the same across all nodes, so we should just call this once in the beginning and reuse the same model device across all the nodes. **Test Plan:** python test/test_quantization.py -k TestQuantizePT2E ghstack-source-id: 3e72b14 Pull Request resolved: #162012

facebook-github-bot · 2025-09-03T14:46:46Z

@andrewor14 has imported this pull request. If you are a Meta employee, you can view this in D79674759.

andrewor14 · 2025-09-03T19:21:03Z

@pytorchbot merge

**Summary:** Previously, we call `assert_and_get_unqiue_device` once per node in both prepare and convert. This is expensive and unnecessary since the model device is the same across all nodes, so we should just call this once in the beginning and reuse the same model device across all the nodes. torchao version of pytorch/pytorch#159901 Note: The prepare path is not completely done yet, since we are blocked on the pytorch PR on being merged. It's different from convert since it still calls utility functions from `torch.ao.quantization.fx`. **Test Plan:** ``` python test/quantization/pt2e/test_quantize_pt2e.py ```

pytorchmergebot · 2025-09-03T19:23:03Z

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

**Summary:** Previously, we call `assert_and_get_unqiue_device` once per node in both prepare and convert. This is expensive and unnecessary since the model device is the same across all nodes, so we should just call this once in the beginning and reuse the same model device across all the nodes. **Test Plan:** python test/test_quantization.py -k TestQuantizePT2E Pull Request resolved: pytorch#159901 Approved by: https://github.com/jerryzh168

andrewor14 requested a review from jerryzh168 August 5, 2025 20:56

pytorch-bot bot added release notes: quantization release notes category release notes: AO frontend labels Aug 5, 2025

andrewor14 requested a review from HDCharles August 5, 2025 20:56

facebook-github-bot added the fx label Aug 5, 2025

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

andrewor14 commented Aug 5, 2025

View reviewed changes

andrewor14 added the topic: improvements topic category label Aug 5, 2025

andrewor14 force-pushed the pt2e-cache-model-device branch from 326006f to 82e8fd3 Compare August 5, 2025 21:05

andrewor14 force-pushed the pt2e-cache-model-device branch from 82e8fd3 to 4eeaee7 Compare August 5, 2025 21:18

andrewor14 mentioned this pull request Aug 5, 2025

[pt2e] Avoid getting model device once per node pytorch/ao#2695

Merged

andrewor14 force-pushed the pt2e-cache-model-device branch 2 times, most recently from 1a3b405 to f1b7d46 Compare August 5, 2025 21:39

andrewor14 force-pushed the pt2e-cache-model-device branch from f1b7d46 to 16077ad Compare August 6, 2025 00:03

andrewor14 force-pushed the pt2e-cache-model-device branch from 16077ad to 27e08a5 Compare August 6, 2025 16:04

andrewor14 force-pushed the pt2e-cache-model-device branch from 27e08a5 to 64aae3a Compare September 2, 2025 21:34

jerryzh168 reviewed Sep 3, 2025

View reviewed changes

andrewor14 force-pushed the pt2e-cache-model-device branch from 64aae3a to c4be20c Compare September 3, 2025 14:45

jerryzh168 approved these changes Sep 3, 2025

View reviewed changes

pytorchmergebot added the merging label Sep 3, 2025

pytorchmergebot added the Merged label Sep 3, 2025

pytorchmergebot closed this in f4c33cd Sep 3, 2025

pytorchmergebot removed the merging label Sep 3, 2025

andrewor14 deleted the pt2e-cache-model-device branch September 3, 2025 19:37

[pt2e] Avoid getting model device once per node #159901

[pt2e] Avoid getting model device once per node #159901

Uh oh!

Conversation

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

Uh oh!

Uh oh!

pytorch-bot bot commented Aug 5, 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/159901

✅ No Failures

Uh oh!

facebook-github-bot commented Aug 5, 2025

Uh oh!

andrewor14 Aug 5, 2025

Choose a reason for hiding this comment

Uh oh!

jerryzh168 Aug 5, 2025

Choose a reason for hiding this comment

Uh oh!

facebook-github-bot commented Aug 5, 2025

Uh oh!

facebook-github-bot commented Aug 5, 2025

Uh oh!

facebook-github-bot commented Aug 5, 2025

Uh oh!

facebook-github-bot commented Aug 6, 2025

Uh oh!

iremyux commented Aug 6, 2025

Uh oh!

facebook-github-bot commented Aug 6, 2025

Uh oh!

facebook-github-bot commented Sep 2, 2025

Uh oh!

jerryzh168 Sep 3, 2025

Choose a reason for hiding this comment

Uh oh!

andrewor14 Sep 3, 2025

Choose a reason for hiding this comment

Uh oh!

facebook-github-bot commented Sep 3, 2025

Uh oh!

andrewor14 commented Sep 3, 2025

Uh oh!

pytorchmergebot commented Sep 3, 2025

Merge started

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

andrewor14 commented Aug 5, 2025 •

edited by pytorch-bot bot

Loading

pytorch-bot bot commented Aug 5, 2025 •

edited

Loading