[MPS] Add Python Module Bindings for the MPS backend #94417

razarmehr · 2023-02-08T18:01:15Z

This PR is a prerequisite for the upcoming Memory Leak Detection PR.
Enable global manual seeding via torch.manual_seed() + test case
Add torch.mps.synchronize() to wait for MPS stream to finish + test case
Enable the following python interfaces for MPS:
torch.mps.[get_rng_state(), set_rng_state(), synchronize(), manual_seed(), seed()]
Added some test cases in test_mps.py
Added mps.rst to document the torch.mps module.
Fixed the failure with test_public_bindings.py

Description of new files added:

torch/csrc/mps/Module.cpp: implements torch._C module functions for torch.mps and torch.backends.mps.
torch/mps/__init__.py: implements Python bindings for torch.mps module.

pytorch-bot · 2023-02-08T18:01:18Z

🔗 Helpful Links

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

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

❌ 1 Failures

As of commit 2ba30c2:

NEW FAILURES - The following jobs have failed:

linux-bionic-cuda11.7-py3.10-gcc7-sm86 / test (default, 4, 4, linux.g5.4xlarge.nvidia.gpu)

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

aten/src/ATen/detail/MPSHooksInterface.h

torch/csrc/Module.cpp

torch/backends/mps/__init__.py

torch/csrc/api/include/torch/mps.h

torch/csrc/mps/Module.cpp

torch/mps/__init__.py

razarmehr · 2023-02-08T20:52:51Z

I'm working on a fix to address CI failures on other backends caused by this PR.

malfet

Some parts of the code is re-exporting MPS properties which are already available in torch.backends.mps package and also copy-n-pastes anti-patterns from old days of CUDA.(like lazy-init)
Please request re-review when this feedback is addressed.

aten/src/ATen/detail/MPSHooksInterface.h

torch/mps/__init__.py

torch/backends/mps/__init__.py

torch/mps/__init__.py

torch/csrc/mps/Module.cpp

test/test_mps.py

torch/csrc/Module.cpp

albanD · 2023-02-10T15:40:59Z

Sounds pretty good. Thanks for the changes. Only small things.

albanD

Sounds good!
Thanks for the updates.

razarmehr · 2023-02-10T18:26:39Z

@pytorchbot rebase

pytorchmergebot · 2023-02-10T18:29:09Z

@pytorchbot successfully started a rebase job. Check the current status here

pytorchmergebot · 2023-02-10T18:29:13Z

Tried to rebase and push PR #94417, but it was already up to date

razarmehr · 2023-02-10T23:16:34Z

The dynamo CI failure is unrelated to this PR and is due to following failed issue #89395.

FAIL [0.001s]: test_coalesce_reference_cycle_cpu_float64 (__main__.TestSparseCPU)

razarmehr · 2023-02-10T23:16:54Z

@pytorchbot merge -f "Lint and MPS tests are green"

pytorchmergebot · 2023-02-10T23:18:37Z

Merge started

Your change will be merged immediately since you used the force (-f) flag, bypassing any CI checks (ETA: 1-5 minutes).

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

- Enable global manual seeding via torch.manual_seed() + test case - Add torch.mps.synchronize() to wait for MPS stream to finish + test case - Enable the following python interfaces for MPS: torch.mps.get_rng_state() torch.mps.set_rng_state() torch.mps.is_available() torch.mps.synchronize() torch.mps.manual_seed() torch.mps.seed() torch.mps.is_initialized() torch.mps.init()

- This patch fixes a regression caused by recent MPS module interface. Since we now compile torch._C.is_mps_available() only if USE_MPS is defined, then it may cause failures on CUDA (and other devices when USE_MPS is not defined) if we upstream. So, this patch checks if is_mps_available is implemented first and then calls it. - Also use the unique name `default_mps_generator` to avoid conflicts with CPU default generator

- Removed USE_MPS condition when adding python binding definitions - Replaced the global attribute setting of default_mps_generator with torch.C_._mps_get_default_generator() - Removed redundant hasattr() checks - Error out in MPSHooks::isOnMacOS13orNewer() if MPS not available - Removed is_available() and is_macos13_or_newer() from torch.mps (already exists in torch.backends.mps)

This should fix the failure with test_public_bindings.py

- Remove import torch.mps from test_mps - Add torch::mps namespace

pytorchmergebot · 2023-02-12T16:11:51Z

Successfully rebased MPS_Binding_Module onto refs/remotes/origin/viable/strict, please pull locally before adding more changes (for example, via git checkout MPS_Binding_Module && git pull --rebase)

razarmehr · 2023-02-12T16:44:33Z

Thanks @huydhn. The commit 2ba30c2 should fix the problem with the bad forking problem. I tested with pytest -k test_fd_limit_exceeded test_dataloader.py, and it's all good now.
Since that dataloader test isn't part of the CI checks on the PR, I will re-merge the PR with that fix (after all CIs are green), and will monitor PyTorch HUD for any regressions.

huydhn · 2023-02-12T17:56:56Z

You can add ciflow/trunk and it will run all MacOS tests on trunk in the PR for you, including the failed test above. It's also done automatically if the PR is merged normally with @pytorchbot merge instead of force merging

razarmehr · 2023-02-12T20:30:06Z

@pytorchbot merge -g

pytorchmergebot · 2023-02-12T20:31:48Z

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

kulinseth

Looks good.

razarmehr requested review from DenisVieriu97, albanD, kulinseth and malfet February 8, 2023 18:01

pytorch-bot bot added ciflow/mps Run MPS tests (subset of trunk) release notes: mps Release notes category labels Feb 8, 2023

pytorchbot added the open source label Feb 8, 2023

albanD reviewed Feb 8, 2023

View reviewed changes

soulitzer added the triaged This issue has been looked at a team member, and triaged and prioritized into an appropriate module label Feb 8, 2023

malfet requested changes Feb 8, 2023

View reviewed changes

razarmehr force-pushed the MPS_Binding_Module branch 2 times, most recently from 5c1d31f to 60437cd Compare February 9, 2023 03:48

razarmehr requested review from albanD and malfet February 9, 2023 20:47

albanD reviewed Feb 10, 2023

View reviewed changes

test/test_mps.py Outdated Show resolved Hide resolved

torch/csrc/Module.cpp Outdated Show resolved Hide resolved

razarmehr force-pushed the MPS_Binding_Module branch from fefd8ae to dc46d04 Compare February 10, 2023 17:19

razarmehr requested a review from albanD February 10, 2023 17:26

albanD approved these changes Feb 10, 2023

View reviewed changes

pytorchmergebot added the Merged label Feb 10, 2023

pytorchmergebot closed this in beb4f5b Feb 10, 2023

razarmehr added 16 commits February 12, 2023 16:11

Check of is_macos13_or_newer is built into torch.C_ bindings

bc21c58

Fix lint issues

ffc8892

Check for MPS availability before adding generator

74d71c4

Fix lint errors

fb0381e

Remove the C++ APIs mps.cpp and mps.h

aafef75

Add terminating nullptr at the end of _MPSModule_methods list

5e4e49e

Add mps.rst for docs

33b87f8

Asd orphan to mps.rst

bd3ec20

Add mps to index.rst

fb9347f

Fix document issue for manual_seed()

75cc059

Add underscore to make _get_default_mps_generator() private()

015d8f5

This should fix the failure with test_public_bindings.py

Add new header for python_functions() of MPS module

752d369

- Remove import torch.mps from test_mps - Add torch::mps namespace

Fix the failure with multiprocessing dataloader caused by bad forking

2ba30c2

pytorchmergebot force-pushed the MPS_Binding_Module branch from cc771b7 to 2ba30c2 Compare February 12, 2023 16:11

huydhn added the ciflow/trunk Trigger trunk jobs on your pull request label Feb 12, 2023

pytorchmergebot closed this in bdd8f51 Feb 12, 2023

kulinseth reviewed Feb 13, 2023

View reviewed changes

kulinseth reopened this Feb 13, 2023

kulinseth approved these changes Feb 13, 2023

View reviewed changes

kulinseth closed this Feb 13, 2023

awaelchli mentioned this pull request Dec 29, 2023

error running the code: ModuleNotFoundError: No module named 'torch.mps' Lightning-AI/pytorch-lightning#18444

Closed

[MPS] Add Python Module Bindings for the MPS backend #94417

[MPS] Add Python Module Bindings for the MPS backend #94417

Uh oh!

Conversation

razarmehr commented Feb 8, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-bot bot commented Feb 8, 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/94417

❌ 1 Failures

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

razarmehr commented Feb 8, 2023

Uh oh!

malfet left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

albanD commented Feb 10, 2023

Uh oh!

albanD left a comment

Choose a reason for hiding this comment

Uh oh!

razarmehr commented Feb 10, 2023

Uh oh!

pytorchmergebot commented Feb 10, 2023

Uh oh!

pytorchmergebot commented Feb 10, 2023

Uh oh!

razarmehr commented Feb 10, 2023

Uh oh!

razarmehr commented Feb 10, 2023

Uh oh!

pytorchmergebot commented Feb 10, 2023

Merge started

Uh oh!

pytorchmergebot commented Feb 12, 2023

Uh oh!

razarmehr commented Feb 12, 2023

Uh oh!

huydhn commented Feb 12, 2023

Uh oh!

razarmehr commented Feb 12, 2023

Uh oh!

pytorchmergebot commented Feb 12, 2023

Merge started

Uh oh!

kulinseth left a comment

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

8 participants

razarmehr commented Feb 8, 2023 •

edited

Loading

pytorch-bot bot commented Feb 8, 2023 •

edited

Loading