Fix BF16 teardown for TPU precision plugin #10990

awaelchli · 2021-12-08T06:02:18Z

What does this PR do?

The env variable gets set in the tpu bf16 precision plugin connect() method but gets deleted again in the TPUSpawnPlugin before training starts. The correct behavior is to set the variable at the beginning of a run and remove it again at the end of a run (teardown). For this reason, we are introducing a teardown() method on the PrecisionPlugin base class (missing so far due to lack of use case).

Before submitting

Was this discussed/approved via a GitHub issue? (not for typos and docs) 2/n Simplify spawn plugins: Spawn immediately #10896
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 update the CHANGELOG? (not for typos, docs, test updates, or internal minor changes/refactorings)

PR review

Anyone in the community is free to review the PR once the tests have passed.
Before you start reviewing make sure you have read 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?

I made sure I had fun coding 🙃

Part of #1 (it's a lie, this is just here to avoid noisy GitHub bot)

cc @Borda @kaushikb11

for more information, see https://pre-commit.ci

…re/simple-spawn

for more information, see https://pre-commit.ci

…tor/spawn/simple-spawn

for more information, see https://pre-commit.ci

…tor/spawn/simple-spawn

Co-authored-by: Rohit Gupta <rohitgr1998@gmail.com>

tests/models/test_tpu.py

tchaton

LGTM !

CHANGELOG.md

pytorch_lightning/plugins/training_type/training_type_plugin.py

tests/models/test_tpu.py

Co-authored-by: thomas chaton <thomas@grid.ai>

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Justus Schock <12886177+justusschock@users.noreply.github.com> Co-authored-by: Rohit Gupta <rohitgr1998@gmail.com> Co-authored-by: Kaushik B <45285388+kaushikb11@users.noreply.github.com> Co-authored-by: thomas chaton <thomas@grid.ai>

awaelchli and others added 30 commits October 20, 2021 04:53

improve spawn queue

d05363f

clean up

d650e26

[pre-commit.ci] auto fixes from pre-commit.com hooks

5fda23a

for more information, see https://pre-commit.ci

Merge branch 'master' into feature/simple-spawn

d6b4a34

fix

bcfb853

[pre-commit.ci] auto fixes from pre-commit.com hooks

97b4bf6

for more information, see https://pre-commit.ci

rename

38b3a54

delete dead code

955b6c8

Merge remote-tracking branch 'origin/feature/simple-spawn' into featu…

13393e8

…re/simple-spawn

clean up

f3216b2

update lite

2d00231

retain the queue interface in hooks

7aa3646

update tests

fb0c0d8

Merge branch 'master' into feature/simple-spawn

1bc59ae

_notebooks

7e6c75e

reset notebooks

b7efc50

avoid circular import

84ca8b4

fix unused imports

965c724

reset debugging script

1aae8dd

typing _ExtraQueue

4b998db

bring changes to tpu_spawn plugin

5871a4b

unify

aa76840

remove dead code

37f9db9

[pre-commit.ci] auto fixes from pre-commit.com hooks

d68cb35

for more information, see https://pre-commit.ci

remove queue from tpu spawn

dd80be9

type annotation for new_process

f97eee8

Merge remote-tracking branch 'origin/feature/simple-spawn' into refac…

ad61d27

…tor/spawn/simple-spawn

[pre-commit.ci] auto fixes from pre-commit.com hooks

459121e

for more information, see https://pre-commit.ci

unused imports

72535ff

Merge remote-tracking branch 'origin/feature/simple-spawn' into refac…

3095da9

…tor/spawn/simple-spawn

mergify bot added ready PRs ready to be merged has conflicts labels Dec 14, 2021

awaelchli and others added 2 commits December 14, 2021 16:42

Update tests/plugins/precision/test_tpu_bf16_plugin.py

4e1fd18

Co-authored-by: Rohit Gupta <rohitgr1998@gmail.com>

Merge branch 'master' into bugfix/bf16-env-variable

4599c33

carmocca reviewed Dec 14, 2021

View reviewed changes

tests/models/test_tpu.py Show resolved Hide resolved

awaelchli mentioned this pull request Dec 16, 2021

Strategies and Accelerator teardown logic clean up #11091

Closed

awaelchli added 2 commits December 21, 2021 03:58

Merge branch 'master' into bugfix/bf16-env-variable

a891411

rm

6a7f462

mergify bot removed the has conflicts label Dec 21, 2021

tchaton approved these changes Dec 21, 2021

View reviewed changes

CHANGELOG.md Outdated Show resolved Hide resolved

CHANGELOG.md Outdated Show resolved Hide resolved

pytorch_lightning/plugins/training_type/training_type_plugin.py Show resolved Hide resolved

tests/models/test_tpu.py Show resolved Hide resolved

kaushikb11 and others added 2 commits December 21, 2021 14:15

Update CHANGELOG.md

7561465

Co-authored-by: thomas chaton <thomas@grid.ai>

Update CHANGELOG.md

990063b

Co-authored-by: thomas chaton <thomas@grid.ai>

mergify bot added the has conflicts label Dec 21, 2021

awaelchli added 2 commits December 21, 2021 15:02

call missing super().teardown()

ce2be66

Merge branch 'master' into bugfix/bf16-env-variable

a93a15f

mergify bot removed the has conflicts label Dec 21, 2021

remove abstract

4c5a480

Borda approved these changes Dec 21, 2021

View reviewed changes

ananthsub approved these changes Dec 21, 2021

View reviewed changes

Merge branch 'master' into bugfix/bf16-env-variable

ac480f4

mergify bot added the has conflicts label Dec 22, 2021

awaelchli added 2 commits December 22, 2021 04:04

Merge branch 'master' into bugfix/bf16-env-variable

db5c6cb

reorder

3528234

mergify bot removed the has conflicts label Dec 22, 2021

awaelchli enabled auto-merge (squash) December 22, 2021 03:13

awaelchli merged commit ba8e7cd into master Dec 22, 2021

awaelchli deleted the bugfix/bf16-env-variable branch December 22, 2021 03:47

awaelchli mentioned this pull request Jan 5, 2022

Weekly Patch Release v1.5.8 #11314

Merged

12 tasks

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Fix BF16 teardown for TPU precision plugin #10990

Fix BF16 teardown for TPU precision plugin #10990

awaelchli commented Dec 8, 2021 •

edited by github-actions bot

tchaton left a comment

Fix BF16 teardown for TPU precision plugin #10990

Fix BF16 teardown for TPU precision plugin #10990

Conversation

awaelchli commented Dec 8, 2021 • edited by github-actions bot

What does this PR do?

Before submitting

PR review

Did you have fun?

tchaton left a comment

Choose a reason for hiding this comment

awaelchli commented Dec 8, 2021 •

edited by github-actions bot