Enable _lazy_clone between CPU and MPS
#148408
Open
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/148408
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 3b1a2c7 with merge base 56e1c23 ( UNSTABLE - The following job is marked as unstable, possibly due to flakiness on trunk:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
Attention! native_functions.yaml was changedIf you are adding a new function or defaulted argument to native_functions.yaml, you cannot use it from pre-existing Python frontend code until our FC window passes (two weeks). Split your PR into two PRs, one which adds the new C++ functionality, and one that makes use of it from Python, and land them two weeks apart. See https://github.com/pytorch/pytorch/wiki/PyTorch's-Python-Frontend-Backward-and-Forward-Compatibility-Policy#forwards-compatibility-fc for more info. Caused by: |
kurtamohler
commented
Mar 8, 2025
kurtamohler
added a commit
to kurtamohler/pytorch
that referenced
this pull request
Mar 11, 2025
ghstack-source-id: 7c251ee Pull Request resolved: pytorch#148408
kurtamohler
commented
Mar 11, 2025
kurtamohler
added a commit
to kurtamohler/pytorch
that referenced
this pull request
Mar 12, 2025
ghstack-source-id: 22e24cc Pull Request resolved: pytorch#148408
kurtamohler
added a commit
to kurtamohler/pytorch
that referenced
this pull request
Mar 12, 2025
ghstack-source-id: 7cb9a35 Pull Request resolved: pytorch#148408
kurtamohler
added a commit
to kurtamohler/pytorch
that referenced
this pull request
Mar 12, 2025
ghstack-source-id: ba80390 Pull Request resolved: pytorch#148408
kurtamohler
added a commit
to kurtamohler/pytorch
that referenced
this pull request
Apr 7, 2025
ghstack-source-id: 97b5249 Pull Request resolved: pytorch#148408
kurtamohler
added a commit
to kurtamohler/pytorch
that referenced
this pull request
Apr 7, 2025
ghstack-source-id: fdcc972 Pull Request resolved: pytorch#148408
albanD
reviewed
Apr 7, 2025
| virtual const void* get_cpu_ptr_from_device_ptr(const void* device_ptr) const; | ||
| virtual void* get_device_ptr_from_cpu_ptr(void* cpu_ptr) const; | ||
| virtual const void* get_device_ptr_from_cpu_ptr(const void* cpu_ptr) const; | ||
| virtual bool has_unified_memory() const; |
There was a problem hiding this comment.
Why do you need to add this new concept here?
I would expect that, in the context of MPS, cpu Tensor is pure cpu, pinned cpu and mps Tensors are unified.
There was a problem hiding this comment.
When CPU and Metal operators access the same location in the shared memory space, they use different addresses. (I wonder if there's a way to make them use the same address space?) So when we lazy clone an MPS tensor to pinned-CPU, we need a way to translate the MPS address into CPU address space and set the output DataPtr to that. Likewise, we need to translate in the other direction for pinned-CPU to MPS lazy clone. These functions provide an API for that
kurtamohler
added a commit
to kurtamohler/pytorch
that referenced
this pull request
Apr 7, 2025
ghstack-source-id: fdcc972 Pull Request resolved: pytorch#148408
kurtamohler
added a commit
to kurtamohler/pytorch
that referenced
this pull request
Apr 7, 2025
ghstack-source-id: f6ae96d Pull Request resolved: pytorch#148408
kurtamohler
added a commit
to kurtamohler/pytorch
that referenced
this pull request
Apr 8, 2025
ghstack-source-id: a863bdd Pull Request resolved: pytorch#148408
Divigroup-RAP
pushed a commit
to Divigroup-RAP/PYTORCH
that referenced
this pull request
Apr 22, 2025
ghstack-source-id: 344f648 Pull Request resolved: pytorch/pytorch#148408
kurtamohler
added a commit
to kurtamohler/pytorch
that referenced
this pull request
Apr 25, 2025
ghstack-source-id: ead3b2f Pull Request resolved: pytorch#148408
kurtamohler
added a commit
to kurtamohler/pytorch
that referenced
this pull request
Apr 25, 2025
ghstack-source-id: ead3b2f Pull Request resolved: pytorch#148408
kurtamohler
added a commit
to kurtamohler/pytorch
that referenced
this pull request
Apr 28, 2025
ghstack-source-id: 3f9e14a Pull Request resolved: pytorch#148408
kurtamohler
added a commit
to kurtamohler/pytorch
that referenced
this pull request
Apr 28, 2025
ghstack-source-id: eb3a36e Pull Request resolved: pytorch#148408
kurtamohler
added a commit
to kurtamohler/pytorch
that referenced
this pull request
Apr 29, 2025
ghstack-source-id: eb3a36e Pull Request resolved: pytorch#148408
kurtamohler
added a commit
to kurtamohler/pytorch
that referenced
this pull request
Apr 29, 2025
ghstack-source-id: eb3a36e Pull Request resolved: pytorch#148408
kurtamohler
added a commit
to kurtamohler/pytorch
that referenced
this pull request
Apr 30, 2025
ghstack-source-id: eb3a36e Pull Request resolved: pytorch#148408
kurtamohler
added a commit
to kurtamohler/pytorch
that referenced
this pull request
Apr 30, 2025
ghstack-source-id: eb3a36e Pull Request resolved: pytorch#148408
kurtamohler
added a commit
to kurtamohler/pytorch
that referenced
this pull request
May 1, 2025
ghstack-source-id: 0179df2 Pull Request resolved: pytorch#148408
kurtamohler
added a commit
to kurtamohler/pytorch
that referenced
this pull request
May 13, 2025
ghstack-source-id: 0179df2 Pull Request resolved: pytorch#148408
kurtamohler
added a commit
to kurtamohler/pytorch
that referenced
this pull request
May 16, 2025
ghstack-source-id: 0179df2 Pull Request resolved: pytorch#148408
kurtamohler
added a commit
to kurtamohler/pytorch
that referenced
this pull request
May 16, 2025
ghstack-source-id: 5e1efe9 Pull Request resolved: pytorch#148408
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Adds
devicearg to_lazy_cloneto enable lazy cloning data from one device to another. At the moment, only the following cases are supported:This PR also adds support for pinned CPU tensors on MPS builds, which was not working properly before.
Stack from ghstack (oldest at bottom):
Tensor.tobetween CPU and MPS #150569_lazy_clonebetween CPU and MPS #148408cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @ipiszy @chenyang78 @kadeng @muchulee8 @amjames @chauhang @aakhundov