Repository: Lightning-AI/pytorch-lightning
File: pytorch_lightning/plugins/training_type/ddp_spawn.py (Python)

Code before:
            register_ddp_comm_hook(
                model=self._model,
                ddp_comm_state=self._ddp_comm_state,
                ddp_comm_hook=self._ddp_comm_hook,
                ddp_comm_wrapper=self._ddp_comm_wrapper,
            )

    def configure_ddp(self) -> None:
        self.pre_configure_ddp()
        self._model = self._setup_model(LightningDistributedModule(self.model))
        self._register_ddp_hooks()

    def determine_ddp_device_ids(self):
        if self.root_device.type == "cpu":
            return None
        return [self.root_device.index]

    def __transfer_distrib_spawn_state_on_fit_end(self, trainer: "pl.Trainer", results: Any) -> None:
        rank_zero_info("Transferring results back to main process ...")
        checkpoint_callback = trainer.checkpoint_callback
        best_model_path = checkpoint_callback.best_model_path if checkpoint_callback else None

        # requires to compute the state_dict on all processes in case Metrics are present
        state_dict = self.lightning_module.state_dict()

        if self.global_rank == 0:
            # save the last weights
            last_path = None
            if trainer.state.fn == TrainerFn.FITTING and best_model_path is not None and len(best_model_path) > 0:
                last_path = re.sub(".ckpt", ".tmp_end.ckpt", best_model_path)
                self.checkpoint_io.save_checkpoint(state_dict, last_path)

            # todo, pass complete checkpoint as state dictionary
            self.mp_queue.put(best_model_path)
            self.mp_queue.put(last_path)
            self.mp_queue.put(results)
            # adds the `callback_metrics` to the queue
            # TODO: Remove the if in v1.7
            if is_overridden("add_to_queue", self.lightning_module):
                self.lightning_module.add_to_queue(self.mp_queue)
            else:
                self.add_to_queue(trainer, self.mp_queue)

    def __recover_child_process_weights(self, best_path, last_path):
        # transfer back the best path to the trainer
        if self.lightning_module.trainer.checkpoint_callback:
            self.lightning_module.trainer.checkpoint_callback.best_model_path = best_path
        # todo, pass also best score

        # load last weights
        if last_path is not None and self.lightning_module.trainer.state.fn == TrainerFn.FITTING:

Review comment (suggestion):
```suggestion
        if not self.global_rank == 0:
            return
```

Code after:
                ddp_comm_hook=self._ddp_comm_hook,
                ddp_comm_wrapper=self._ddp_comm_wrapper,
            )

    def configure_ddp(self) -> None:
        self.pre_configure_ddp()
        self._model = self._setup_model(LightningDistributedModule(self.model))
        self._register_ddp_hooks()

    def determine_ddp_device_ids(self):
        if self.root_device.type == "cpu":
            return None
        return [self.root_device.index]

    def __transfer_distrib_spawn_state_on_fit_end(self, trainer: "pl.Trainer", results: Any) -> None:
        rank_zero_warn("cleaning up ddp environment...")
        checkpoint_callback = trainer.checkpoint_callback
        best_model_path = checkpoint_callback.best_model_path if checkpoint_callback else None

        # requires to compute the state_dict on all processes in case Metrics are present
        state_dict = self.lightning_module.state_dict()

        if self.global_rank != 0:
            return

        # save the last weights
        last_path = None
        if trainer.state.fn == TrainerFn.FITTING and best_model_path is not None and len(best_model_path) > 0:
            last_path = re.sub(".ckpt", ".tmp_end.ckpt", best_model_path)
            self.checkpoint_io.save_checkpoint(state_dict, last_path)

        # todo, pass complete checkpoint as state dictionary
        self.mp_queue.put(best_model_path)
        self.mp_queue.put(last_path)
        self.mp_queue.put(results)
        # adds the `callback_metrics` to the queue
        # TODO: Remove the if in v1.7
        if is_overridden("add_to_queue", self.lightning_module):
            self.lightning_module.add_to_queue(self.mp_queue)
        else:
            self.add_to_queue(trainer, self.mp_queue)

    def __recover_child_process_weights(self, best_path, last_path):
        # transfer back the best path to the trainer
        if self.lightning_module.trainer.checkpoint_callback:
            self.lightning_module.trainer.checkpoint_callback.best_model_path = best_path
        # todo, pass also best score

        # load last weights
        if last_path is not None and self.lightning_module.trainer.state.fn == TrainerFn.FITTING:
            ckpt = self.checkpoint_io.load_checkpoint(last_path, map_location=(lambda storage, loc: storage))