-
Notifications
You must be signed in to change notification settings - Fork 505
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.
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
Improve device auto-detection #7787
Conversation
RFC: Is there a plan to hard code PJRT_DEVICE in the python script to replace the PJRT_DEVICE environment variable? @miladm |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Thanks for the clean up!
@JackCaoG do you see any problem selecting $PJRT_DEVICE
when we import torch_xla
?
seems fine to me as long as we don't init the runtime yet. |
Fixes Lightning-AI#20419 `torch_xla.runtime.using_pjrt()` is removed in pytorch/xla#7787 This PR replaces references to that function with a check to [`device_type()`](https://github.com/pytorch/xla/blob/master/torch_xla/runtime.py#L83) to recreate the behavior of that function, minus the manual initialization
…py for `torch-xla>=2.5` (#20442) * Replace `using_pjrt()` xla runtime `device_type()` check with in xla.py Fixes #20419 `torch_xla.runtime.using_pjrt()` is removed in pytorch/xla#7787 This PR replaces references to that function with a check to [`device_type()`](https://github.com/pytorch/xla/blob/master/torch_xla/runtime.py#L83) to recreate the behavior of that function, minus the manual initialization * Added tests/refactored for version compat * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * precommit --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Resolves #7730.
Now device are automatically detected when we import torch_xla: