Fault Tolerant Manual: Add support for DDP

pytorch_lightning/loops/epoch/evaluation_epoch_loop.py

@@ -178,10 +178,7 @@ def on_save_checkpoint(self) -> Dict:
         state_to_save = "state" if self._has_completed() else "previous_state"
         state: Optional[MergedIteratorState] = getattr(self._data_fetcher.dataloader_iter, state_to_save, None)
         if state:
-            state_dict["dataloader_state_dict"] = asdict(state)
-        state_dict["dataloader_state_dict"] = _collect_states_on_rank_zero_over_collection(
-            state_dict["dataloader_state_dict"]
-        )
+            state_dict["dataloader_state_dict"] = _collect_states_on_rank_zero_over_collection(asdict(state))
         return state_dict
 
     def on_load_checkpoint(self, state_dict: Dict) -> None:

pytorch_lightning/loops/epoch/training_epoch_loop.py

@@ -322,12 +322,8 @@ def on_save_checkpoint(self) -> Dict:
         ):
             return state_dict
 
-        state_dict["dataloader_state_dict"] = self.trainer.train_dataloader.state_dict(
-            has_completed=self._has_completed()
-        )
-
         state_dict["dataloader_state_dict"] = _collect_states_on_rank_zero_over_collection(
-            state_dict["dataloader_state_dict"]
+            self.trainer.train_dataloader.state_dict(has_completed=self._has_completed())
         )
         return state_dict
 

pytorch_lightning/trainer/trainer.py

@@ -2094,14 +2094,17 @@ def _results(self) -> Optional[ResultCollection]:
             return active_loop._results
 
     def _exit_gracefully_on_signal(self) -> None:
-        if _fault_tolerant_training():
-            # the signal should be sent to rank 0
-            should_terminate_gracefully = self.training_type_plugin.broadcast(self._terminate_gracefully)
-            if not should_terminate_gracefully:
-                return
-            caller = inspect.stack()[1]
-            class_name = caller[0].f_locals["self"].__class__.__name__
-            raise ExitGracefullyException(f"Exiting gracefully on {class_name}:{caller.function}")
+        if not _fault_tolerant_training():
+            return
+        if not self._should_terminated_gracefully():
+            return
+        caller = inspect.stack()[1]
+        class_name = caller[0].f_locals["self"].__class__.__name__
+        raise ExitGracefullyException(f"Exiting gracefully on {class_name}:{caller.function}")
+
+    def _should_terminated_gracefully(self) -> bool:
+        value = torch.tensor(self._terminate_gracefully, device=self.training_type_plugin.root_device)
+        return self.training_type_plugin.reduce(value, reduce_op="sum") > 0
 
     @property
     def weights_summary(self) -> Optional[str]: