Avoid torchelastic warning message when importing lightning (#15610)

Author: awaelchli

src/lightning_lite/CHANGELOG.md

@@ -50,6 +50,6 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 - Fix an issue with the SLURM `srun` detection causing permission errors ([#15485](https://github.com/Lightning-AI/lightning/issues/15485))
 
--
+- Fixed the import of `lightning_lite` causing a warning 'Redirects are currently not supported in Windows or MacOs' ([#15610](https://github.com/PyTorchLightning/pytorch-lightning/issues/15610))
 
 -

src/lightning_lite/strategies/deepspeed.py

@@ -18,7 +18,7 @@
 import platform
 from contextlib import contextmanager
 from pathlib import Path
-from typing import Any, Dict, Generator, Iterable, List, Mapping, Optional, Tuple, Union
+from typing import Any, Dict, Generator, Iterable, List, Mapping, Optional, Tuple, TYPE_CHECKING, Union
 
 import torch
 from lightning_utilities.core.imports import RequirementCache
@@ -38,7 +38,7 @@
 from lightning_lite.utilities.types import _LRScheduler, _PATH, ReduceLROnPlateau
 
 _DEEPSPEED_AVAILABLE = RequirementCache("deepspeed")
-if _DEEPSPEED_AVAILABLE:
+if TYPE_CHECKING and _DEEPSPEED_AVAILABLE:
     import deepspeed
 
 
@@ -271,6 +271,9 @@ def __init__(
                 reduce_bucket_size=reduce_bucket_size,
                 sub_group_size=sub_group_size,
             )
+
+        import deepspeed
+
         self._config_initialized = False
         deepspeed.utils.logging.logger.setLevel(logging_level)
 
@@ -327,6 +330,9 @@ def module_sharded_context(self) -> Generator[None, None, None]:
         # Current limitation in Lite: The config needs to be fully determined at the time of calling the
         # context manager, which happens at the start of `Lite.run()`. Later modificatoins through e.g. `Lite.setup()`
         # won't have an effect here.
+
+        import deepspeed
+
         if self.zero_stage_3:
             assert self._config_initialized
 
@@ -405,6 +411,8 @@ def _setup_module_and_optimizer(
 
         This calls :func:`deepspeed.initialize` internally.
         """
+        import deepspeed
+
         model_parameters = filter(lambda p: p.requires_grad, model.parameters())
         deepspeed_engine, deepspeed_optimizer, _, _ = deepspeed.initialize(
             args=argparse.Namespace(device_rank=self.root_device.index),
@@ -432,6 +440,8 @@ def _setup_distributed(self) -> None:
             self._config_initialized = True
 
     def _init_deepspeed_distributed(self) -> None:
+        import deepspeed
+
         assert self.cluster_environment is not None
         if platform.system() != "Windows":
             # do not set env variables on windows, allow deepspeed to control setup
@@ -453,6 +463,8 @@ def _set_node_environment_variables(self) -> None:
         os.environ["LOCAL_RANK"] = str(self.local_rank)
 
     def _set_deepspeed_activation_checkpointing(self) -> None:
+        import deepspeed
+
         assert isinstance(self.config, dict)
         if self.config.get("activation_checkpointing"):
             checkpoint_config = self.config["activation_checkpointing"]
@@ -573,6 +585,7 @@ def _restore_zero_state(self, module: Module, ckpt: Mapping[str, Any]) -> None:
         Args:
             ckpt: The ckpt file.
         """
+        import deepspeed
 
         def load(module: torch.nn.Module, prefix: str = "") -> None:
 

src/pytorch_lightning/CHANGELOG.md

@@ -71,7 +71,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 - Fixed model state transfer in multiprocessing launcher when running multi-node ([#15567](https://github.com/Lightning-AI/lightning/pull/15567))
 
-- Fixed manual optimization raising `AttributeError` with Bagua Strategy ([[#12534](https://github.com/PyTorchLightning/pytorch-lightning/issues/12534)])
+- Fixed manual optimization raising `AttributeError` with Bagua Strategy ([#12534](https://github.com/PyTorchLightning/pytorch-lightning/issues/12534))
+
+- Fixed the import of `pytorch_lightning` causing a warning 'Redirects are currently not supported in Windows or MacOs' ([#15610](https://github.com/PyTorchLightning/pytorch-lightning/issues/15610))
 
 
 ## [1.8.0] - 2022-11-01

src/pytorch_lightning/strategies/deepspeed.py

@@ -19,7 +19,7 @@
 import platform
 from collections import OrderedDict
 from pathlib import Path
-from typing import Any, Dict, Generator, List, Mapping, Optional, Tuple, Union
+from typing import Any, Dict, Generator, List, Mapping, Optional, Tuple, TYPE_CHECKING, Union
 
 import torch
 from lightning_utilities.core.apply_func import apply_to_collection
@@ -52,7 +52,7 @@
 warning_cache = WarningCache()
 
 _DEEPSPEED_AVAILABLE = RequirementCache("deepspeed")
-if _DEEPSPEED_AVAILABLE:
+if TYPE_CHECKING and _DEEPSPEED_AVAILABLE:
     import deepspeed
 
 
@@ -319,6 +319,8 @@ def __init__(
                 reduce_bucket_size=reduce_bucket_size,
                 sub_group_size=sub_group_size,
             )
+        import deepspeed
+
         self._config_initialized = False
         deepspeed.utils.logging.logger.setLevel(logging_level)
 
@@ -368,6 +370,8 @@ def setup(self, trainer: "pl.Trainer") -> None:
         self.barrier()
 
     def _init_deepspeed_distributed(self) -> None:
+        import deepspeed
+
         assert self.cluster_environment is not None
         if platform.system() != "Windows":
             # do not set env variables on windows, allow deepspeed to control setup
@@ -428,6 +432,8 @@ def _setup_model_and_optimizer(
 
         This calls :func:`deepspeed.initialize` internally.
         """
+        import deepspeed
+
         model_parameters = filter(lambda p: p.requires_grad, model.parameters())
         deepspeed_engine, deepspeed_optimizer, _, _ = deepspeed.initialize(
             args=argparse.Namespace(device_rank=self.root_device.index),
@@ -527,6 +533,8 @@ def _initialize_deepspeed_train(self, model: Module) -> None:
 
     @contextlib.contextmanager
     def model_sharded_context(self) -> Generator[None, None, None]:
+        import deepspeed
+
         if self.zero_stage_3:
             assert self._config_initialized
 
@@ -547,6 +555,8 @@ def model_sharded_context(self) -> Generator[None, None, None]:
             yield
 
     def _set_deepspeed_activation_checkpointing(self) -> None:
+        import deepspeed
+
         assert isinstance(self.config, dict)
         if self.config.get("activation_checkpointing"):
             checkpoint_config = self.config["activation_checkpointing"]
@@ -559,6 +569,8 @@ def _set_deepspeed_activation_checkpointing(self) -> None:
             )
 
     def _initialize_deepspeed_inference(self, model: Module) -> None:
+        import deepspeed
+
         assert isinstance(self.config, dict)
 
         # todo: this is required for DeepSpeed throughput timers
@@ -639,6 +651,8 @@ def _format_batch_size_and_grad_accum_config(self) -> None:
             self.config["gradient_clipping"] = self.lightning_module.trainer.gradient_clip_val or 0.0
 
     def _auto_select_batch_size(self) -> int:
+        import deepspeed
+
         # train_micro_batch_size_per_gpu is used for throughput logging purposes
         # by default we try to use the batch size of the loader
         assert self.lightning_module is not None
@@ -842,6 +856,7 @@ def _restore_zero_state(self, ckpt: Mapping[str, Any]) -> None:
         Args:
             ckpt: The ckpt file.
         """
+        import deepspeed
 
         assert self.lightning_module is not None
 

tests/tests_lite/utilities/test_imports.py

@@ -15,6 +15,8 @@
 import sys
 from textwrap import dedent
 
+from tests_lite.helpers.runif import RunIf
+
 from lightning_lite.strategies.deepspeed import _DEEPSPEED_AVAILABLE
 from lightning_lite.strategies.fairscale import _FAIRSCALE_AVAILABLE
 
@@ -46,3 +48,24 @@ def test_import_lightning_lite_with_torch_dist_unavailable():
     )
     # run in complete isolation
     assert subprocess.call([sys.executable, "-c", code]) == 0
+
+
+@RunIf(deepspeed=True)
+def test_import_deepspeed_lazily():
+    """Test that we are importing deepspeed only when necessary."""
+    code = dedent(
+        """
+        import lightning_lite
+        import sys
+
+        assert 'deepspeed' not in sys.modules
+        from lightning_lite.strategies import DeepSpeedStrategy
+        from lightning_lite.plugins import DeepSpeedPrecision
+        assert 'deepspeed' not in sys.modules
+
+        import deepspeed
+        assert 'deepspeed' in sys.modules
+        """
+    )
+    # run in complete isolation
+    assert subprocess.call([sys.executable, "-c", code]) == 0

tests/tests_pytorch/utilities/test_imports.py

@@ -25,6 +25,7 @@
 
 from pytorch_lightning.strategies.bagua import _BAGUA_AVAILABLE
 from pytorch_lightning.utilities import _APEX_AVAILABLE, _HOROVOD_AVAILABLE, _OMEGACONF_AVAILABLE, _POPTORCH_AVAILABLE
+from tests_pytorch.helpers.runif import RunIf
 
 
 def test_imports():
@@ -157,3 +158,24 @@ def test_import_pytorch_lightning_with_torch_dist_unavailable():
     )
     # run in complete isolation
     assert subprocess.call([sys.executable, "-c", code]) == 0
+
+
+@RunIf(deepspeed=True)
+def test_import_deepspeed_lazily():
+    """Test that we are importing deepspeed only when necessary."""
+    code = dedent(
+        """
+        import pytorch_lightning
+        import sys
+
+        assert 'deepspeed' not in sys.modules
+        from pytorch_lightning.strategies import DeepSpeedStrategy
+        from pytorch_lightning.plugins import DeepSpeedPrecisionPlugin
+        assert 'deepspeed' not in sys.modules
+
+        import deepspeed
+        assert 'deepspeed' in sys.modules
+        """
+    )
+    # run in complete isolation
+    assert subprocess.call([sys.executable, "-c", code]) == 0