Fixes around `Strategy.set_world_ranks` #16966

awaelchli · 2023-03-06T14:28:57Z

What does this PR do?

Avoid calling set_world_ranks for XLAStrategy.

cc @Borda @carmocca @justusschock @awaelchli @JackCaoG @steventk-g @Liyang90

awaelchli · 2023-03-17T12:34:43Z

src/lightning/pytorch/strategies/xla.py

@@ -211,13 +211,10 @@ def reduce(

    def setup_distributed(self) -> None:
        self._launched = True
-        self.set_world_ranks()
        rank_zero_only.rank = self.global_rank

    def set_world_ranks(self) -> None:


The accelerator connector will call this method. But we don't want the XLA strategy setting the rank (yet). Since XLAStrategy inherits from DDPStrategy, we need to leave it empty here.

awaelchli · 2023-03-17T12:34:59Z

src/lightning/pytorch/strategies/xla.py

@@ -211,13 +211,10 @@ def reduce(

    def setup_distributed(self) -> None:
        self._launched = True
-        self.set_world_ranks()


the line below already covers what was done in this method.

src/lightning/pytorch/strategies/xla.py

github-actions · 2023-04-13T14:51:37Z

⛈️ Required checks status: Has failure 🔴

Warning
This job will need to be re-run to merge your PR. If you do not have write access to the repository, you can ask Lightning-AI/lai-frameworks to re-run it. If you push a new commit, all of CI will re-trigger.

Groups summary

🟢 pytorch_lightning: Tests workflow

Check ID	Status
pl-cpu (macOS-11, lightning, 3.8, 1.11)	success	✅
pl-cpu (macOS-11, lightning, 3.9, 1.12)	success	✅
pl-cpu (macOS-11, lightning, 3.10, 1.13)	success	✅
pl-cpu (macOS-11, lightning, 3.10, 2.0)	success	✅
pl-cpu (macOS-11, lightning, 3.8, 1.11, oldest)	success	✅
pl-cpu (ubuntu-20.04, lightning, 3.8, 1.11)	success	✅
pl-cpu (ubuntu-20.04, lightning, 3.9, 1.12)	success	✅
pl-cpu (ubuntu-20.04, lightning, 3.10, 1.13)	success	✅
pl-cpu (ubuntu-20.04, lightning, 3.10, 2.0)	success	✅
pl-cpu (ubuntu-20.04, lightning, 3.8, 1.11, oldest)	success	✅
pl-cpu (windows-2022, lightning, 3.8, 1.11)	success	✅
pl-cpu (windows-2022, lightning, 3.9, 1.12)	success	✅
pl-cpu (windows-2022, lightning, 3.10, 1.13)	success	✅
pl-cpu (windows-2022, lightning, 3.10, 2.0)	success	✅
pl-cpu (windows-2022, lightning, 3.8, 1.11, oldest)	success	✅
pl-cpu (macOS-11, pytorch, 3.8, 1.13)	success	✅
pl-cpu (ubuntu-20.04, pytorch, 3.8, 1.13)	success	✅
pl-cpu (windows-2022, pytorch, 3.8, 1.13)	success	✅

These checks are required after the changes to src/lightning/fabric/connector.py, src/lightning/fabric/plugins/environments/lsf.py, src/lightning/fabric/strategies/ddp.py, src/lightning/fabric/strategies/deepspeed.py, src/lightning/fabric/strategies/fsdp.py, src/lightning/fabric/strategies/xla.py, src/lightning/pytorch/strategies/ddp.py, src/lightning/pytorch/strategies/deepspeed.py, src/lightning/pytorch/strategies/fsdp.py, src/lightning/pytorch/strategies/xla.py.

🟢 pytorch_lightning: Azure GPU

Check ID	Status
pytorch-lightning (GPUs)	success	✅

These checks are required after the changes to src/lightning/pytorch/strategies/ddp.py, src/lightning/pytorch/strategies/deepspeed.py, src/lightning/pytorch/strategies/fsdp.py, src/lightning/pytorch/strategies/xla.py, src/lightning/fabric/connector.py, src/lightning/fabric/plugins/environments/lsf.py, src/lightning/fabric/strategies/ddp.py, src/lightning/fabric/strategies/deepspeed.py, src/lightning/fabric/strategies/fsdp.py, src/lightning/fabric/strategies/xla.py.

🟢 fabric: Docs

Check ID	Status
make-doctest (fabric)	success	✅
make-html (fabric)	success	✅

These checks are required after the changes to src/lightning/fabric/connector.py, src/lightning/fabric/plugins/environments/lsf.py, src/lightning/fabric/strategies/ddp.py, src/lightning/fabric/strategies/deepspeed.py, src/lightning/fabric/strategies/fsdp.py, src/lightning/fabric/strategies/xla.py.

🟢 pytorch_lightning: Docs

Check ID	Status
make-doctest (pytorch)	success	✅
make-html (pytorch)	success	✅

These checks are required after the changes to src/lightning/pytorch/strategies/ddp.py, src/lightning/pytorch/strategies/deepspeed.py, src/lightning/pytorch/strategies/fsdp.py, src/lightning/pytorch/strategies/xla.py.

🟢 lightning_fabric: CPU workflow

Check ID	Status
fabric-cpu (macOS-11, lightning, 3.8, 1.11)	success	✅
fabric-cpu (macOS-11, lightning, 3.9, 1.12)	success	✅
fabric-cpu (macOS-11, lightning, 3.10, 1.13)	success	✅
fabric-cpu (macOS-11, lightning, 3.10, 2.0)	success	✅
fabric-cpu (macOS-11, lightning, 3.8, 1.11, oldest)	success	✅
fabric-cpu (ubuntu-20.04, lightning, 3.8, 1.11)	success	✅
fabric-cpu (ubuntu-20.04, lightning, 3.9, 1.12)	success	✅
fabric-cpu (ubuntu-20.04, lightning, 3.10, 1.13)	success	✅
fabric-cpu (ubuntu-20.04, lightning, 3.10, 2.0)	success	✅
fabric-cpu (ubuntu-20.04, lightning, 3.8, 1.11, oldest)	success	✅
fabric-cpu (windows-2022, lightning, 3.8, 1.11)	success	✅
fabric-cpu (windows-2022, lightning, 3.9, 1.12)	success	✅
fabric-cpu (windows-2022, lightning, 3.10, 1.13)	success	✅
fabric-cpu (windows-2022, lightning, 3.10, 2.0)	success	✅
fabric-cpu (windows-2022, lightning, 3.8, 1.11, oldest)	success	✅
fabric-cpu (macOS-11, fabric, 3.8, 1.13)	success	✅
fabric-cpu (ubuntu-20.04, fabric, 3.8, 1.13)	success	✅
fabric-cpu (windows-2022, fabric, 3.8, 1.13)	success	✅

These checks are required after the changes to src/lightning/fabric/connector.py, src/lightning/fabric/plugins/environments/lsf.py, src/lightning/fabric/strategies/ddp.py, src/lightning/fabric/strategies/deepspeed.py, src/lightning/fabric/strategies/fsdp.py, src/lightning/fabric/strategies/xla.py, tests/tests_fabric/test_connector.py.

🟢 lightning_fabric: Azure GPU

Check ID	Status
lightning-fabric (GPUs)	success	✅

These checks are required after the changes to src/lightning/fabric/connector.py, src/lightning/fabric/plugins/environments/lsf.py, src/lightning/fabric/strategies/ddp.py, src/lightning/fabric/strategies/deepspeed.py, src/lightning/fabric/strategies/fsdp.py, src/lightning/fabric/strategies/xla.py, tests/tests_fabric/test_connector.py.

🟢 mypy

Check ID	Status
mypy	success	✅

These checks are required after the changes to src/lightning/fabric/connector.py, src/lightning/fabric/plugins/environments/lsf.py, src/lightning/fabric/strategies/ddp.py, src/lightning/fabric/strategies/deepspeed.py, src/lightning/fabric/strategies/fsdp.py, src/lightning/fabric/strategies/xla.py, src/lightning/pytorch/strategies/ddp.py, src/lightning/pytorch/strategies/deepspeed.py, src/lightning/pytorch/strategies/fsdp.py, src/lightning/pytorch/strategies/xla.py.

🔴 install

Check ID	Status
install-pkg (ubuntu-22.04, app, 3.8)	success	✅
install-pkg (ubuntu-22.04, app, 3.10)	success	✅
install-pkg (ubuntu-22.04, fabric, 3.8)	success	✅
install-pkg (ubuntu-22.04, fabric, 3.10)	success	✅
install-pkg (ubuntu-22.04, pytorch, 3.8)	success	✅
install-pkg (ubuntu-22.04, pytorch, 3.10)	success	✅
install-pkg (ubuntu-22.04, lightning, 3.8)	success	✅
install-pkg (ubuntu-22.04, lightning, 3.10)	success	✅
install-pkg (ubuntu-22.04, notset, 3.8)	success	✅
install-pkg (ubuntu-22.04, notset, 3.10)	success	✅
install-pkg (macOS-12, app, 3.8)	success	✅
install-pkg (macOS-12, app, 3.10)	success	✅
install-pkg (macOS-12, fabric, 3.8)	success	✅
install-pkg (macOS-12, fabric, 3.10)	success	✅
install-pkg (macOS-12, pytorch, 3.8)	success	✅
install-pkg (macOS-12, pytorch, 3.10)	success	✅
install-pkg (macOS-12, lightning, 3.8)	success	✅
install-pkg (macOS-12, lightning, 3.10)	success	✅
install-pkg (macOS-12, notset, 3.8)	success	✅
install-pkg (macOS-12, notset, 3.10)	success	✅
install-pkg (windows-2022, app, 3.8)	success	✅
install-pkg (windows-2022, app, 3.10)	failure	❌
install-pkg (windows-2022, fabric, 3.8)	success	✅
install-pkg (windows-2022, fabric, 3.10)	success	✅
install-pkg (windows-2022, pytorch, 3.8)	success	✅
install-pkg (windows-2022, pytorch, 3.10)	success	✅
install-pkg (windows-2022, lightning, 3.8)	success	✅
install-pkg (windows-2022, lightning, 3.10)	failure	❌
install-pkg (windows-2022, notset, 3.8)	success	✅
install-pkg (windows-2022, notset, 3.10)	failure	❌

These checks are required after the changes to src/lightning/fabric/connector.py, src/lightning/fabric/plugins/environments/lsf.py, src/lightning/fabric/strategies/ddp.py, src/lightning/fabric/strategies/deepspeed.py, src/lightning/fabric/strategies/fsdp.py, src/lightning/fabric/strategies/xla.py, src/lightning/pytorch/strategies/ddp.py, src/lightning/pytorch/strategies/deepspeed.py, src/lightning/pytorch/strategies/fsdp.py, src/lightning/pytorch/strategies/xla.py.

🟢 link-check

Check ID	Status
check-md-links / markdown-link-check	success	✅

These checks are required after the changes to src/lightning/fabric/CHANGELOG.md, src/lightning/pytorch/CHANGELOG.md.

Thank you for your contribution! 💜

Note
This comment is automatically generated and updates for 60 minutes every 180 seconds. If you have any other questions, contact carmocca for help.

awaelchli

The added changes LGTM @carmocca

* don't call set_world_ranks in xla strategy * update * fabric and other strategies * CHANGELOG * Typos * Reuse test --------- Co-authored-by: Carlos Mocholí <[email protected]> (cherry picked from commit 50662eb)

awaelchli added 2 commits March 6, 2023 15:26

don't call set_world_ranks in xla strategy

979fca0

update

475a113

github-actions bot added the pl Generic label for PyTorch Lightning package label Mar 6, 2023

awaelchli added the accelerator: tpu Tensor Processing Unit label Mar 17, 2023

awaelchli added this to the v1.9.x milestone Mar 17, 2023

awaelchli commented Mar 17, 2023

View reviewed changes

carmocca added strategy: xla and removed accelerator: tpu Tensor Processing Unit labels Apr 12, 2023

carmocca reviewed Apr 12, 2023

View reviewed changes

src/lightning/pytorch/strategies/xla.py Show resolved Hide resolved

carmocca reviewed Apr 12, 2023

View reviewed changes

src/lightning/pytorch/strategies/xla.py Show resolved Hide resolved

carmocca added 2 commits April 13, 2023 16:19

Merge branch 'master' into bugfix/xla-rank

7bd503a

fabric and other strategies

f58b3eb

github-actions bot added the fabric lightning.fabric.Fabric label Apr 13, 2023

carmocca changed the title ~~WIP: Avoid calling set_world_ranks in main process for XLA strategy~~ Avoid calling set_world_ranks in main process for XLA strategy Apr 13, 2023

carmocca added 3 commits April 13, 2023 16:41

CHANGELOG

4b685b5

Typos

9f42d7b

Reuse test

4344965

carmocca changed the title ~~Avoid calling set_world_ranks in main process for XLA strategy~~ Fixes around Strategy.set_world_ranks Apr 13, 2023

carmocca approved these changes Apr 13, 2023

View reviewed changes

carmocca marked this pull request as ready for review April 13, 2023 14:51

carmocca requested review from justusschock and williamFalcon as code owners April 13, 2023 14:51

carmocca added bug Something isn't working strategy labels Apr 13, 2023

carmocca mentioned this pull request Apr 13, 2023

[TPU] Add support for PJRT from PyTorch/XLA 2.0 #17352

Merged

7 tasks

awaelchli commented Apr 13, 2023

View reviewed changes

justusschock approved these changes Apr 13, 2023

View reviewed changes

mergify bot added the ready PRs ready to be merged label Apr 13, 2023

carmocca enabled auto-merge (squash) April 13, 2023 15:09

carmocca disabled auto-merge April 13, 2023 15:44

lantiga merged commit 50662eb into master Apr 13, 2023

lantiga deleted the bugfix/xla-rank branch April 13, 2023 15:45

awaelchli added a commit that referenced this pull request Aug 29, 2023

Fix incomplete cherry-pick for #16966

dec9b9f

lantiga pushed a commit that referenced this pull request Aug 30, 2023

Fix incomplete cherry-pick for #16966

efa83da

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Fixes around `Strategy.set_world_ranks` #16966

Fixes around `Strategy.set_world_ranks` #16966

awaelchli commented Mar 6, 2023 •

edited by github-actions bot

Loading

awaelchli Mar 17, 2023

awaelchli Mar 17, 2023

github-actions bot commented Apr 13, 2023 •

edited

Loading

awaelchli left a comment

Fixes around Strategy.set_world_ranks #16966

Fixes around Strategy.set_world_ranks #16966

Conversation

awaelchli commented Mar 6, 2023 • edited by github-actions bot Loading

What does this PR do?

awaelchli Mar 17, 2023

Choose a reason for hiding this comment

awaelchli Mar 17, 2023

Choose a reason for hiding this comment

github-actions bot commented Apr 13, 2023 • edited Loading

⛈️ Required checks status: Has failure 🔴

Groups summary

awaelchli left a comment

Choose a reason for hiding this comment

Fixes around `Strategy.set_world_ranks` #16966

Fixes around `Strategy.set_world_ranks` #16966

awaelchli commented Mar 6, 2023 •

edited by github-actions bot

Loading

github-actions bot commented Apr 13, 2023 •

edited

Loading