[Fix] TPU Training Type Plugin #6816

kaushikb11 · 2021-04-04T07:29:22Z

What does this PR do?

Fix TPU Colab Hang
TPU Training Plugin Housekeeping
Update Accelerator Connector logic for TPUs.

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 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?

Make sure you had fun coding 🙃

codecov · 2021-04-04T07:48:39Z

Codecov Report

Merging #6816 (f944e62) into master (6dc1078) will decrease coverage by 4%.
The diff coverage is 96%.

@@           Coverage Diff           @@
##           master   #6816    +/-   ##
=======================================
- Coverage      91%     87%    -4%     
=======================================
  Files         192     192            
  Lines       12192   12146    -46     
=======================================
- Hits        11149   10588   -561     
- Misses       1043    1558   +515

tchaton

TPUSpawn post_dispath shoudn't be be removed entirely.

pytorch_lightning/plugins/training_type/tpu_spawn.py

tests/models/test_tpu.py

…ghtning into update/tpu

carmocca

Awesome cleanup!

Is there any way we can test the hang is resolved?

pytorch_lightning/plugins/training_type/single_tpu.py

kaushikb11 · 2021-04-06T07:34:51Z

Is there any way we can test the hang is resolved?

@carmocca Previously, we used to have additional logic for Kaggle/Colab https://github.com/PyTorchLightning/pytorch-lightning/blob/master/pytorch_lightning/plugins/training_type/tpu_spawn.py#L130

As it is no longer required, the existing tests should be enough to test the training flow.

* update readme by v1.2.x (#6728) * [bugfix] Add support for omegaconf and tpu (#6741) * fix_hydra * update changelog Co-authored-by: Your Name <you@example.com> * [docs] Update Bolts link (#6743) * Update Bolts link * Update Bolts link * formt Co-authored-by: Jirka Borovec <Borda@users.noreply.github.com> * Update logic for checking TPUs availability (#6767) * Update logic for checking TPUs availability * fix flake8 * add fix * resolve bug (#6781) * Fix validation progress counter with check_val_every_n_epoch > 1 (#5952) Co-authored-by: rohitgr7 <rohitgr1998@gmail.com> Co-authored-by: Carlos Mocholí <carlossmocholi@gmail.com> * Remove extinct parameters from lightning_module.rst (#6801) Fixes #6800 * Update TPU docs for installation (#6794) * fix boolean check on iterable dataset when len not defined (#6828) * fix iterable dataset len check * update predict and validate * add validate to test * add changelog * add predict * Sanitize `None` params during pruning (#6836) * sanitize none params during pruning * amend * Fix `unfreeze_and_add_param_group` expects `modules` rather than `module` (#6822) * Enforce an epoch scheduler interval when using SWA (#6588) Co-authored-by: Carlos Mocholi <carlossmocholi@gmail.com> * Fix DPP + SyncBN (#6838) * Fix DPP + SyncBN Ensure that model is already on correct GPU before applying SyncBN conversion * Fix order of SyncBN for ddp_spawn * [Fix] TPU Training Type Plugin (#6816) * Fix support for symlink save_dir in TensorBoardLogger (#6730) * Add test for symlink support and initial fix * Respond to comment and add docstring * Update CHANGELOG.md * Simplify * Update pytorch_lightning/utilities/cloud_io.py Co-authored-by: Carlos Mocholí <carlossmocholi@gmail.com> * Make `LightningLocalFileSystem` protected Co-authored-by: Carlos Mocholí <carlossmocholi@gmail.com> * Fixed missing arguments in `lr_find` call (#6784) There seem to be 3 arguments missing in the `lr_find` call in the tunining.py file. * Update Changelog & version * Fix TPU tests for checkpoint Skip advanced profiler for torch > 1.8 Skip pytorch profiler for torch > 1.8 Fix save checkpoint logic for TPUs Co-authored-by: Jirka Borovec <Borda@users.noreply.github.com> Co-authored-by: thomas chaton <thomas@grid.ai> Co-authored-by: Your Name <you@example.com> Co-authored-by: Akihiro Nitta <nitta@akihironitta.com> Co-authored-by: Yuan-Hang Zhang <sailordiary@users.noreply.github.com> Co-authored-by: rohitgr7 <rohitgr1998@gmail.com> Co-authored-by: Carlos Mocholí <carlossmocholi@gmail.com> Co-authored-by: Elizaveta Logacheva <elimohl@gmail.com> Co-authored-by: Adrian Wälchli <aedu.waelchli@gmail.com> Co-authored-by: Karthik Prasad <prasadkr@uci.edu> Co-authored-by: Sadiq Jaffer <sadiq@toao.com> Co-authored-by: Michael Baumgartner <m.baumgartner@ymail.com> Co-authored-by: Eugene Khvedchenya <ekhvedchenya@gmail.com> Co-authored-by: Ethan Harris <ewah1g13@soton.ac.uk> Co-authored-by: Tharindu Hasthika <tharindubathigama@gmail.com>

…ter) to github/third-party/PyTorchLightning/pytorch-lightning Summary: ### New commit log messages ## [UnReleased] - 2021-MM-DD ### Added - Added more explicit exception message when trying to execute `trainer.test()` or `trainer.validate()` with `fast_dev_run=True` ([#6667](Lightning-AI/pytorch-lightning#6667)) - Added `LightningCLI` class to provide simple reproducibility with minimum boilerplate training cli. ([#4492](Lightning-AI/pytorch-lightning#4492)) - Trigger warning when non-metric logged value with multi processes hasn't been reduced ([#6417](Lightning-AI/pytorch-lightning#6417)) - Added `gradient_clip_algorithm` argument to Trainer for gradient clipping by value ([#6123](Lightning-AI/pytorch-lightning#6123)). - Added a way to print to terminal without breaking up the progress bar ([#5470](Lightning-AI/pytorch-lightning#5470)) - Added support to checkpoint after training steps in `ModelCheckpoint` callback ([#6146](Lightning-AI/pytorch-lightning#6146)) - Added `checkpoint` parameter to callback's `on_save_checkpoint` hook ([#6072](Lightning-AI/pytorch-lightning#6072)) - Added `RunningStage.SANITY_CHECKING` ([#4945](Lightning-AI/pytorch-lightning#4945)) - Added `TrainerState.{FITTING,VALIDATING,TESTING,PREDICTING,TUNING}` ([#4945](Lightning-AI/pytorch-lightning#4945)) - Added `Trainer.validate()` method to perform one evaluation epoch over the validation set ([#4948](Lightning-AI/pytorch-lightning#4948)) - Added `LightningEnvironment` for Lightning-specific DDP ([#5915](Lightning-AI/pytorch-lightning#5915)) - Added `teardown()` hook to LightningDataModule ([#4673](Lightning-AI/pytorch-lightning#4673)) - Added `auto_insert_metric_name` parameter to `ModelCheckpoint` ([#6277](Lightning-AI/pytorch-lightning#6277)) - Added arg to `self.log` that enables users to give custom names when dealing with multiple dataloaders ([#6274](Lightning-AI/pytorch-lightning#6274)) - Added `teardown` method to `BaseProfiler` to enable subclasses defining post-profiling steps outside of `__del__` ([#6370](Lightning-AI/pytorch-lightning#6370)) - Added `setup` method to `BaseProfiler` to enable subclasses defining pre-profiling steps for every process ([#6633](Lightning-AI/pytorch-lightning#6633)) - Added no return warning to predict ([#6139](Lightning-AI/pytorch-lightning#6139)) - Added `Trainer.predict` config validation ([#6543](Lightning-AI/pytorch-lightning#6543)) - Added `AbstractProfiler` interface ([#6621](Lightning-AI/pytorch-lightning#6621)) - Added support for including module names for forward in the autograd trace of `PyTorchProfiler` ([#6349](Lightning-AI/pytorch-lightning#6349)) - Added support for the PyTorch 1.8.1 autograd profiler ([#6618](Lightning-AI/pytorch-lightning#6618)) - Added `outputs` parameter to callback's `on_validation_epoch_end` & `on_test_epoch_end` hooks ([#6120](Lightning-AI/pytorch-lightning#6120)) - Added `configure_sharded_model` hook ([#6679](Lightning-AI/pytorch-lightning#6679)) - Added support for `precision=64`, enabling training with double precision ([#6595](Lightning-AI/pytorch-lightning#6595)) - Added support for DDP communication hooks ([#6736](Lightning-AI/pytorch-lightning#6736)) - Added `artifact_location` argument to `MLFlowLogger` which will be passed to the `MlflowClient.create_experiment` call ([#6677](Lightning-AI/pytorch-lightning#6677)) - Added `model` parameter to precision plugins' `clip_gradients` signature ([#6764](Lightning-AI/pytorch-lightning#6764)) ### Changed - Renamed `pytorch_lightning.callbacks.swa` to `pytorch_lightning.callbacks.stochastic_weight_avg` ([#6259](Lightning-AI/pytorch-lightning#6259)) - Refactor `RunningStage` and `TrainerState` usage ([#4945](Lightning-AI/pytorch-lightning#4945)) - Changed `trainer.evaluating` to return `True` if validating or testing ([#4945](Lightning-AI/pytorch-lightning#4945)) - Changed `setup()` and `teardown()` stage argument to take any of `{fit,validate,test,predict}` ([#6386](Lightning-AI/pytorch-lightning#6386)) - Changed profilers to save separate report files per state and rank ([#6621](Lightning-AI/pytorch-lightning#6621)) - Changed `PyTorchProfiler` to use `torch.autograd.profiler.record_function` to record functions ([#6349](Lightning-AI/pytorch-lightning#6349)) ### Deprecated - `period` has been deprecated in favor of `every_n_val_epochs` in the `ModelCheckpoint` callback ([#6146](Lightning-AI/pytorch-lightning#6146)) - Deprecated `trainer.running_sanity_check` in favor of `trainer.sanity_checking` ([#4945](Lightning-AI/pytorch-lightning#4945)) - Deprecated `Profiler(output_filename)` in favor of `dirpath` and `filename` ([#6621](Lightning-AI/pytorch-lightning#6621)) - Deprecated `PytorchProfiler(profiled_functions)` in favor of `record_functions` ([#6349](Lightning-AI/pytorch-lightning#6349)) - Deprecated metrics in favor of `torchmetrics` ([#6505](Lightning-AI/pytorch-lightning#6505), [#6530](Lightning-AI/pytorch-lightning#6530), [#6540](Lightning-AI/pytorch-lightning#6540), [#6547](Lightning-AI/pytorch-lightning#6547), [#6515](Lightning-AI/pytorch-lightning#6515), [#6572](Lightning-AI/pytorch-lightning#6572), [#6573](Lightning-AI/pytorch-lightning#6573), [#6584](Lightning-AI/pytorch-lightning#6584), [#6636](Lightning-AI/pytorch-lightning#6636), [#6637](Lightning-AI/pytorch-lightning#6637), [#6649](Lightning-AI/pytorch-lightning#6649), [#6659](Lightning-AI/pytorch-lightning#6659), ) ### Removed - Removed support for passing a bool value to `profiler` argument of Trainer ([#6164](Lightning-AI/pytorch-lightning#6164)) - Removed no return warning from val/test step ([#6139](Lightning-AI/pytorch-lightning#6139)) - Removed passing a `ModelCheckpoint` instance to `Trainer(checkpoint_callback)` ([#6166](Lightning-AI/pytorch-lightning#6166)) - Removed deprecated Trainer argument `enable_pl_optimizer` and `automatic_optimization` ([#6163](Lightning-AI/pytorch-lightning#6163)) - Removed deprecated metrics ([#6161](Lightning-AI/pytorch-lightning#6161)) * from `pytorch_lightning.metrics.functional.classification` removed `to_onehot`, `to_categorical`, `get_num_classes`, `roc`, `multiclass_roc`, `average_precision`, `precision_recall_curve`, `multiclass_precision_recall_curve` * from `pytorch_lightning.metrics.functional.reduction` removed `reduce`, `class_reduce` - Removed deprecated `ModelCheckpoint` arguments `prefix`, `mode="auto"` ([#6162](Lightning-AI/pytorch-lightning#6162)) - Removed `mode='auto'` from `EarlyStopping` ([#6167](Lightning-AI/pytorch-lightning#6167)) - Removed legacy references for magic keys in the `Result` object ([#6016](Lightning-AI/pytorch-lightning#6016)) - Removed deprecated `LightningModule` `hparams` setter ([#6207](Lightning-AI/pytorch-lightning#6207)) - Removed legacy code to log or include metrics in the progress bar by returning them in a dict with the `"log"/"progress_bar"` magic keys. Use `self.log` instead ([#6734](Lightning-AI/pytorch-lightning#6734)) - Removed `optimizer_idx` argument from `training_step` in manual optimization ([#6093](Lightning-AI/pytorch-lightning#6093)) ### Fixed - Set better defaults for `rank_zero_only.rank` when training is launched with SLURM and torchelastic ([#6802](Lightning-AI/pytorch-lightning#6802)) - Made the `Plugin.reduce` method more consistent across all Plugins to reflect a mean-reduction by default ([#6011](Lightning-AI/pytorch-lightning#6011)) - Move lightning module to correct device type when using LightningDistributedWrapper ([#6070](Lightning-AI/pytorch-lightning#6070)) - Do not print top-k verbose log with `ModelCheckpoint(monitor=None)` ([#6109](Lightning-AI/pytorch-lightning#6109)) - Fixed csv extension check ([#6436](Lightning-AI/pytorch-lightning#6436)) - Fixed `ModelCheckpoint(monitor=None, save_last=True)` not saving checkpoints ([#6136](Lightning-AI/pytorch-lightning#6136)) - Fixed `ModelCheckpoint(save_top_k=0, save_last=True)` not saving the `last` checkpoint ([#6136](Lightning-AI/pytorch-lightning#6136)) - Fixed `.teardown(stage='fit')` getting called during `trainer.test` ([#6386](Lightning-AI/pytorch-lightning#6386)) - Fixed `.on_fit_{start,end}()` getting called during `trainer.test` ([#6386](Lightning-AI/pytorch-lightning#6386)) - Fixed LightningModule `all_gather` on cpu tensors ([#6416](Lightning-AI/pytorch-lightning#6416)) - Fixed torch distributed not available in setup hook for DDP ([#6506](Lightning-AI/pytorch-lightning#6506)) - Fixed `EarlyStopping` logic when `min_epochs` or `min_steps` requirement is not met ([#6705](Lightning-AI/pytorch-lightning#6705)) ## [1.2.7] - 2021-04-06 ### Fixed - Fixed resolve a bug with omegaconf and xm.save ([#6741](Lightning-AI/pytorch-lightning#6741)) - Fixed an issue with IterableDataset when __len__ is not defined ([#6828](Lightning-AI/pytorch-lightning#6828)) - Sanitize None params during pruning ([#6836](Lightning-AI/pytorch-lightning#6836)) - Enforce an epoch scheduler interval when using SWA ([#6588](Lightning-AI/pytorch-lightning#6588)) - Fixed TPU Colab hang issue, post training ([#6816](Lightning-AI/pytorch-lightning#6816)) - Fixed a bug where `TensorBoardLogger` would give a warning and not log correctly to a symbolic link `save_dir` ([#6730](Lightning-AI/pytorch-lightning#6730)) ## [1.2.6] - 2021-03-30 ### Changed - Changed the behavior of `on_epoch_start` to run at the beginning of validation & test epoch ([#6498](Lightning-AI/pytorch-lightning#6498)) ### Removed - Removed legacy code to include `step` dictionary returns in `callback_metrics`. Use `self.log_dict` instead. ([#6682](Lightning-AI/pytorch-lightning#6682)) ### Fixed - Fixed `DummyLogger.log_hyperparams` raising a `TypeError` when running with `fast_dev_run=True` ([#6398](Lightning-AI/pytorch-lightning#6398)) - Fixed error on TPUs when there was no `ModelCheckpoint` ([#6654](Lightning-AI/pytorch-lightning#6654)) - Fixed `trainer.test` freeze on TPUs ([#6654](Lightning-AI/pytorch-lightning#6654)) - Fixed a bug where gradients were disabled after calling `Trainer.predict` ([#6657](Lightning-AI/pytorch-lightning#6657)) - Fixed bug where no TPUs were detected in a TPU pod env ([#6719](Lightning-AI/pytorch-lightning#6719)) ## [1.2.5] - 2021-03-23 ### Changed - Update Gradient Clipping for the TPU Accelerator ([#6576](Lightning-AI/pytorch-lightning#6576)) - Refactored setup for typing friendly ([#6590](Lightning-AI/pytorch-lightning#6590)) ### Fixed - Fixed a bug where `all_gather` would not work correctly with `tpu_cores=8` ([#6587](Lightning-AI/pytorch-lightning#6587)) - Fixed comparing required versions ([#6434](Lightning-AI/pytorch-lightning#6434)) - Fixed duplicate logs appearing in console when using the python logging module ([#6275](Lightning-AI/pytorch-lightning#6275)) - Added Autocast in validation, test and predict modes for Native AMP ([#6565](Lightning-AI/pytorch-lightning#6565)) Reviewed By: shuyingsunshine21 Differential Revision: D27528929 fbshipit-source-id: 311c88f71461c2c79bbf185e28d7a6d683ccc26f

Update TPU Training Type Plugin

5518f76

kaushikb11 added 11 commits April 4, 2021 21:31

Update post dispatch

90578b5

Update TPU Spawn

ad87ffd

fix reduce boolean decision

5e6eb38

Update TPU Spawn

3dba5b9

Update type hint

197ab09

Fix Acc connector for TPUs

db8ab87

Update Single device tpu

ae04ea8

fix code format issue

2f40f2e

fix

3ee4a8b

tpu spawn interactive

30ae218

fix tests

b3fbad5

kaushikb11 added the accelerator: tpu Tensor Processing Unit label Apr 5, 2021

kaushikb11 added this to the 1.2.x milestone Apr 5, 2021

kaushikb11 marked this pull request as ready for review April 5, 2021 11:44

kaushikb11 requested review from awaelchli, Borda, carmocca, justusschock, SeanNaren, tchaton and williamFalcon as code owners April 5, 2021 11:44

kaushikb11 added the bug Something isn't working label Apr 5, 2021

tchaton requested changes Apr 5, 2021

View reviewed changes

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

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

tchaton approved these changes Apr 5, 2021

View reviewed changes

update changelog

a81bcfa

mergify bot added the has conflicts label Apr 5, 2021

Merge branch 'master' into update/tpu

240b287

mergify bot removed the has conflicts label Apr 5, 2021

