Fix DDP on XLA #16020

Liyang90 · 2022-12-12T18:32:47Z

Fix DDP on XLA by inserting xm.reduce_gradients in closure instead of using xm.optimizer_step. Adding a xm.mark_step() call after optimizer.step for performance.

What does this PR do?

It fixes DDP correctness when using tpu accelerator. The related issue is 15878

It also includes fix for "Issue 1" in 15884

Fixes #\15878

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

None

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?

Make sure you had fun coding 🙃

Fix DDP on XLA by inserting `xm.reduce_gradients` in `closure` instead of using `xm.optimizer_step`. Adding a `xm.mark_step()` call after `optimizer.step` for performance.

src/pytorch_lightning/plugins/precision/tpu.py

Borda

lgth, just can we add a test for this discovery?

awaelchli · 2023-01-03T10:24:35Z

@Borda I added the test.

Co-authored-by: awaelchli <[email protected]>

* [App] Introduce basic auth to Lightning CLI (#16105) * Introduce basic auth to Lightning CLI for app creation * Parsing creds added * Adding auth field to app instance body * Adding tests * Adding changelog entry * Adding more tests * Update runtime.py * Setting auth on update * Fix test * Update lightning-cloud dep * Apply suggestions from code review Co-authored-by: Adrian Wälchli <[email protected]> * Update runtime.py * Fix for release * Update base.txt Co-authored-by: Jirka Borovec <[email protected]> Co-authored-by: Adrian Wälchli <[email protected]> * docs: fix order of on_fit_start() hook (#16180) Fixes #16170 * docs: updated broken links (#16191) Co-authored-by: Shashwat <shashwat> Fixes #16186 * Always use the local rank zero imports (#16178) * update BYOC documentation with AWS details (#16044) Co-authored-by: Adrian Wälchli <[email protected]> * Using internal ip + port in a load balancer instead of URL exposed (#16119) Co-authored-by: thomas chaton <[email protected]> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> * Remove CUDA_LAUNCH_BLOCKING from Lite tests (#16177) * simplify torch.Tensor (#16190) * Fix inclusion of `model_parallel` document (#16197) fix link to gpu/advanced section * Update numpy requirement from <1.23.1,>=1.17.2 to >=1.17.2,<1.24.1 in /requirements (#16199) Update numpy requirement in /requirements Updates the requirements on [numpy](https://github.com/numpy/numpy) to permit the latest version. - [Release notes](https://github.com/numpy/numpy/releases) - [Changelog](https://github.com/numpy/numpy/blob/main/doc/RELEASE_WALKTHROUGH.rst) - [Commits](numpy/numpy@v1.17.2...v1.24.0) --- updated-dependencies: - dependency-name: numpy dependency-type: direct:production ... Signed-off-by: dependabot[bot] <[email protected]> Signed-off-by: dependabot[bot] <[email protected]> Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> * Update s3fs requirement from <2022.8.3,>=2022.5.0 to >=2022.5.0,<2022.11.1 in /requirements (#16198) Update s3fs requirement in /requirements Updates the requirements on [s3fs](https://github.com/fsspec/s3fs) to permit the latest version. - [Release notes](https://github.com/fsspec/s3fs/releases) - [Changelog](https://github.com/fsspec/s3fs/blob/main/release-procedure.md) - [Commits](fsspec/s3fs@2022.5.0...2022.11.0) --- updated-dependencies: - dependency-name: s3fs dependency-type: direct:production ... Signed-off-by: dependabot[bot] <[email protected]> Signed-off-by: dependabot[bot] <[email protected]> Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> * ci: upload only with release (#16194) * Update endpoint access examples: added info about accessing auth-protected apps (#16145) * Skip a failing Bagua test for manual optimization (#16225) * Fix type error when dividing chunk size in colossalai strategy (#16212) Co-authored-by: awaelchli <[email protected]> * Fix DDP on XLA (#16020) Co-authored-by: awaelchli <[email protected]> * Added support and test for custom artifact names in WandbLogger (#16173) * Update ipython[all] requirement from <8.6.1 to <8.7.1 in /requirements (#16220) Updates the requirements on [ipython[all]](https://github.com/ipython/ipython) to permit the latest version. - [Release notes](https://github.com/ipython/ipython/releases) - [Commits](ipython/ipython@rel-0.8.4...8.7.0) --- updated-dependencies: - dependency-name: ipython[all] dependency-type: direct:production ... Signed-off-by: dependabot[bot] <[email protected]> Signed-off-by: dependabot[bot] <[email protected]> Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> * Improvements to checkpoint migration (#16233) Co-authored-by: Justus Schock <[email protected]> * Support arbitrary `Optimizable`s as optimizers (#16189) * Differences with master Signed-off-by: dependabot[bot] <[email protected]> Co-authored-by: Dmitry Frolov <[email protected]> Co-authored-by: Jirka Borovec <[email protected]> Co-authored-by: Adrian Wälchli <[email protected]> Co-authored-by: David Gilbertson <[email protected]> Co-authored-by: Shashwat Agrawal <[email protected]> Co-authored-by: Raphael Randschau <[email protected]> Co-authored-by: thomas chaton <[email protected]> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> Co-authored-by: HELSON <[email protected]> Co-authored-by: Liyang90 <[email protected]> Co-authored-by: Manan Goel <[email protected]> Co-authored-by: Justus Schock <[email protected]>

Fix DDP on XLA

37f8166

Fix DDP on XLA by inserting `xm.reduce_gradients` in `closure` instead of using `xm.optimizer_step`. Adding a `xm.mark_step()` call after `optimizer.step` for performance.

Liyang90 requested review from williamFalcon, awaelchli, carmocca and justusschock as code owners December 12, 2022 18:32

github-actions bot added the pl Generic label for PyTorch Lightning package label Dec 12, 2022

awaelchli approved these changes Dec 12, 2022

View reviewed changes

src/pytorch_lightning/plugins/precision/tpu.py Show resolved Hide resolved

awaelchli added bug Something isn't working accelerator: tpu Tensor Processing Unit labels Dec 12, 2022

awaelchli added this to the v1.8.x milestone Dec 12, 2022

awaelchli added 2 commits December 12, 2022 22:39

fix mypy error

de01b1a

add changelog entry

a9ad469

Borda requested changes Dec 13, 2022

View reviewed changes

awaelchli mentioned this pull request Dec 14, 2022

TQDMProgressBar refresh forces TPU to recompile compute graph #15743

Open

mergify bot added the has conflicts label Dec 20, 2022

Merge branch 'master' into xla_ddp

2e4c131

mergify bot removed the has conflicts label Dec 20, 2022

awaelchli added the community This PR is from the community label Dec 22, 2022

mergify bot added the has conflicts label Dec 22, 2022

Merge branch 'master' into xla_ddp

36693a7

mergify bot removed the has conflicts label Dec 28, 2022

add test

101c6c2

awaelchli requested a review from Borda January 3, 2023 08:58

Borda approved these changes Jan 3, 2023

View reviewed changes

Borda enabled auto-merge (squash) January 3, 2023 11:47

Merge branch 'master' into xla_ddp

29b9e7d

mergify bot added the ready PRs ready to be merged label Jan 3, 2023

justusschock approved these changes Jan 3, 2023

View reviewed changes

Merge branch 'master' into xla_ddp

f16dd75

mergify bot added has conflicts and removed ready PRs ready to be merged labels Jan 3, 2023

Merge branch 'master' into xla_ddp

9f9e470

mergify bot added ready PRs ready to be merged and removed has conflicts ready PRs ready to be merged labels Jan 3, 2023

Borda merged commit ec641f1 into Lightning-AI:master Jan 3, 2023

carmocca pushed a commit that referenced this pull request Jan 4, 2023

Fix DDP on XLA (#16020)

d8e3435

Co-authored-by: awaelchli <[email protected]>

Liyang90 deleted the xla_ddp branch January 31, 2023 00:04

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Fix DDP on XLA #16020

Fix DDP on XLA #16020

Liyang90 commented Dec 12, 2022 •

edited

Loading

Borda left a comment

awaelchli commented Jan 3, 2023

Fix DDP on XLA #16020

Fix DDP on XLA #16020

Conversation

Liyang90 commented Dec 12, 2022 • edited Loading

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?

Borda left a comment

Choose a reason for hiding this comment

awaelchli commented Jan 3, 2023

Liyang90 commented Dec 12, 2022 •

edited

Loading