Fairscale import updates

src/pytorch_lightning/overrides/fairscale.py

@@ -14,43 +14,35 @@
 from typing import Optional, Union
 
 import torch.nn as nn
-from lightning_utilities.core.imports import module_available
 
 import pytorch_lightning as pl
 from pytorch_lightning.overrides.base import (
     _LightningModuleWrapperBase,
     _LightningPrecisionModuleWrapperBase,
     unwrap_lightning_module,
 )
-from pytorch_lightning.utilities.imports import _IS_WINDOWS
 from pytorch_lightning.utilities.rank_zero import rank_zero_deprecation
 
-_FAIRSCALE_AVAILABLE = not _IS_WINDOWS and module_available("fairscale.nn")
 
+class LightningShardedDataParallel(_LightningModuleWrapperBase):
+    def __init__(
+        self,
+        forward_module: Optional[Union["pl.LightningModule", _LightningPrecisionModuleWrapperBase]] = None,
+        pl_module: Optional[Union["pl.LightningModule", _LightningPrecisionModuleWrapperBase]] = None,
+    ) -> None:
+        self._validate_init_arguments(pl_module, forward_module)
+        super().__init__(forward_module=(pl_module or forward_module))
 
-if _FAIRSCALE_AVAILABLE:  # pragma: no-cover
+
+def unwrap_lightning_module_sharded(wrapped_model: nn.Module) -> "pl.LightningModule":
     from fairscale.nn.data_parallel.sharded_ddp import ShardedDataParallel
 
-    class LightningShardedDataParallel(_LightningModuleWrapperBase):
-        def __init__(
-            self,
-            forward_module: Optional[Union["pl.LightningModule", _LightningPrecisionModuleWrapperBase]] = None,
-            pl_module: Optional[Union["pl.LightningModule", _LightningPrecisionModuleWrapperBase]] = None,
-        ) -> None:
-            self._validate_init_arguments(pl_module, forward_module)
-            super().__init__(forward_module=(pl_module or forward_module))
-
-    def unwrap_lightning_module_sharded(wrapped_model: nn.Module) -> "pl.LightningModule":
-        rank_zero_deprecation(
-            "The function `unwrap_lightning_module_sharded` is deprecated in v1.8.0 and will be removed in v1.10.0."
-            " Access the `LightningModule` directly through the strategy attribute `Strategy.lightning_module`."
-        )
-        model = wrapped_model
-        if isinstance(model, ShardedDataParallel):
-            model = model.module
-
-        return unwrap_lightning_module(model, _suppress_warning=True)
-
-else:
-    LightningShardedDataParallel = ...  # type: ignore[assignment,misc]
-    unwrap_lightning_module_sharded = ...  # type: ignore[assignment]
+    rank_zero_deprecation(
+        "The function `unwrap_lightning_module_sharded` is deprecated in v1.8.0 and will be removed in v1.10.0."
+        " Access the `LightningModule` directly through the strategy attribute `Strategy.lightning_module`."
+    )
+    model = wrapped_model
+    if isinstance(model, ShardedDataParallel):
+        model = model.module
+
+    return unwrap_lightning_module(model, _suppress_warning=True)

src/pytorch_lightning/plugins/precision/sharded_native_amp.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 from typing import Optional, Union
 
-from pytorch_lightning.overrides.fairscale import _FAIRSCALE_AVAILABLE
+from lightning_lite.strategies.fairscale import _FAIRSCALE_AVAILABLE
 from pytorch_lightning.plugins.precision.native_amp import NativeMixedPrecisionPlugin
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
 

src/pytorch_lightning/strategies/ddp.py

@@ -30,6 +30,7 @@
 
 import pytorch_lightning as pl
 from lightning_lite.plugins import CheckpointIO, ClusterEnvironment
+from lightning_lite.strategies.fairscale import _FAIRSCALE_AVAILABLE
 from lightning_lite.utilities.distributed import distributed_available, get_default_process_group_backend_for_device
 from lightning_lite.utilities.distributed import group as _group
 from lightning_lite.utilities.distributed import init_dist_connection, ReduceOp, sync_ddp_if_available
@@ -39,7 +40,6 @@
 from pytorch_lightning.overrides import LightningDistributedModule
 from pytorch_lightning.overrides.base import _LightningPrecisionModuleWrapperBase
 from pytorch_lightning.overrides.distributed import prepare_for_backward
-from pytorch_lightning.overrides.fairscale import _FAIRSCALE_AVAILABLE
 from pytorch_lightning.plugins.precision import PrecisionPlugin
 from pytorch_lightning.strategies.launchers.subprocess_script import _SubprocessScriptLauncher
 from pytorch_lightning.strategies.parallel import ParallelStrategy

src/pytorch_lightning/strategies/fully_sharded.py

@@ -19,10 +19,10 @@
 
 import pytorch_lightning as pl
 from lightning_lite.plugins import CheckpointIO, ClusterEnvironment
+from lightning_lite.strategies.fairscale import _FAIRSCALE_AVAILABLE
 from lightning_lite.utilities.enums import PrecisionType
 from lightning_lite.utilities.optimizer import optimizers_to_device
 from pytorch_lightning.overrides.base import _LightningModuleWrapperBase
-from pytorch_lightning.overrides.fairscale import _FAIRSCALE_AVAILABLE
 from pytorch_lightning.plugins.precision import PrecisionPlugin
 from pytorch_lightning.strategies.ddp import DDPStrategy
 from pytorch_lightning.trainer.states import TrainerFn

src/pytorch_lightning/strategies/sharded.py

@@ -19,11 +19,11 @@
 from torch.optim import Optimizer
 
 import pytorch_lightning as pl
+from lightning_lite.strategies.fairscale import _FAIRSCALE_AVAILABLE
 from lightning_lite.utilities.enums import PrecisionType
 from lightning_lite.utilities.optimizer import optimizers_to_device
 from pytorch_lightning.core.optimizer import LightningOptimizer
 from pytorch_lightning.overrides.base import _LightningModuleWrapperBase, _LightningPrecisionModuleWrapperBase
-from pytorch_lightning.overrides.fairscale import _FAIRSCALE_AVAILABLE
 from pytorch_lightning.strategies.ddp import DDPStrategy
 from pytorch_lightning.trainer.states import TrainerFn
 from pytorch_lightning.utilities.exceptions import MisconfigurationException

src/pytorch_lightning/strategies/sharded_spawn.py

@@ -19,9 +19,9 @@
 from torch.optim import Optimizer
 
 import pytorch_lightning as pl
+from lightning_lite.strategies.fairscale import _FAIRSCALE_AVAILABLE
 from lightning_lite.utilities.optimizer import optimizers_to_device
 from pytorch_lightning.overrides.base import _LightningModuleWrapperBase, _LightningPrecisionModuleWrapperBase
-from pytorch_lightning.overrides.fairscale import _FAIRSCALE_AVAILABLE
 from pytorch_lightning.strategies.ddp_spawn import DDPSpawnStrategy
 from pytorch_lightning.trainer.states import TrainerFn
 from pytorch_lightning.utilities.exceptions import MisconfigurationException

tests/tests_pytorch/helpers/runif.py

@@ -20,9 +20,9 @@
 from packaging.version import Version
 from pkg_resources import get_distribution
 
+from lightning_lite.strategies.fairscale import _FAIRSCALE_AVAILABLE
 from pytorch_lightning.accelerators.mps import MPSAccelerator
 from pytorch_lightning.callbacks.progress.rich_progress import _RICH_AVAILABLE
