[TPU] v4 support #17227

carmocca · 2023-03-29T20:47:40Z

What does this PR do?

Updates the codebase and test suite to support the v4 nodes out of the box.

In these, "auto" needs to default to 4, not 8

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

src/lightning/fabric/accelerators/tpu.py

src/lightning/pytorch/accelerators/tpu.py

github-actions · 2023-03-29T21:04:44Z

⚡ Required checks status: All passing 🟢

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/accelerators/tpu.py, src/lightning/fabric/connector.py, src/lightning/fabric/plugins/environments/xla.py, src/lightning/fabric/strategies/single_device.py, src/lightning/fabric/strategies/single_tpu.py, src/lightning/fabric/utilities/device_parser.py, src/lightning/pytorch/accelerators/tpu.py, src/lightning/pytorch/strategies/single_device.py, src/lightning/pytorch/strategies/single_tpu.py, src/lightning/pytorch/strategies/xla.py, src/lightning/pytorch/trainer/connectors/accelerator_connector.py, tests/tests_pytorch/accelerators/test_tpu.py, tests/tests_pytorch/callbacks/test_device_stats_monitor.py, tests/tests_pytorch/conftest.py, tests/tests_pytorch/models/test_tpu.py, tests/tests_pytorch/profilers/test_xla_profiler.py, tests/tests_pytorch/strategies/test_registry.py, tests/tests_pytorch/trainer/connectors/test_accelerator_connector.py, tests/tests_pytorch/trainer/properties/test_estimated_stepping_batches.py.

🟢 pytorch_lightning: Azure GPU

Check ID	Status
pytorch-lightning (GPUs)	success	✅

These checks are required after the changes to src/lightning/pytorch/accelerators/tpu.py, src/lightning/pytorch/strategies/single_device.py, src/lightning/pytorch/strategies/single_tpu.py, src/lightning/pytorch/strategies/xla.py, src/lightning/pytorch/trainer/connectors/accelerator_connector.py, tests/tests_pytorch/accelerators/test_tpu.py, tests/tests_pytorch/callbacks/test_device_stats_monitor.py, tests/tests_pytorch/conftest.py, tests/tests_pytorch/models/test_tpu.py, tests/tests_pytorch/profilers/test_xla_profiler.py, tests/tests_pytorch/strategies/test_registry.py, tests/tests_pytorch/trainer/connectors/test_accelerator_connector.py, tests/tests_pytorch/trainer/properties/test_estimated_stepping_batches.py, src/lightning/fabric/accelerators/tpu.py, src/lightning/fabric/connector.py, src/lightning/fabric/plugins/environments/xla.py, src/lightning/fabric/strategies/single_device.py, src/lightning/fabric/strategies/single_tpu.py, src/lightning/fabric/utilities/device_parser.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/accelerators/tpu.py, src/lightning/fabric/connector.py, src/lightning/fabric/plugins/environments/xla.py, src/lightning/fabric/strategies/single_device.py, src/lightning/fabric/strategies/single_tpu.py, src/lightning/fabric/utilities/device_parser.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/accelerators/tpu.py, src/lightning/pytorch/strategies/single_device.py, src/lightning/pytorch/strategies/single_tpu.py, src/lightning/pytorch/strategies/xla.py, src/lightning/pytorch/trainer/connectors/accelerator_connector.py, docs/source-pytorch/accelerators/tpu_advanced.rst, docs/source-pytorch/accelerators/tpu_basic.rst, docs/source-pytorch/accelerators/tpu_faq.rst, docs/source-pytorch/accelerators/tpu_intermediate.rst.

🟢 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/accelerators/tpu.py, src/lightning/fabric/connector.py, src/lightning/fabric/plugins/environments/xla.py, src/lightning/fabric/strategies/single_device.py, src/lightning/fabric/strategies/single_tpu.py, src/lightning/fabric/utilities/device_parser.py, tests/tests_fabric/accelerators/test_tpu.py, tests/tests_fabric/conftest.py, tests/tests_fabric/plugins/environments/test_xla.py, tests/tests_fabric/strategies/test_xla.py, tests/tests_fabric/test_connector.py, tests/tests_fabric/test_fabric.py, tests/tests_fabric/utilities/test_device_parser.py.

🟢 lightning_fabric: Azure GPU

Check ID	Status
lightning-fabric (GPUs)	success	✅

These checks are required after the changes to src/lightning/fabric/accelerators/tpu.py, src/lightning/fabric/connector.py, src/lightning/fabric/plugins/environments/xla.py, src/lightning/fabric/strategies/single_device.py, src/lightning/fabric/strategies/single_tpu.py, src/lightning/fabric/utilities/device_parser.py, tests/tests_fabric/accelerators/test_tpu.py, tests/tests_fabric/conftest.py, tests/tests_fabric/plugins/environments/test_xla.py, tests/tests_fabric/strategies/test_xla.py, tests/tests_fabric/test_connector.py, tests/tests_fabric/test_fabric.py, tests/tests_fabric/utilities/test_device_parser.py.

🟢 mypy

Check ID	Status
mypy	success	✅

These checks are required after the changes to src/lightning/fabric/accelerators/tpu.py, src/lightning/fabric/connector.py, src/lightning/fabric/plugins/environments/xla.py, src/lightning/fabric/strategies/single_device.py, src/lightning/fabric/strategies/single_tpu.py, src/lightning/fabric/utilities/device_parser.py, src/lightning/pytorch/accelerators/tpu.py, src/lightning/pytorch/strategies/single_device.py, src/lightning/pytorch/strategies/single_tpu.py, src/lightning/pytorch/strategies/xla.py, src/lightning/pytorch/trainer/connectors/accelerator_connector.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)	success	✅
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)	success	✅
install-pkg (windows-2022, notset, 3.8)	success	✅
install-pkg (windows-2022, notset, 3.10)	success	✅

