Fixes various typing errors in pytorch_lightning/strategies/deepspeed.py (#13832)

Co-authored-by: otaj <[email protected]>
Co-authored-by: Carlos Mocholí <[email protected]>
Co-authored-by: Rohit Gupta <[email protected]>

Author: 4 people

pyproject.toml

@@ -63,7 +63,6 @@ module = [
     "pytorch_lightning.profilers.simple",
     "pytorch_lightning.strategies.ddp",
     "pytorch_lightning.strategies.ddp_spawn",
-    "pytorch_lightning.strategies.deepspeed",
     "pytorch_lightning.strategies.fully_sharded",
     "pytorch_lightning.strategies.ipu",
     "pytorch_lightning.strategies.sharded",

src/pytorch_lightning/plugins/precision/deepspeed.py

@@ -26,9 +26,10 @@
 from pytorch_lightning.utilities.model_helpers import is_overridden
 from pytorch_lightning.utilities.warnings import WarningCache
 
+_DEEPSPEED_AVAILABLE = _RequirementAvailable("deepspeed")
 _DEEPSPEED_GREATER_EQUAL_0_6 = _RequirementAvailable("deepspeed>=0.6.0")
 if TYPE_CHECKING:
-    if pl.strategies.deepspeed._DEEPSPEED_AVAILABLE:
+    if _DEEPSPEED_AVAILABLE:
         import deepspeed
 
 warning_cache = WarningCache()

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, cast, Dict, Generator, List, Mapping, Optional, Tuple, Union
 
 import torch
 from torch import Tensor
@@ -48,12 +48,12 @@
 from pytorch_lightning.utilities.optimizer import optimizers_to_device
 from pytorch_lightning.utilities.rank_zero import rank_zero_info
 from pytorch_lightning.utilities.seed import reset_seed
-from pytorch_lightning.utilities.types import _PATH, LRSchedulerConfig, LRSchedulerTypeUnion, STEP_OUTPUT
+from pytorch_lightning.utilities.types import _LRScheduler, _PATH, LRSchedulerConfig, ReduceLROnPlateau, STEP_OUTPUT
 from pytorch_lightning.utilities.warnings import rank_zero_warn, WarningCache
 
 warning_cache = WarningCache()
 
-_DEEPSPEED_AVAILABLE: bool = _RequirementAvailable("deepspeed")
+_DEEPSPEED_AVAILABLE = _RequirementAvailable("deepspeed")
 if _DEEPSPEED_AVAILABLE:
     import deepspeed
 
@@ -76,7 +76,7 @@ def __init__(
         super().__init__(pl_module)
         self.precision = precision
 
-    def forward(self, *inputs, **kwargs):
+    def forward(self, *inputs: Any, **kwargs: Any) -> Any:
         inputs = apply_to_collection(inputs, Tensor, function=self._batch_to)
         return super().forward(*inputs, **kwargs)
 
@@ -123,7 +123,7 @@ def __init__(
         reduce_bucket_size: int = 200_000_000,
         zero_allow_untested_optimizer: bool = True,
         logging_batch_size_per_gpu: Union[str, int] = "auto",
-        config: Optional[Union[Path, str, dict]] = None,
+        config: Optional[Union[_PATH, Dict[str, Any]]] = None,
         logging_level: int = logging.WARN,
         parallel_devices: Optional[List[torch.device]] = None,
         cluster_environment: Optional[ClusterEnvironment] = None,
@@ -142,7 +142,7 @@ def __init__(
     ) -> None:
         """Provides capabilities to run training using the DeepSpeed library, with training optimizations for large
         billion parameter models. `For more information: https://pytorch-
-        lightning.readthedocs.io/en/latest/advanced/advanced_gpu.html#deepspeed`.
+        lightning.readthedocs.io/en/stable/advanced/model_parallel.html#deepspeed`.
 
         .. warning:: ``DeepSpeedStrategy`` is in beta and subject to change.
 
@@ -331,7 +331,7 @@ def __init__(
         self.hysteresis = hysteresis
         self.min_loss_scale = min_loss_scale
 
-    def _load_config(self, config):
+    def _load_config(self, config: Optional[Union[_PATH, Dict[str, Any]]]) -> Optional[Dict[str, Any]]:
         if config is None and self.DEEPSPEED_ENV_VAR in os.environ:
             rank_zero_info(f"Loading DeepSpeed config from set {self.DEEPSPEED_ENV_VAR} environment variable")
             config = os.environ[self.DEEPSPEED_ENV_VAR]
@@ -342,9 +342,10 @@ def _load_config(self, config):
                 )
             with open(config) as f:
                 config = json.load(f)
+        assert isinstance(config, dict) or config is None
         return config
 
-    def setup_distributed(self):
+    def setup_distributed(self) -> None:
         reset_seed()
 
         # determine which process we are and world size
@@ -357,8 +358,10 @@ def setup_distributed(self):
             self._config_initialized = True
 
     def setup(self, trainer: "pl.Trainer") -> None:
+        assert self.accelerator is not None
         self.accelerator.setup(trainer)
         # we set the device so that optimizers can be created with distributed comms.
+        assert self.lightning_module is not None
         self.lightning_module._device = self.root_device
         self.setup_optimizers(trainer)
         self.setup_precision_plugin()
@@ -367,6 +370,7 @@ def setup(self, trainer: "pl.Trainer") -> None:
         self.barrier()
 
     def _init_deepspeed_distributed(self) -> None:
+        assert self.cluster_environment is not None
         if platform.system() != "Windows":
             # do not set env variables on windows, allow deepspeed to control setup
             self._set_node_environment_variables()
@@ -378,14 +382,15 @@ def _init_deepspeed_distributed(self) -> None:
         self._process_group_backend = self._get_process_group_backend()
         deepspeed.init_distributed(self._process_group_backend, distributed_port=self.cluster_environment.main_port)
 
-    def _get_process_group_backend(self):
+    def _get_process_group_backend(self) -> str:
         return (
             self._process_group_backend
             or _get_process_group_backend_from_env()
             or get_default_process_group_backend_for_device(self.root_device)
         )
 
     def _set_node_environment_variables(self) -> None:
+        assert self.cluster_environment is not None
         os.environ["MASTER_ADDR"] = self.cluster_environment.main_address
         os.environ["MASTER_PORT"] = str(self.cluster_environment.main_port)
         os.environ["RANK"] = str(self.global_rank)
@@ -396,7 +401,9 @@ def _set_node_environment_variables(self) -> None:
     def restore_checkpoint_after_setup(self) -> bool:
         return True
 
-    def _setup_model_and_optimizers(self, model: Module, optimizers: List[Optimizer]) -> Tuple[Module, List[Optimizer]]:
+    def _setup_model_and_optimizers(
+        self, model: Module, optimizers: List[Optimizer]
+    ) -> Tuple["deepspeed.DeepSpeedEngine", List[Optimizer]]:
         """Setup a model and multiple optimizers together.
 
         Currently only a single optimizer is supported.
@@ -414,14 +421,18 @@ def _setup_model_and_optimizers(self, model: Module, optimizers: List[Optimizer]
         # train_micro_batch_size_per_gpu is used for throughput logging purposes
         # normally we set this to the batch size, but it is not available here unless the user provides it
         # as part of the config
+        assert self.config is not None
         self.config.setdefault("train_micro_batch_size_per_gpu", 1)
         self.model, optimizer = self._setup_model_and_optimizer(model, optimizers[0])
         self._set_deepspeed_activation_checkpointing()
         return self.model, [optimizer]
 
     def _setup_model_and_optimizer(
-        self, model: Module, optimizer: Optimizer, lr_scheduler: Optional[LRSchedulerTypeUnion] = None
-    ):
+        self,
+        model: Module,
+        optimizer: Optional[Optimizer],
+        lr_scheduler: Optional[Union[_LRScheduler, ReduceLROnPlateau]] = None,
+    ) -> Tuple["deepspeed.DeepSpeedEngine", Optimizer]:
         """Initialize one model and one optimizer with an optional learning rate scheduler.
 
         This calls :func:`deepspeed.initialize` internally.
@@ -431,14 +442,15 @@ def _setup_model_and_optimizer(
             args=argparse.Namespace(device_rank=self.root_device.index),
             config=self.config,
             model=model,
-            model_parameters=model_parameters,  # type: ignore
+            model_parameters=model_parameters,
             optimizer=optimizer,
             lr_scheduler=lr_scheduler,
             dist_init_required=False,
         )
         return deepspeed_engine, deepspeed_optimizer
 
-    def init_deepspeed(self):
+    def init_deepspeed(self) -> None:
+        assert self.lightning_module is not None
         # deepspeed handles gradient clipping internally
         if is_overridden("configure_gradient_clipping", self.lightning_module, pl.LightningModule):
             rank_zero_warn(
@@ -464,6 +476,7 @@ def init_deepspeed(self):
                 "DeepSpeed currently does not support different `accumulate_grad_batches` at different epochs."
             )
 
+        assert isinstance(self.model, (pl.LightningModule, _LightningPrecisionModuleWrapperBase))
         model = LightningDeepSpeedModule(pl_module=self.model, precision=self.precision_plugin.precision)
 
         if self.lightning_module.trainer and self.lightning_module.trainer.training:
@@ -472,6 +485,7 @@ def init_deepspeed(self):
             self._initialize_deepspeed_inference(model)
 
     def _init_optimizers(self) -> Tuple[Optimizer, Optional[LRSchedulerConfig], Optional[int]]:
+        assert self.lightning_module is not None
         optimizers, lr_schedulers, optimizer_frequencies = _init_optimizers_and_lr_schedulers(self.lightning_module)
         if len(optimizers) > 1 or len(lr_schedulers) > 1:
             raise MisconfigurationException(
@@ -485,10 +499,13 @@ def _init_optimizers(self) -> Tuple[Optimizer, Optional[LRSchedulerConfig], Opti
 
     @property
     def zero_stage_3(self) -> bool:
-        return self.config.get("zero_optimization") and self.config.get("zero_optimization").get("stage") == 3
+        assert isinstance(self.config, dict)
+        zero_optimization = self.config.get("zero_optimization")
+        return zero_optimization is not None and zero_optimization.get("stage") == 3
 
-    def _initialize_deepspeed_train(self, model):
+    def _initialize_deepspeed_train(self, model: Module) -> None:
         optimizer, scheduler = None, None
+        assert isinstance(self.config, dict)
         if "optimizer" in self.config:
             rank_zero_info(
                 "You have specified an optimizer and/or scheduler within the DeepSpeed config."
@@ -538,7 +555,8 @@ def model_sharded_context(self) -> Generator[None, None, None]:
         with model_parallel_context:
             yield
 
-    def _set_deepspeed_activation_checkpointing(self):
+    def _set_deepspeed_activation_checkpointing(self) -> None:
+        assert isinstance(self.config, dict)
         if self.config.get("activation_checkpointing"):
             checkpoint_config = self.config["activation_checkpointing"]
             deepspeed.checkpointing.configure(
@@ -549,8 +567,9 @@ def _set_deepspeed_activation_checkpointing(self):
                 profile=checkpoint_config.get("profile"),
             )
 
-    def _initialize_deepspeed_inference(self, model):
+    def _initialize_deepspeed_inference(self, model: Module) -> None:
         # todo: Currently DeepSpeed requires optimizers at inference to partition weights correctly
+        assert isinstance(self.config, dict)
         optimizer, scheduler = None, None
         if "optimizer" not in self.config:
             rank_zero_info(
@@ -585,13 +604,15 @@ def _initialize_deepspeed_inference(self, model):
         self.model = model
 
     @property
-    def lightning_module(self):
+    def lightning_module(self) -> Optional["pl.LightningModule"]:
         # the model may not be wrapped with DeepEngine & LightningDeepSpeedModule if calling this too early
         module = getattr(self.model, "module", self.model)
-        return module.module if isinstance(module, LightningDeepSpeedModule) else module
+        module = module.module if isinstance(module, LightningDeepSpeedModule) else module
+        assert isinstance(module, pl.LightningModule) or module is None
+        return module
 
     @property
-    def distributed_sampler_kwargs(self):
+    def distributed_sampler_kwargs(self) -> Dict[str, int]:
         distributed_sampler_kwargs = dict(num_replicas=self.world_size, rank=self.global_rank)
         return distributed_sampler_kwargs
 
@@ -616,17 +637,18 @@ def handles_gradient_accumulation(self) -> bool:
         """Whether the plugin handles gradient accumulation internally."""
         return True
 
-    def _format_config(self):
+    def _format_config(self) -> None:
         if self.config is None:
             raise MisconfigurationException(
                 "To use DeepSpeed you must pass in a DeepSpeed config dict, or a path to a JSON config."
-                " See: https://pytorch-lightning.readthedocs.io/en/latest/advanced/advanced_gpu.html#deepspeed"
+                " See: https://pytorch-lightning.readthedocs.io/en/stable/advanced/model_parallel.html#deepspeed"
             )
         self._format_batch_size_and_grad_accum_config()
         self._format_precision_config()
 
-    def _format_batch_size_and_grad_accum_config(self):
+    def _format_batch_size_and_grad_accum_config(self) -> None:
         # todo: using lite, we do not support these variables within the config
+        assert isinstance(self.config, dict)
         if self.lightning_module is None:
             return
 
@@ -642,16 +664,17 @@ def _format_batch_size_and_grad_accum_config(self):
         if "gradient_clipping" not in self.config:
             self.config["gradient_clipping"] = self.lightning_module.trainer.gradient_clip_val or 0.0
 
-    def _auto_select_batch_size(self):
+    def _auto_select_batch_size(self) -> int:
         # 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
         batch_size = 1
         train_dl_source = self.lightning_module.trainer._data_connector._train_dataloader_source
         if train_dl_source.is_defined():
             try:
                 train_dataloader = train_dl_source.dataloader()
                 if hasattr(train_dataloader, "batch_sampler"):
-                    batch_size = train_dataloader.batch_sampler.batch_size
+                    batch_size = train_dataloader.batch_sampler.batch_size  # type: ignore[union-attr]
             # broad exception on purpose as `source.dataloader()` will fail if the dataloader requires `setup`
             # to have been called before
             except Exception:
@@ -664,6 +687,7 @@ def _auto_select_batch_size(self):
         return batch_size
 
     def _format_precision_config(self) -> None:
+        assert isinstance(self.config, dict)
         if self.precision_plugin.precision in (PrecisionType.HALF, PrecisionType.MIXED):
             if "fp16" not in self.config and self.precision_plugin.amp_type == AMPType.NATIVE:
                 # FP16 is a DeepSpeed standalone AMP implementation
@@ -707,7 +731,7 @@ def _create_default_config(
         single_submit: bool,
         overlap_events: bool,
         thread_count: int,
-        **zero_kwargs,
+        **zero_kwargs: Any,
     ) -> Dict:
         cfg = {
             "activation_checkpointing": {
@@ -753,7 +777,7 @@ def _create_default_config(
         return cfg
 
     @property
-    def deepspeed_engine(self):
+    def deepspeed_engine(self) -> "deepspeed.DeepSpeedEngine":
         return self.model
 
     @property
@@ -786,7 +810,7 @@ def save_checkpoint(self, checkpoint: Dict, filepath: _PATH, storage_options: Op
                 "When saving the DeepSpeed Stage 3 checkpoint, "
                 "each worker will save a shard of the checkpoint within a directory. "
                 "If a single file is required after training, "
-                "see https://pytorch-lightning.readthedocs.io/en/latest/advanced/advanced_gpu.html#"
+                "see https://pytorch-lightning.readthedocs.io/en/stable/advanced/model_parallel.html#"
                 "deepspeed-zero-stage-3-single-file for instructions."
             )
         # Use deepspeed's internal checkpointing function to handle partitioned weights across processes
@@ -799,10 +823,12 @@ def load_checkpoint(self, checkpoint_path: _PATH) -> Dict[str, Any]:
         if self.load_full_weights and self.zero_stage_3:
             # Broadcast to ensure we load from the rank 0 checkpoint
             # This doesn't have to be the case when using deepspeed sharded checkpointing
-            checkpoint_path = self.broadcast(checkpoint_path)
+            checkpoint_path = cast(_PATH, self.broadcast(checkpoint_path))
             return super().load_checkpoint(checkpoint_path)
 
         # Rely on deepspeed to load the checkpoint and necessary information
+        assert self.lightning_module is not None
+
         from pytorch_lightning.trainer.states import TrainerFn
 
         is_fitting = self.lightning_module.trainer.state.fn == TrainerFn.FITTING
@@ -818,6 +844,7 @@ def load_checkpoint(self, checkpoint_path: _PATH) -> Dict[str, Any]:
 
     @property
     def lightning_restore_optimizer(self) -> bool:
+        assert self.lightning_module is not None
         # managed by DeepSpeed
         if self.load_full_weights and self.zero_stage_3 and self.lightning_module.trainer.state.fn == TrainerFn.FITTING:
             rank_zero_warn(
@@ -842,11 +869,13 @@ def _restore_zero_state(self, ckpt: Mapping[str, Any]) -> None:
             ckpt: The ckpt file.
         """
 
-        def load(module: torch.nn.Module, prefix=""):
+        assert self.lightning_module is not None
+
+        def load(module: torch.nn.Module, prefix: str = "") -> None:
 
-            missing_keys = []
-            unexpected_keys = []
-            error_msgs = []
+            missing_keys: List[str] = []
+            unexpected_keys: List[str] = []
+            error_msgs: List[str] = []
             state_dict = ckpt["state_dict"]
 
             # copy state_dict so _load_from_state_dict can modify it
@@ -914,14 +943,17 @@ def register_strategies(cls, strategy_registry: Dict) -> None:
             offload_optimizer_device="nvme",
         )
 
-    def validation_step(self, *args, **kwargs) -> Optional[STEP_OUTPUT]:
+    def validation_step(self, *args: Any, **kwargs: Any) -> Optional[STEP_OUTPUT]:
+        assert self.model is not None
         with self.precision_plugin.val_step_context():
             return self.model(*args, **kwargs)
 
-    def test_step(self, *args, **kwargs) -> Optional[STEP_OUTPUT]:
+    def test_step(self, *args: Any, **kwargs: Any) -> Optional[STEP_OUTPUT]:
+        assert self.model is not None
         with self.precision_plugin.test_step_context():
             return self.model(*args, **kwargs)
 
-    def predict_step(self, *args, **kwargs) -> STEP_OUTPUT:
+    def predict_step(self, *args: Any, **kwargs: Any) -> STEP_OUTPUT:
+        assert self.model is not None
         with self.precision_plugin.predict_step_context():
             return self.model(*args, **kwargs)