Repository: Lightning-AI/pytorch-lightning
File: src/lightning/fabric/strategies/fsdp.py (Python)

Code before:
    def precision(self, precision: Optional[FSDPPrecision]) -> None:
        if precision is not None and not isinstance(precision, FSDPPrecision):
            raise TypeError(f"The FSDP strategy can only work with the `FSDPPrecision` plugin, found {precision}")
        self._precision = precision

    @override
    def _configure_launcher(self) -> None:
        assert self.cluster_environment is not None
        if not self.cluster_environment.creates_processes_externally:
            self._launcher = _SubprocessScriptLauncher(self.cluster_environment, self.num_processes, self.num_nodes)

    @override
    def setup_environment(self) -> None:
        super().setup_environment()
        self._setup_distributed()

        # if 'device_mesh' in the `_fsdp_kwargs` is provided as a tuple, update it into the `DeviceMesh` object here
        if isinstance(self._fsdp_kwargs.get("device_mesh"), tuple):
            from torch.distributed.device_mesh import init_device_mesh

            self._fsdp_kwargs["device_mesh"] = init_device_mesh("cuda", self._fsdp_kwargs["device_mesh"])

    @override
    def setup_module_and_optimizers(
        self, module: Module, optimizers: list[Optimizer], scheduler: Optional[_LRScheduler] = None
    ) -> Tuple[Module, list[Optimizer], Optional[_LRScheduler]]:
        """Wraps the model into a :class:`~torch.distributed.fsdp.fully_sharded_data_parallel.FullyShardedDataParallel`
        module and sets `use_orig_params=True` to keep the reference to the original parameters in the optimizer."""
        use_orig_params = self._fsdp_kwargs.get("use_orig_params")
        if use_orig_params is False:
            raise ValueError(
                f"You set `{type(self).__name__}(use_orig_params=False)` but this is not supported when"
                " setting the model and optimizer up jointly. Either set it to `True` or set the objects"
                " up in this order: Create the model, call `setup_module`, create the optimizer,"
                " call `setup_optimizer`."
            )
        module = self.setup_module(module)
        return module, optimizers, scheduler

    @override
    def setup_module(self, module: Module) -> Module:
        """Wraps the model into a :class:`~torch.distributed.fsdp.fully_sharded_data_parallel.FullyShardedDataParallel`
        module."""
        from torch.distributed.fsdp import FullyShardedDataParallel

        if any(isinstance(mod, FullyShardedDataParallel) for mod in module.modules()):
            # The user has wrapped their submodules manually, don't apply the auto wrap policy.
            if _has_meta_device_parameters_or_buffers(module):
                rank_zero_warn(
                    "The model is already wrapped in `FSDP` but there are still parameters on the meta device."
                )

Review comment (performance):
```suggestion
    ) -> tuple[Module, list[Optimizer], Optional[_LRScheduler]]:
```

Code after:
    def precision(self, precision: Optional[Precision]) -> None:
        if precision is not None and not isinstance(precision, FSDPPrecision):
            raise TypeError(f"The FSDP strategy can only work with the `FSDPPrecision` plugin, found {precision}")
        self._precision = precision

    @override
    def _configure_launcher(self) -> None:
        assert self.cluster_environment is not None
        if not self.cluster_environment.creates_processes_externally:
            self._launcher = _SubprocessScriptLauncher(self.cluster_environment, self.num_processes, self.num_nodes)

    @override
    def setup_environment(self) -> None:
        super().setup_environment()
        self._setup_distributed()

        # if 'device_mesh' in the `_fsdp_kwargs` is provided as a tuple, update it into the `DeviceMesh` object here
        if isinstance(self._fsdp_kwargs.get("device_mesh"), tuple):
            from torch.distributed.device_mesh import init_device_mesh

            self._fsdp_kwargs["device_mesh"] = init_device_mesh("cuda", self._fsdp_kwargs["device_mesh"])

    @override
    def setup_module_and_optimizers(
        self, module: Module, optimizers: list[Optimizer], scheduler: Optional["_LRScheduler"] = None
    ) -> tuple[Module, list[Optimizer], Optional["_LRScheduler"]]:
        """Wraps the model into a :class:`~torch.distributed.fsdp.fully_sharded_data_parallel.FullyShardedDataParallel`
        module and sets `use_orig_params=True` to keep the reference to the original parameters in the optimizer."""
        use_orig_params = self._fsdp_kwargs.get("use_orig_params")
        if use_orig_params is False:
            raise ValueError(
                f"You set `{type(self).__name__}(use_orig_params=False)` but this is not supported when"
                " setting the model and optimizer up jointly. Either set it to `True` or set the objects"
                " up in this order: Create the model, call `setup_module`, create the optimizer,"
                " call `setup_optimizer`."
            )
        module = self.setup_module(module)
        return module, optimizers, scheduler

    @override
    def setup_module(self, module: Module) -> Module:
        """Wraps the model into a :class:`~torch.distributed.fsdp.fully_sharded_data_parallel.FullyShardedDataParallel`
        module."""
        from torch.distributed.fsdp import FullyShardedDataParallel

        if any(isinstance(mod, FullyShardedDataParallel) for mod in module.modules()):
            # The user has wrapped their submodules manually, don't apply the auto wrap policy.
            if _has_meta_device_parameters_or_buffers(module):
                rank_zero_warn(
                    "The model is already wrapped in `FSDP` but there are still parameters on the meta device."
                )