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

Code before:

    def configure_ddp(self):
        if self.main_rpc_process:
            self.pre_configure_ddp()

            self._model = DistributedDataParallel(
                LightningDistributedModule(self.model),
                device_ids=self.determine_ddp_device_ids(),
                process_group=mpu.get_data_parallel_group(),
                **self._ddp_kwargs,
            )
            # Plugin handle backwards across processes. Currently not supported for DDP + pipe parallel
            self._model.require_backward_grad_sync = False

    @rank_zero_only
    def rpc_save_model(self, save_model_fn, last_filepath, trainer, pl_module) -> None:
        model = self.lightning_module
        if not hasattr(model.sequential_module, "foreach_worker"):
            return
        current_layers = pl_module.sequential_module
        model.sequential_module.foreach_worker(
            save_layers_on_all_rank_zero_workers, {"gpus_per_model": self.gpus_per_model}, include_self=True
        )
        pl_module.sequential_module = load_sequential_from_saved_layers(self.gpus_per_model)
        save_model_fn(last_filepath, trainer, pl_module)
        pl_module.sequential_module = current_layers

    def worker_optimizer_step(self, model: LightningModule, opt_idx: int, *args, **kwargs) -> None:
        model.sequential_module.foreach_worker(
            run_optimizer, {
                "opt_idx": opt_idx,
                "args": args,
                "kwargs": kwargs
            }, include_self=False
        )

    @property
    def distributed_sampler_kwargs(self):
        return dict(
            num_replicas=mpu.get_data_parallel_world_size(),
            rank=mpu.get_data_parallel_rank(),
        )

    @property
    def data_parallel_group(self):
        return mpu.get_data_parallel_group()

    def set_main_rpc_process(self):
        self.main_rpc_process = torch_distrib.get_rank(group=mpu.get_pipeline_parallel_group()) == 0

    @property

Review comment (suggestion):
not sure how much this makes sense since the `model` is not used...
cc: @tchaton 

Code after:

    def configure_ddp(self):
        if self.main_rpc_process:
            self.pre_configure_ddp()

            self._model = DistributedDataParallel(
                LightningDistributedModule(self.model),
                device_ids=self.determine_ddp_device_ids(),
                process_group=mpu.get_data_parallel_group(),
                **self._ddp_kwargs,
            )
            # Plugin handle backwards across processes. Currently not supported for DDP + pipe parallel
            self._model.require_backward_grad_sync = False

    @rank_zero_only
    def rpc_save_model(self, save_model_fn, last_filepath, trainer) -> None:
        model = self.lightning_module
        if not hasattr(model.sequential_module, "foreach_worker"):
            return
        current_layers = model.sequential_module
        model.sequential_module.foreach_worker(
            save_layers_on_all_rank_zero_workers, {"gpus_per_model": self.gpus_per_model}, include_self=True
        )
        model.sequential_module = load_sequential_from_saved_layers(self.gpus_per_model)
        save_model_fn(last_filepath, trainer)
        model.sequential_module = current_layers

    def worker_optimizer_step(self, model: LightningModule, opt_idx: int, *args, **kwargs) -> None:
        model.sequential_module.foreach_worker(
            run_optimizer, {
                "opt_idx": opt_idx,
                "args": args,
                "kwargs": kwargs
            }, include_self=False
        )

    @property
    def distributed_sampler_kwargs(self):
        return dict(
            num_replicas=mpu.get_data_parallel_world_size(),
            rank=mpu.get_data_parallel_rank(),
        )

    @property
    def data_parallel_group(self):
        return mpu.get_data_parallel_group()

    def set_main_rpc_process(self):
        self.main_rpc_process = torch_distrib.get_rank(group=mpu.get_pipeline_parallel_group()) == 0

    @property