These checks are required after the changes to src/lightning/fabric/accelerators/tpu.py, src/lightning/fabric/connector.py, src/lightning/fabric/plugins/environments/xla.py, src/lightning/fabric/strategies/single_device.py, src/lightning/fabric/strategies/single_tpu.py, src/lightning/fabric/utilities/device_parser.py, src/lightning/pytorch/accelerators/tpu.py, src/lightning/pytorch/strategies/single_device.py, src/lightning/pytorch/strategies/single_tpu.py, src/lightning/pytorch/strategies/xla.py, src/lightning/pytorch/trainer/connectors/accelerator_connector.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.

codecov · 2023-03-29T23:27:16Z

Codecov Report

Merging #17227 (b9ff9a7) into master (0c02c44) will decrease coverage by 9%.
The diff coverage is 84%.

Additional details and impacted files

@@            Coverage Diff            @@
##           master   #17227     +/-   ##
=========================================
- Coverage      68%      59%     -9%     
=========================================
  Files         433      410     -23     
  Lines       31659    31377    -282     
=========================================
- Hits        21626    18569   -3057     
- Misses      10033    12808   +2775

docs/source-pytorch/accelerators/tpu_basic.rst

tests/tests_fabric/test_fabric.py

tests/tests_fabric/utilities/test_device_parser.py

src/lightning/fabric/accelerators/tpu.py

src/lightning/pytorch/accelerators/tpu.py

tests/tests_fabric/accelerators/test_tpu.py

tests/tests_fabric/test_fabric.py

tests/tests_pytorch/accelerators/test_tpu.py

tests/tests_fabric/test_connector.py

carmocca added feature Is an improvement or enhancement accelerator: tpu Tensor Processing Unit labels Mar 29, 2023

carmocca added this to the 2.1 milestone Mar 29, 2023

carmocca self-assigned this Mar 29, 2023

github-actions bot added fabric lightning.fabric.Fabric pl Generic label for PyTorch Lightning package labels Mar 29, 2023

carmocca commented Mar 29, 2023

View reviewed changes

src/lightning/fabric/accelerators/tpu.py Show resolved Hide resolved

src/lightning/fabric/accelerators/tpu.py Show resolved Hide resolved

src/lightning/fabric/accelerators/tpu.py Show resolved Hide resolved

src/lightning/pytorch/accelerators/tpu.py Show resolved Hide resolved

carmocca marked this pull request as ready for review March 29, 2023 21:04

carmocca requested review from awaelchli, justusschock and williamFalcon as code owners March 29, 2023 21:04

carmocca requested review from edenlightning, lantiga and Borda as code owners March 29, 2023 23:26

Update

d656b2e

carmocca force-pushed the carmocca/tpu-env branch from bccbbe5 to d656b2e Compare March 30, 2023 00:21

accident

a2ac44f

carmocca commented Mar 30, 2023

View reviewed changes

docs/source-pytorch/accelerators/tpu_basic.rst Show resolved Hide resolved

carmocca added 3 commits March 30, 2023 03:39

Fixes

037df23

Update

ee33a21

3.8 skips

33354e8

carmocca force-pushed the carmocca/tpu-env branch from 1c27645 to 33354e8 Compare March 30, 2023 02:46

carmocca added 4 commits March 30, 2023 05:36

Fix fabric tests

23bb0b6

CPU

427ea66

Fix fabric test

7ba20ee

Non-standalone

105f44a

carmocca commented Mar 30, 2023

View reviewed changes

tests/tests_fabric/test_fabric.py Show resolved Hide resolved

carmocca commented Mar 30, 2023

View reviewed changes

tests/tests_fabric/utilities/test_device_parser.py Show resolved Hide resolved

awaelchli approved these changes Mar 30, 2023

View reviewed changes

carmocca added 2 commits March 30, 2023 14:55

windowss

d4d06b6

Adrian review. CHANGELOG

4940545

carmocca changed the title ~~TPU v4 support~~ [TPU] v4 support Mar 30, 2023

Merge branch 'master' into carmocca/tpu-env

2027a82

Borda approved these changes Apr 11, 2023

View reviewed changes

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

carmocca enabled auto-merge (squash) April 11, 2023 19:14

mergify bot added has conflicts and removed ready PRs ready to be merged labels Apr 11, 2023

Merge branch 'master' into carmocca/tpu-env

b9ff9a7

mergify bot added ready PRs ready to be merged and removed has conflicts ready PRs ready to be merged labels Apr 11, 2023

carmocca merged commit 0489f2e into master Apr 11, 2023

carmocca deleted the carmocca/tpu-env branch April 11, 2023 22:24

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[TPU] v4 support #17227

[TPU] v4 support #17227

carmocca commented Mar 29, 2023 •

edited by github-actions bot

Loading

github-actions bot commented Mar 29, 2023 •

edited

Loading

codecov bot commented Mar 29, 2023 •

edited

Loading

[TPU] v4 support #17227

[TPU] v4 support #17227

Conversation

carmocca commented Mar 29, 2023 • edited by github-actions bot Loading

What does this PR do?

github-actions bot commented Mar 29, 2023 • edited Loading

⚡ Required checks status: All passing 🟢

Groups summary

codecov bot commented Mar 29, 2023 • edited Loading

Codecov Report

carmocca commented Mar 29, 2023 •

edited by github-actions bot

Loading

github-actions bot commented Mar 29, 2023 •

edited

Loading

codecov bot commented Mar 29, 2023 •

edited

Loading