-from pytorch_lightning.overrides.fairscale import _FAIRSCALE_AVAILABLE
 from pytorch_lightning.strategies.bagua import _BAGUA_AVAILABLE
 from pytorch_lightning.strategies.deepspeed import _DEEPSPEED_AVAILABLE
 from pytorch_lightning.utilities.imports import (

tests/tests_pytorch/plugins/precision/test_sharded_precision.py

@@ -15,7 +15,7 @@
 import pytest
 import torch
 
-from pytorch_lightning.overrides.fairscale import _FAIRSCALE_AVAILABLE
+from lightning_lite.strategies.fairscale import _FAIRSCALE_AVAILABLE
 from pytorch_lightning.plugins import ShardedNativeMixedPrecisionPlugin
 from tests_pytorch.helpers.runif import RunIf
 

tests/tests_pytorch/strategies/test_ddp_fully_sharded_with_full_state_dict.py

@@ -5,10 +5,10 @@
 import pytest
 import torch
 
+from lightning_lite.strategies.fairscale import _FAIRSCALE_AVAILABLE
 from pytorch_lightning import Trainer
 from pytorch_lightning.callbacks import ModelCheckpoint
 from pytorch_lightning.demos.boring_classes import BoringModel
-from pytorch_lightning.overrides.fairscale import _FAIRSCALE_AVAILABLE
 from pytorch_lightning.plugins import FullyShardedNativeMixedPrecisionPlugin
 from pytorch_lightning.strategies import DDPFullyShardedStrategy
 from pytorch_lightning.utilities.exceptions import MisconfigurationException

tests/tests_pytorch/strategies/test_ddp_strategy.py

@@ -20,9 +20,9 @@
 from torch.nn.parallel import DistributedDataParallel
 
 from lightning_lite.plugins.environments import ClusterEnvironment, LightningEnvironment
+from lightning_lite.strategies.fairscale import _FAIRSCALE_AVAILABLE
 from pytorch_lightning import LightningModule, Trainer
 from pytorch_lightning.demos.boring_classes import BoringModel
-from pytorch_lightning.overrides.fairscale import _FAIRSCALE_AVAILABLE
 from pytorch_lightning.strategies import DDPStrategy
 from pytorch_lightning.trainer.states import TrainerFn
 from pytorch_lightning.utilities.imports import _TORCH_GREATER_EQUAL_1_10

tests/tests_pytorch/strategies/test_sharded_strategy.py

@@ -5,9 +5,9 @@
 import pytest
 import torch
 
+from lightning_lite.strategies.fairscale import _FAIRSCALE_AVAILABLE
 from pytorch_lightning import LightningModule, Trainer
 from pytorch_lightning.demos.boring_classes import BoringModel
-from pytorch_lightning.overrides.fairscale import _FAIRSCALE_AVAILABLE
 from pytorch_lightning.strategies import DDPShardedStrategy, DDPSpawnShardedStrategy
 from pytorch_lightning.trainer.states import TrainerFn
 from tests_pytorch.helpers.runif import RunIf

tests/tests_pytorch/utilities/test_imports.py

@@ -20,9 +20,7 @@
 from lightning_utilities.core.imports import compare_version, module_available, RequirementCache
 from torch.distributed import is_available
 
-from pytorch_lightning.overrides.fairscale import _FAIRSCALE_AVAILABLE
 from pytorch_lightning.strategies.bagua import _BAGUA_AVAILABLE
-from pytorch_lightning.strategies.deepspeed import _DEEPSPEED_AVAILABLE
 from pytorch_lightning.utilities import _APEX_AVAILABLE, _HOROVOD_AVAILABLE, _OMEGACONF_AVAILABLE, _POPTORCH_AVAILABLE
 
 
@@ -41,20 +39,6 @@ def test_imports():
     else:
         assert _BAGUA_AVAILABLE
 
-    try:
-        import deepspeed  # noqa
-    except ModuleNotFoundError:
-        assert not _DEEPSPEED_AVAILABLE
-    else:
-        assert _DEEPSPEED_AVAILABLE
-
-    try:
-        import fairscale.nn  # noqa
-    except ModuleNotFoundError:
-        assert not _FAIRSCALE_AVAILABLE
-    else:
-        assert _FAIRSCALE_AVAILABLE
-
     try:
         import horovod.torch  # noqa
     except ModuleNotFoundError: