[fix] Add barriers before and after setup hook is run

CHANGELOG.md

@@ -9,6 +9,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Added
 
+- Added synchronization points before and after `setup` hooks are run ([#7202](https://github.com/PyTorchLightning/pytorch-lightning/pull/7202))
+
+
 - Added a `teardown` hook to `ClusterEnvironment` ([#6942](https://github.com/PyTorchLightning/pytorch-lightning/pull/6942))
 
 

pytorch_lightning/plugins/training_type/ddp.py

@@ -285,7 +285,7 @@ def post_dispatch(self) -> None:
  self.cluster_environment.teardown()
 
  def barrier(self, *args, **kwargs):
- if torch_distrib.is_initialized():
+ if torch_distrib.is_available() and torch_distrib.is_initialized():
  torch_distrib.barrier()
 
  def broadcast(self, obj: object, src: int = 0) -> object:

pytorch_lightning/trainer/trainer.py

@@ -1111,6 +1111,7 @@ def tune(
  def call_setup_hook(self, model: LightningModule) -> None:
  assert self.state.running, f"TrainerState: {self.state}"
  state = self._setup_state
+ self.accelerator.barrier("pre_setup")
 
  if self.datamodule is not None:
  called = getattr(self.datamodule, f'has_setup_{state}')
@@ -1119,6 +1120,7 @@ def call_setup_hook(self, model: LightningModule) -> None:
 
  self.setup(model, stage=state)
  model.setup(stage=state)
+ self.accelerator.barrier("post_setup")
 
  def call_configure_sharded_model(self, model: LightningModule) -> None:
  # Call configure sharded model hook if accelerator requests. In some cases