8000 Add torch.accelerator.device_index as accelerator's device switch context by guangyey · Pull Request #148864 · pytorch/pytorch · GitHub
[go: up one dir, main page]
Skip to content

Add torch.accelerator.device_index as accelerator's device switch context #148864

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Closed
wants to merge 22 commits into from
Closed

Add torch.accelerator.device_index as accelerator's device switch context #148864

wants to merge 22 commits into from

Conversation

guangyey
Copy link
Collaborator
@guangyey guangyey commented Mar 10, 2025 •
edited
Loading

Stack from ghstack (oldest at bottom):

Motivation

We propose adding support for the Python with statement on torch.accelerator.device_index to enable device switching functionality. This enhancement would simplify writing device-agnostic code and provide benefits across all accelerators. Its device-specific counterparts include torch.cuda.device and torch.cuda._DeviceGuard.

Design Philosophy
It accepts either an Int or None as input. When None is passed, no device switch is performed. Supporting None is important for compatibility, as it's possible to encounter None values from torch.device.index.

Therefore, with this PR, we can do like this

src = 0
dst = 1
# Set src to current device
torch.accelerator.set_device_index(src)
with torch.accelerator.device_index(dst):
    # Inside with statement, we set dst to current device
    assert torch.accelerator.get_device_index() == dst
# Here the current device should be src
assert torch.accelerator.get_device_index() == src

cc @albanD @EikanWang

@pytorch-bot PyTorch Bot
< 8000 /path> Copy link
pytorch-bot bot commented Mar 10, 2025 •
edited
Loading

🔗 Helpful Links

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

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

❌ 3 New Failures, 1 Unrelated Failure

As of commit b43545c with merge base fc6e37c (image):

NEW FAILURES - The following jobs have failed:

FLAKY - The following job failed but was likely due to flakiness present on trunk:

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

guangyey added a commit that referenced this pull request Mar 10, 2025
@guangyey
Support torch.device as accelerator's device switch context
fe0479d 
ghstack-source-id: 752b700
Pull Request resolved: #148864
@guangyey guangyey added ciflow/trunk Trigger trunk jobs on your pull request release notes: python_frontend python frontend release notes category ciflow/xpu Run XPU CI tasks labels Mar 10, 2025
guangyey added a commit that referenced this pull request Mar 10, 2025
@guangyey
Support torch.device as accelerator's device switch context
9919e36 
ghstack-source-id: 9994801
Pull Request resolved: #148864
@guangyey guangyey added the ciflow/rocm Trigger "default" config CI on ROCm label Mar 10, 2025
guangyey added a commit that referenced this pull request Mar 10, 2025
@guangyey
Support torch.device as accelerator's device switch context
21536ee 
ghstack-source-id: 01a7f2c
Pull Request resolved: #148864
@guangyey guangyey added the module: accelerator Issues related to the shared accelerator API label Mar 10, 2025
@guangyey guangyey requested a review from albanD March 10, 2025 08:41
@guangyey guangyey mentioned this pull request Mar 10, 2025
Closed
guangyey added 4 commits March 10, 2025 11:18
@guangyey
Update
959d056 
[ghstack-poisoned]
@guangyey
Update
1a410d5 
[ghstack-poisoned]
@guangyey
Update
1b49b97 
[ghstack-poisoned]
@guangyey
Update
1e42e7c 
[ghstack-poisoned]
@guangyey
8000 Copy link
Collaborator Author

@pytorchbot rebase

@pytorchmergebot
Copy link
Collaborator

@pytorchbot started a rebase job onto refs/remotes/origin/viable/strict. Check the current status here

@guangyey
Update
2a36109 
[ghstack-poisoned]
@pytorchmergebot
Copy link
Collaborator

Successfully rebased gh/guangyey/126/orig onto refs/remotes/origin/viable/strict, please pull locally before adding more changes (for example, via ghstack checkout https://github.com/pytorch/pytorch/pull/148864)

pytorchmergebot pushed a commit that referenced this pull request Mar 12, 2025
@guangyey
Support torch.device as accelerator's device switch context
57413d3 
ghstack-source-id: b294e55
Pull Request resolved: #148864
@guangyey
Copy link
Collaborator Author

@pytorchbot rebase

@pytorchmergebot
Copy link
Collaborator

@pytorchbot started a rebase job onto refs/remotes/origin/viable/strict. Check the current status here

@guangyey
Update
2e2986c 
[ghstack-poisoned]
@pytorchmergebot
Copy link
Collaborator

Successfully rebased gh/guangyey/126/orig onto refs/remotes/origin/viable/strict, please pull locally before adding more changes (for example, via ghstack checkout https://github.com/pytorch/pytorch/pull/148864)

pytorchmergebot pushed a commit that referenced this pull request Mar 14, 2025
@guangyey
Support torch.device as accelerator's device switch context
056285e 
ghstack-source-id: d289cde
Pull Request resolved: #148864
@guangyey
Copy link
Collaborator Author

@albanD May I know if this PR is reasonable for you.

@guangyey
Update
5087a56 
[ghstack-poisoned]
Divigroup-RAP pushed a commit to Divigroup-RAP/PYTORCH that referenced this pull request Apr 22, 2025
@guangyey
Support torch.device as accelerator's device switch context
59185f9 
ghstack-source-id: d07fc95
Pull Request resolved: pytorch/pytorch#148864
@pytorch-bot pytorch-bot bot added the topic: not user facing topic category label Apr 23, 2025
@guangyey guangyey changed the title Support torch.device as accelerator's device switch context Add torch.accelerator.device_index as accelerator's device switch context Apr 23, 2025
@guangyey guangyey changed the title Add torch.accelerator.device_index as accelerator's device switch context [WIP] Add torch.accelerator.device_index as accelerator's device switch context Apr 23, 2025
guangyey added 3 commits April 23, 2025 12:59
@guangyey
Update
1d53ded 
[ghstack-poisoned]
@guangyey
Update
d27f975 
[ghstack-poisoned]
@guangyey
Update
42487d5 
[ghstack-poisoned]
albanD
albanD approved these changes Apr 23, 2025
View reviewed changes
Copy link
Collaborator
@albanD albanD left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks!

guangyey added 6 commits April 23, 2025 16:20
@guangyey
Update
effb3fc 
[ghstack-poisoned]
@guangyey
Update
16d6080 
[ghstack-poisoned]
@guangyey
Update 8000
dad33cc 
[ghstack-poisoned]
@guangyey
Update
f62d254 
[ghstack-poisoned]
@guangyey
Update
c0829f7 
[ghstack-poisoned]
@guangyey
Update
fe60ddc 
[ghstack-poisoned]
@guangyey guangyey changed the title [WIP] Add torch.accelerator.device_index as accelerator's device switch context Add torch.accelerator.device_index as accelerator's device switch context Apr 24, 2025
guangyey added 3 commits April 24, 2025 09:21
@guangyey
Update
c1ee27a 
[ghstack-poisoned]
@guangyey
Update
921ce45 
[ghstack-poisoned]
@guangyey
Update
6ae0a78 
[ghstack-poisoned]
@pytorchmergebot
Copy link
Collaborator< 6D40 /span>

Starting merge as part of PR stack under #148880

@guangyey
Update
b43545c 
[ghstack-poisoned]
@pytorchmergebot
Copy link
Collaborator

Rebased gh/guangyey/127/orig onto refs/remotes/origin/viable/strict because #148880 was rebased, please pull locally before adding more changes (for example, via ghstack checkout https://github.com/pytorch/pytorch/pull/148864)

@pytorchmergebot
Copy link
Collaborator

Starting merge as part of PR stack under #148880

1 similar comment
@pytorchmergebot
Copy link
Collaborator

Starting merge as part of PR stack under #148880

pytorchmergebot pushed a commit that referenced this pull request Apr 25, 2025
@guangyey @pytorchmergebot
Refactor to use torch.accelerator.device_index instead of torch.cuda.…
ad81eeb 
…device for generic device context manager (#148880)

Pull Request resolved: #148880
Approved by: https://github.com/EikanWang, https://github.com/albanD
ghstack dependencies: #148864
@guangyey guangyey mentioned this pull request May 5, 2025
Closed
pytorchmergebot pushed a commit that referenced this pull request May 6, 2025
@guangyey @pytorchmergebot
Fix typo on test_multi_device_context_manager for XPU (#152812)
35c727e 
# Motivation
Align #152474, fix the typo on UT for XPU introduced by #148864
Pull Request resolved: #152812
Approved by: https://github.com/EikanWang, https://github.com/Skylion007
@guangyey guangyey mentioned this pull request May 7, 2025
Closed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@albanD albanD albanD approved these changes

@EikanWang EikanWang Awaiting requested review from EikanWang EikanWang is a code owner

@gujinghui gujinghui Awaiting requested review from gujinghui gujinghui is a code owner

Assignees
No one assigned
Labels
ciflow/rocm Trigger "default" config CI on ROCm ciflow/trunk Trigger trunk jobs on your pull request ciflow/xpu Run XPU CI tasks Merged module: accelerator Issues related to the shared accelerator API open source release notes: python_frontend python frontend release notes category topic: not user facing topic category
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
4 participants
@guangyey @pytorchmergebot @albanD @pytorchbot
0