Disable non blocking to device with MPS by j0rd1smit · Pull Request #14368 · Lightning-AI/pytorch-lightning

j0rd1smit · 2022-08-23T19:12:57Z

What does this PR do?

This PR ensures that the race condition bug in Pytorch (pytorch/pytorch#83015) does not affect lightning. As discussed with @justusschock in #13285 this PR disables non-blocking moves to MPS devices since it can result in tensors with different values after the moves to the MPS device. To verify it works now, the test case test_data_is_not_changed_after_move_to_mps_device has been added.

Fixes #13285

Does your PR introduce any breaking changes? If yes, please list them.

No.

Before submitting

Was this discussed/approved via a GitHub issue? (not for typos and docs)
Did you read the contributor guideline, Pull Request section?
Did you make sure your PR does only one thing, instead of bundling different changes together?
Did you make sure to update the documentation with your changes? (if necessary)
Did you write any new necessary tests? (not for typos and docs)
Did you verify new and existing tests pass locally with your changes?
Did you list all the breaking changes introduced by this pull request?
Did you update the CHANGELOG? (not for typos, docs, test updates, or minor internal changes/refactors)

PR review

Anyone in the community is welcome to review the PR.
Before you start reviewing, make sure you have read the review guidelines. In short, see the following bullet-list:

Is this pull request ready for review? (if not, please submit in draft mode)
Check that all items from Before submitting are resolved
Make sure the title is self-explanatory and the description concisely explains the PR
Add labels and milestones (and optionally projects) to the PR so it can be classified

Did you have fun?

Yes, this was my first ever PR 😀. It was quite fun to do. However, since it is my first time, I hope I did everything OK. If not, let me know. I'm happy to learn.

…e_mps

carmocca

Somebody will need to manually run the CI on an MPS device

codecov · 2022-08-23T23:04:42Z

Codecov Report

Merging #14368 (39af0b9) into master (33a5ed9) will decrease coverage by 3%.
The diff coverage is 100%.

❗ Current head 39af0b9 differs from pull request most recent head 523b3a8. Consider uploading reports for the commit 523b3a8 to get more accurate results

@@            Coverage Diff             @@
##           master   #14368      +/-   ##
==========================================
- Coverage      79%      76%      -3%     
==========================================
  Files         111      332     +221     
  Lines        7258    26894   +19636     
==========================================
+ Hits         5740    20432   +14692     
- Misses       1518     6462    +4944

awaelchli · 2022-08-24T07:16:33Z

Somebody will need to manually run the CI on an MPS device

@carmocca Easier said than done :) See #14012
But perhaps we can run some selected tests related to device transfers

justusschock · 2022-08-25T09:49:37Z

@j0rd1smit seems like you fixed it. I think it's reasonable to update the type, but I'd do so in another PR as this might require further typing changes for mypy. If you want to tackle that as well, feel free to do so :)

…e_mps

j0rd1smit · 2022-08-25T11:14:44Z

@justusschock agreed that should be a separate PR. So guess this PR is ready.

…e_mps

* disable non-blocking for mps due to race condition bug * fixed typo * fixed: unknown mps device for non arm systems * Removed unrobust test case * moved _MPS_DEVICES such that we used in apply_func * Resolve circular dependencies * Comment rewording * changed torchElasticEnvironment to a global import * simplified if statement to blocking device type * Added change to CHANGELOG * Update src/pytorch_lightning/utilities/apply_func.py * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fixed mypy not detecting casting of device * Moved check into if statement to mainain original behavior Co-authored-by: Carlos Mocholí <carlossmocholi@gmail.com> Co-authored-by: Justus Schock <12886177+justusschock@users.noreply.github.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Jirka Borovec <Borda@users.noreply.github.com>

j0rd1smit added 2 commits August 23, 2022 20:53

disable non-blocking for mps due to race condition bug

92ea1a0

fixed typo

437a366

j0rd1smit requested review from Borda, awaelchli, carmocca, kaushikb11, rohitgr7, tchaton and williamFalcon as code owners August 23, 2022 19:12

github-actions Bot added the pl Generic label for PyTorch Lightning package label Aug 23, 2022

fixed: unknown mps device for non arm systems

401f97a

j0rd1smit changed the title ~~Bugfix/13285 disable non blocking to device mps~~ Aug 23, 2022

carmocca reviewed Aug 23, 2022

View reviewed changes

Comment thread src/pytorch_lightning/utilities/apply_func.py Outdated

Comment thread tests/tests_pytorch/accelerators/test_mps.py Outdated

Comment thread tests/tests_pytorch/accelerators/test_mps.py Outdated

carmocca added community This PR is from the community accelerator: mps Apple Silicon GPU labels Aug 23, 2022

carmocca added this to the pl:1.7.x milestone Aug 23, 2022

carmocca added the bug Something isn't working label Aug 23, 2022

j0rd1smit added 2 commits August 23, 2022 22:23

Removed unrobust test case

a112f95

moved _MPS_DEVICES such that we used in apply_func

1e63b36

j0rd1smit requested a review from justusschock as a code owner August 23, 2022 20:35

j0rd1smit and others added 3 commits August 23, 2022 22:35

Merge branch 'master' into bugfix/13285_disable_non_blocking_to_devic…

e56026d

…e_mps

Resolve circular dependencies

02220a7

Comment rewording

fa5ea1b

carmocca approved these changes Aug 23, 2022

View reviewed changes

carmocca changed the title ~~[WIP] Bugfix/13285 disable non blocking to device mps~~ Aug 23, 2022

carmocca self-assigned this Aug 23, 2022

justusschock approved these changes Aug 24, 2022

View reviewed changes

Comment thread src/pytorch_lightning/utilities/apply_func.py Outdated

Comment thread src/pytorch_lightning/utilities/apply_func.py Outdated

Comment thread src/pytorch_lightning/utilities/apply_func.py Outdated

mergify Bot added the ready to be merged PRs ready to be merged label Aug 24, 2022

Merge branch 'master' into bugfix/13285_disable_non_blocking_to_devic…

39af0b9

…e_mps

akihironitta approved these changes Aug 25, 2022

View reviewed changes

mergify Bot added has conflicts and removed ready to be merged PRs ready to be merged labels Aug 25, 2022

Merge branch 'master' into bugfix/13285_disable_non_blocking_to_devic…

523b3a8

…e_mps

mergify Bot added ready to be merged PRs ready to be merged and removed has conflicts ready to be merged PRs ready to be merged labels Aug 26, 2022

justusschock enabled auto-merge (squash) August 26, 2022 06:00

mergify Bot added has conflicts and removed ready to be merged PRs ready to be merged labels Aug 26, 2022

Merge branch 'master' into bugfix/13285_disable_non_blocking_to_devic…

f364736

…e_mps

mergify Bot added ready to be merged PRs ready to be merged has conflicts and removed has conflicts ready to be merged PRs ready to be merged labels Aug 26, 2022

Merge branch 'master' into bugfix/13285_disable_non_blocking_to_devic…

fa87e28

…e_mps

mergify Bot added ready to be merged PRs ready to be merged and removed has conflicts ready to be merged PRs ready to be merged labels Aug 26, 2022

awaelchli reviewed Aug 26, 2022

View reviewed changes

Comment thread src/pytorch_lightning/utilities/apply_func.py

justusschock merged commit cced335 into Lightning-AI:master Aug 26, 2022

Navigation Menu

Uh oh!

Disable non blocking to device with MPS - #14368

Disable non blocking to device with MPS#14368
justusschock merged 19 commits into
Lightning-AI:masterfrom
j0rd1smit:bugfix/13285_disable_non_blocking_to_device_mps

j0rd1smit commented Aug 23, 2022

Uh oh!

Uh oh!

Uh oh!

carmocca left a comment

codecov Bot commented Aug 23, 2022 •

edited

Loading

Uh oh!

Uh oh!

Uh oh!

awaelchli commented Aug 24, 2022

justusschock commented Aug 25, 2022

j0rd1smit commented Aug 25, 2022

Uh oh!

Labels

6 participants

Uh oh!

Conversation

j0rd1smit commented Aug 23, 2022

What does this PR do?

Does your PR introduce any breaking changes? If yes, please list them.

Before submitting

PR review

Did you have fun?

Uh oh!

Uh oh!

Uh oh!

carmocca left a comment

Choose a reason for hiding this comment

codecov Bot commented Aug 23, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

Uh oh!

Uh oh!

awaelchli commented Aug 24, 2022

justusschock commented Aug 25, 2022

j0rd1smit commented Aug 25, 2022

Uh oh!

Labels

6 participants

codecov Bot commented Aug 23, 2022 •

edited

Loading