[HotFix] Resolve TPU Training

.gitignore

@@ -155,3 +155,5 @@ cifar-10-batches-py
 # ctags
 tags
 data
+MNIST
+runs

CHANGELOG.md

@@ -288,6 +288,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Fixed missing `process_dataloader` call for `TPUSpawn` when in distributed mode ([#6015](https://github.com/PyTorchLightning/pytorch-lightning/pull/6015))
 
 
+- Fixed synchrnization issues with TPUs Training ([#6027](https://github.com/PyTorchLightning/pytorch-lightning/pull/6027))
+
+
 ## [1.1.8] - 2021-02-08
 
 ### Fixed

dockers/tpu-tests/tpu_test_cases.jsonnet

@@ -25,6 +25,12 @@ local tputests = base.BaseTest {
           pytorch_lightning/utilities/xla_device_utils.py \
           tests/accelerators/test_tpu_backend.py \
           tests/models/test_tpu.py
+      # Takes too long
+      # python pl_examples/domain_templates/computer_vision_fine_tuning.py \
+      #  --tpu_cores 8 \
+      #  --epochs 15 \
+      #  --limit_train_batches 8 \
+      #  --limit_val_batches 8
       test_exit_code=$?
       echo "\n||| END PYTEST LOGS |||\n"
       coverage xml

pl_examples/domain_templates/computer_vision_fine_tuning.py

@@ -55,7 +55,7 @@
 import pytorch_lightning as pl
 from pl_examples import cli_lightning_logo
 from pytorch_lightning import _logger as log
-from pytorch_lightning import LightningDataModule
+from pytorch_lightning import LightningDataModule, seed_everything
 from pytorch_lightning.callbacks.finetuning import BaseFinetuning
 from pytorch_lightning.utilities import rank_zero_info
 
@@ -71,7 +71,7 @@ def __init__(self, milestones: tuple = (5, 10), train_bn: bool = False):
         self.train_bn = train_bn
 
     def freeze_before_training(self, pl_module: pl.LightningModule):
-        self.freeze(modules=pl_module.feature_extractor, train_bn=self.train_bn)
+        self.freeze(modules=pl_module.feature_extractor, train_bn=False)
 
     def finetune_function(self, pl_module: pl.LightningModule, epoch: int, optimizer: Optimizer, opt_idx: int):
         if epoch == self.milestones[0]:
@@ -148,7 +148,7 @@ def val_dataloader(self):
     def add_model_specific_args(parent_parser):
         parser = argparse.ArgumentParser(parents=[parent_parser])
         parser.add_argument(
-            "--num-workers", default=0, type=int, metavar="W", help="number of CPU workers", dest="num_workers"
+            "--num-workers", default=2, type=int, metavar="W", help="number of CPU workers", dest="num_workers"
         )
         parser.add_argument(
             "--batch-size", default=8, type=int, metavar="W", help="number of sample in a batch", dest="batch_size"
@@ -276,10 +276,23 @@ def add_model_specific_args(parent_parser):
             help="Name (as in ``torchvision.models``) of the feature extractor",
         )
         parser.add_argument(
-            "--epochs", default=15, type=int, metavar="N", help="total number of epochs", dest="nb_epochs"
+            "--epochs", default=5, type=int, metavar="N", help="total number of epochs", dest="nb_epochs"
+        )
+        parser.add_argument(
+            "--limit_train_batches",
+            default=1.0,
+            type=float,
+            help="How much of training dataset to check (floats = percent, int = num_batches)"
+        )
+        parser.add_argument(
+            "--limit_val_batches",
+            default=1.0,
+            type=float,
+            help="How much of validation dataset to check (floats = percent, int = num_batches)"
         )
         parser.add_argument("--batch-size", default=8, type=int, metavar="B", help="batch size", dest="batch_size")
         parser.add_argument("--gpus", type=int, default=0, help="number of gpus to use")
+        parser.add_argument("--tpu_cores", type=int, default=None, help="number of tpu cores to use")
         parser.add_argument(
             "--lr", "--learning-rate", default=1e-3, type=float, metavar="LR", help="initial learning rate", dest="lr"
         )
@@ -300,7 +313,7 @@ def add_model_specific_args(parent_parser):
             dest="train_bn",
         )
         parser.add_argument(
-            "--milestones", default=[2, 4], type=list, metavar="M", help="List of two epochs milestones"
+            "--milestones", default=[5, 10], type=list, metavar="M", help="List of two epochs milestones"
         )
         return parser
 
@@ -315,6 +328,7 @@ def main(args: argparse.Namespace) -> None:
         For the sake of the example, the images dataset will be downloaded
         to a temporary directory.
     """
+    seed_everything(42)
 
     datamodule = CatDogImageDataModule(
         dl_path=os.path.join(args.root_data_path, 'data'), batch_size=args.batch_size, num_workers=args.num_workers
@@ -326,6 +340,7 @@ def main(args: argparse.Namespace) -> None:
         weights_summary=None,
         progress_bar_refresh_rate=1,
         num_sanity_val_steps=0,
+        tpu_cores=args.tpu_cores,
         gpus=args.gpus,
         max_epochs=args.nb_epochs,
         callbacks=[finetuning_callback]

pytorch_lightning/accelerators/accelerator.py

@@ -11,10 +11,11 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from typing import Any, Callable, Optional, TYPE_CHECKING, Union
+from typing import Any, Callable, Iterable, Optional, TYPE_CHECKING, Union
 
 import torch
 from torch.optim import Optimizer
+from torch.utils.data import DataLoader
 
 from pytorch_lightning.core import LightningModule
 from pytorch_lightning.plugins.precision import (
@@ -388,3 +389,11 @@ def all_gather(self, tensor: Union[torch.Tensor], group: Optional[Any] = None, s
             A tensor of shape (world_size, batch, ...)
         """
         return all_gather_ddp_if_available(tensor, group=group, sync_grads=sync_grads)
+
+    def process_dataloader(self, dataloader: Union[Iterable, DataLoader]) -> Union[Iterable, DataLoader]:
+        """Wraps the dataloader if necessary
+
+        Args:
+            dataloader: iterable. Ideally of type: :class:`torch.utils.data.DataLoader`
+        """
+        return self.training_type_plugin.process_dataloader(dataloader)

pytorch_lightning/callbacks/model_checkpoint.py

@@ -554,6 +554,12 @@ def _save_top_k_checkpoints(self, trainer, pl_module, metrics):
         epoch = metrics.get("epoch")
         step = metrics.get("step")
 
+        # when `val_loss` is being logged and no ModelCheckpoint is being provided
+        # `val_loss` will be selected for monitor and need to be reduced to
+        # prevent processes divergence
+        if self.monitor == "val_loss":
+            current = trainer.training_type_plugin.reduce(current, reduce_op="mean")
+
         if self.check_monitor_top_k(current):
             self._update_best_and_save(current, epoch, step, trainer, pl_module, metrics)
         elif self.verbose:

pytorch_lightning/loggers/tensorboard.py

@@ -234,6 +234,9 @@ def save(self) -> None:
 
     @rank_zero_only
     def finalize(self, status: str) -> None:
+        self.close()
+
+    def close(self):
         self.experiment.flush()
         self.experiment.close()
         self.save()

pytorch_lightning/plugins/training_type/tpu_spawn.py

@@ -46,10 +46,6 @@ def create_mp_queue(self):
     def distributed_sampler_kwargs(self) -> dict:
         return dict(num_replicas=xm.xrt_world_size(), rank=xm.get_ordinal())
 
-    @property
-    def should_finalize(self):
-        return self.world_size == 1
-
     @property
     def is_distributed(self):
         return self.world_size != 1
@@ -179,6 +175,14 @@ def reduce_early_stopping_decision(self, should_stop: bool) -> bool:
         should_stop = int(stop.item()) == self.world_size
         return should_stop
 
+    def reduce(self, output, group: Optional[Any] = None, reduce_op: str = None):
+        if not isinstance(output, torch.Tensor):
+            output = torch.tensor(output, device=self.device)
+        output = xm.mesh_reduce('reduce', output, sum)
+        if isinstance(reduce_op, str) and reduce_op.lower() == "mean":
+            output /= self.world_size
+        return output
+
     def post_dispatch(self) -> None:
         # TODO: Check if trainer references can be resolved otherwise
         model = self.lightning_module
@@ -213,6 +217,10 @@ def __load_weights_on_main_process(self) -> None:
 
         self._model = model
 
+    def _close_logger(self, trainer) -> None:
+        if hasattr(trainer, "logger"):
+            trainer.logger.close()
+
     @property
     def xmp_spawn_kwargs(self):
         return {
@@ -225,9 +233,11 @@ def start_training(self, trainer) -> None:
         # todo: precision pluging is call in accelerator setup and should be moved
         if 'XLA_USE_BF16' in os.environ:
             del os.environ["XLA_USE_BF16"]
+        self._close_logger(trainer)
         xmp.spawn(self.new_process, **self.xmp_spawn_kwargs)
 
     def start_testing(self, trainer) -> None:
+        self._close_logger(trainer)
         xmp.spawn(self.new_process, **self.xmp_spawn_kwargs)
 
     def start_predicting(self, trainer) -> None:

pytorch_lightning/plugins/training_type/training_type_plugin.py

@@ -35,10 +35,6 @@ def __init__(self) -> None:
         self._results = None
         self.global_rank = 0
 
-    @property
-    def should_finalize(self):
-        return True
-
     @property
     @abstractmethod
     def on_gpu(self) -> bool:

pytorch_lightning/trainer/trainer.py

@@ -711,7 +711,7 @@ def run_evaluation(self, max_batches=None, on_epoch=False):
         for dataloader_idx, dataloader in enumerate(dataloaders):
             # bookkeeping
             dl_outputs = []
-            dataloader = self.training_type_plugin.process_dataloader(dataloader)
+            dataloader = self.accelerator.process_dataloader(dataloader)
             dl_max_batches = self.evaluation_loop.max_batches[dataloader_idx]
 
             for batch_idx, batch in enumerate(dataloader):
@@ -823,7 +823,7 @@ def run_predict(self):
 
         # run validation/testing
         for dataloader_idx, dataloader in enumerate(dataloaders):
-            dataloader = self.training_type_plugin.process_dataloader(dataloader)
+            dataloader = self.accelerator.process_dataloader(dataloader)
             dl_max_batches = self.predict_loop.max_batches[dataloader_idx]
 
             for batch_idx, batch in enumerate(dataloader):

pytorch_lightning/trainer/training_loop.py

@@ -140,7 +140,7 @@ def on_train_end(self):
         # todo: TPU 8 cores hangs in flush with TensorBoard. Might do for all loggers.
         # It might be related to xla tensors blocked when moving the cpu
         # kill loggers
-        if self.trainer.logger is not None and self.trainer.training_type_plugin.should_finalize:
+        if self.trainer.logger is not None:
             self.trainer.logger.finalize("success")
 
         # summarize profile results
@@ -502,7 +502,7 @@ def tbptt_split_batch(self, batch):
 
     def run_training_epoch(self):
         # modify dataloader if needed (ddp, etc...)
-        train_dataloader = self.trainer.training_type_plugin.process_dataloader(self.trainer.train_dataloader)
+        train_dataloader = self.trainer.accelerator.process_dataloader(self.trainer.train_dataloader)
 
         # track epoch output
         epoch_output = [[] for _ in range(self.num_optimizers)]

tests/helpers/utils.py

@@ -91,7 +91,7 @@ def wrapper(*args, **kwargs):
 
         def inner_f(queue, **kwargs):
             try:
-                func(**kwargs)
+                func(*args, **kwargs)
                 queue.put(1)
             except Exception:
                 _trace = traceback.format_exc()

tests/models/test_tpu.py

@@ -264,9 +264,6 @@ def test_distributed_backend_set_when_using_tpu(tmpdir, tpu_cores):
 
 
 @pytest.mark.skipif(not _TPU_AVAILABLE, reason="test requires TPU machine")
-@pytest.mark.skipif(
-    not os.getenv("PL_RUNNING_SPECIAL_TESTS", '0') == '1', reason="test should be run outside of pytest"
-)
 @pl_multi_process_test
 def test_broadcast_on_tpu():
     """ Checks if an object from the master process is broadcasted to other processes correctly"""