awaelchli approved these changes Apr 5, 2021

View reviewed changes

tests/models/test_tpu.py Outdated Show resolved Hide resolved

kaushikb11 added 2 commits April 6, 2021 01:39

Update values

3c912cb

Merge branch 'update/tpu' of https://github.com/kaushikb11/pytorch-li…

0b4a4ab

…ghtning into update/tpu

carmocca reviewed Apr 6, 2021

View reviewed changes

pytorch_lightning/plugins/training_type/single_tpu.py Outdated Show resolved Hide resolved

pytorch_lightning/plugins/training_type/single_tpu.py Outdated Show resolved Hide resolved

mergify bot added the has conflicts label Apr 6, 2021

kaushikb11 and others added 2 commits April 6, 2021 13:01

refactor

851a8d9

Merge branch 'master' into update/tpu

f944e62

mergify bot removed the has conflicts label Apr 6, 2021

carmocca approved these changes Apr 6, 2021

View reviewed changes

kaushikb11 merged commit cf8e828 into Lightning-AI:master Apr 6, 2021

kaushikb11 deleted the update/tpu branch April 6, 2021 09:32

kaushikb11 mentioned this pull request Apr 6, 2021

1.2.x cherries 🍒 #6083

Closed

kaushikb11 added a commit to kaushikb11/pytorch-lightning that referenced this pull request Apr 6, 2021

[Fix] TPU Training Type Plugin (Lightning-AI#6816)

3643954

lexierule pushed a commit that referenced this pull request Apr 7, 2021

[Fix] TPU Training Type Plugin (#6816)

c7422d4

kaushikb11 mentioned this pull request Apr 7, 2021

Fix save pretrained for TPUs minimaxir/aitextgen#105

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[Fix] TPU Training Type Plugin #6816

[Fix] TPU Training Type Plugin #6816

kaushikb11 commented Apr 4, 2021 •

edited

codecov bot commented Apr 4, 2021 •

edited

tchaton left a comment

carmocca left a comment

kaushikb11 commented Apr 6, 2021 •

edited

[Fix] TPU Training Type Plugin #6816

[Fix] TPU Training Type Plugin #6816

Conversation

kaushikb11 commented Apr 4, 2021 • edited

What does this PR do?

Before submitting

PR review

Did you have fun?

codecov bot commented Apr 4, 2021 • edited

Codecov Report

tchaton left a comment

Choose a reason for hiding this comment

carmocca left a comment

Choose a reason for hiding this comment

kaushikb11 commented Apr 6, 2021 • edited

kaushikb11 commented Apr 4, 2021 •

edited

codecov bot commented Apr 4, 2021 •

edited

kaushikb11 commented Apr 6, 2021 •

edited