Add @override for files in src/lightning/fabric/strategies (#19185)

Author: VictorPrins

src/lightning/fabric/strategies/ddp.py

@@ -21,6 +21,7 @@
 from torch import Tensor
 from torch.nn import Module
 from torch.nn.parallel.distributed import DistributedDataParallel
+from typing_extensions import override
 
 from lightning.fabric.accelerators.accelerator import Accelerator
 from lightning.fabric.plugins.collectives.torch_collective import default_pg_timeout
@@ -78,6 +79,7 @@ def __init__(
         self._ddp_kwargs = kwargs
 
     @property
+    @override
     def root_device(self) -> torch.device:
         assert self.parallel_devices is not None
         return self.parallel_devices[self.local_rank]
@@ -96,24 +98,28 @@ def num_processes(self) -> int:
         return len(self.parallel_devices) if self.parallel_devices is not None else 0
 
     @property
+    @override
     def distributed_sampler_kwargs(self) -> Dict[str, Any]:
         return {"num_replicas": (self.num_nodes * self.num_processes), "rank": self.global_rank}
 
     @property
     def process_group_backend(self) -> Optional[str]:
         return self._process_group_backend
 
+    @override
     def _configure_launcher(self) -> None:
         assert self.cluster_environment is not None
         if self._start_method == "popen":
             self._launcher = _SubprocessScriptLauncher(self.cluster_environment, self.num_processes, self.num_nodes)
         else:
             self._launcher = _MultiProcessingLauncher(self, start_method=self._start_method)
 
+    @override
     def setup_environment(self) -> None:
         self._setup_distributed()
         super().setup_environment()
 
+    @override
     def setup_module(self, module: Module) -> DistributedDataParallel:
         """Wraps the model into a :class:`~torch.nn.parallel.distributed.DistributedDataParallel` module."""
         device_ids = self._determine_ddp_device_ids()
@@ -122,9 +128,11 @@ def setup_module(self, module: Module) -> DistributedDataParallel:
         with ctx:
             return DistributedDataParallel(module=module, device_ids=device_ids, **self._ddp_kwargs)
 
+    @override
     def module_to_device(self, module: Module) -> None:
         module.to(self.root_device)
 
+    @override
     def all_reduce(
         self, tensor: Tensor, group: Optional[Any] = None, reduce_op: Optional[Union[ReduceOp, str]] = "mean"
     ) -> Tensor:
@@ -144,6 +152,7 @@ def all_reduce(
             return _sync_ddp_if_available(tensor, group, reduce_op=reduce_op)
         return tensor
 
+    @override
     def barrier(self, *args: Any, **kwargs: Any) -> None:
         if not _distributed_is_initialized():
             return
@@ -152,6 +161,7 @@ def barrier(self, *args: Any, **kwargs: Any) -> None:
         else:
             torch.distributed.barrier()
 
+    @override
     def broadcast(self, obj: TBroadcast, src: int = 0) -> TBroadcast:
         if not _distributed_is_initialized():
             return obj
@@ -160,11 +170,13 @@ def broadcast(self, obj: TBroadcast, src: int = 0) -> TBroadcast:
         torch.distributed.broadcast_object_list(obj, src, group=_group.WORLD)
         return obj[0]
 
+    @override
     def get_module_state_dict(self, module: Module) -> Dict[str, Union[Any, Tensor]]:
         if isinstance(module, DistributedDataParallel):
             module = module.module
         return super().get_module_state_dict(module)
 
+    @override
     def load_module_state_dict(
         self, module: Module, state_dict: Dict[str, Union[Any, Tensor]], strict: bool = True
     ) -> None:
@@ -173,6 +185,7 @@ def load_module_state_dict(
         super().load_module_state_dict(module=module, state_dict=state_dict, strict=strict)
 
     @classmethod
+    @override
     def register_strategies(cls, strategy_registry: _StrategyRegistry) -> None:
         entries = (
             ("ddp", "popen"),
@@ -210,6 +223,7 @@ def _determine_ddp_device_ids(self) -> Optional[List[int]]:
 
 
 class _DDPBackwardSyncControl(_BackwardSyncControl):
+    @override
     def no_backward_sync(self, module: Module) -> ContextManager:
         """Blocks gradient synchronization inside the :class:`~torch.nn.parallel.distributed.DistributedDataParallel`
         wrapper."""

src/lightning/fabric/strategies/deepspeed.py

@@ -25,6 +25,7 @@
 from lightning_utilities.core.imports import RequirementCache
 from torch.nn import Module
 from torch.optim import Optimizer
+from typing_extensions import override
 
 from lightning.fabric.accelerators import Accelerator, CUDAAccelerator
 from lightning.fabric.plugins.environments.cluster_environment import ClusterEnvironment
@@ -299,13 +300,15 @@ def zero_stage_3(self) -> bool:
         return zero_optimization is not None and zero_optimization.get("stage") == 3
 
     @property
+    @override
     def distributed_sampler_kwargs(self) -> Dict[str, int]:
         return {"num_replicas": self.world_size, "rank": self.global_rank}
 
     @property
     def model(self) -> "DeepSpeedEngine":
         return self._deepspeed_engine
 
+    @override
     def setup_module_and_optimizers(
         self, module: Module, optimizers: List[Optimizer]
     ) -> Tuple["DeepSpeedEngine", List[Optimizer]]:
@@ -328,6 +331,7 @@ def setup_module_and_optimizers(
         self._set_deepspeed_activation_checkpointing()
         return self._deepspeed_engine, [optimizer]
 
+    @override
     def setup_module(self, module: Module) -> "DeepSpeedEngine":
         """Set up a module for inference (no optimizers).
 
@@ -337,6 +341,7 @@ def setup_module(self, module: Module) -> "DeepSpeedEngine":
         self._deepspeed_engine, _ = self._initialize_engine(module)
         return self._deepspeed_engine
 
+    @override
     def setup_optimizer(self, optimizer: Optimizer) -> Optimizer:
         """Optimizers can only be set up jointly with the model in this strategy.
 
@@ -345,6 +350,7 @@ def setup_optimizer(self, optimizer: Optimizer) -> Optimizer:
         """
         raise NotImplementedError(self._err_msg_joint_setup_required())
 
+    @override
     def module_init_context(self, empty_init: Optional[bool] = None) -> ContextManager:
         if self.zero_stage_3 and empty_init is False:
             raise NotImplementedError(
@@ -357,6 +363,7 @@ def module_init_context(self, empty_init: Optional[bool] = None) -> ContextManag
         stack.enter_context(module_sharded_ctx)
         return stack
 
+    @override
     def module_sharded_context(self) -> ContextManager:
         # Current limitation in Fabric: The config needs to be fully determined at the time of calling the context
         # manager. Later modifications through e.g. `Fabric.setup()` won't have an effect here.
@@ -370,6 +377,7 @@ def module_sharded_context(self) -> ContextManager:
             config_dict_or_path=self.config,
         )
 
+    @override
     def save_checkpoint(
         self,
         path: _PATH,
@@ -434,6 +442,7 @@ def save_checkpoint(
         # use deepspeed's internal checkpointing function to handle partitioned weights across processes
         engine.save_checkpoint(path, client_state=state, tag="checkpoint")
 
+    @override
     def load_checkpoint(
         self,
         path: _PATH,
@@ -514,6 +523,7 @@ def load_checkpoint(
         _move_state_into(source=client_state, destination=state, keys=keys)
         return client_state
 
+    @override
     def clip_gradients_norm(
         self,
         module: "DeepSpeedEngine",
@@ -527,6 +537,7 @@ def clip_gradients_norm(
             "Make sure to set the `gradient_clipping` value in your Config."
         )
 
+    @override
     def clip_gradients_value(
         self, module: "DeepSpeedEngine", optimizer: Optimizer, clip_val: Union[float, int]
     ) -> None:
@@ -536,6 +547,7 @@ def clip_gradients_value(
         )
 
     @classmethod
+    @override
     def register_strategies(cls, strategy_registry: _StrategyRegistry) -> None:
         strategy_registry.register("deepspeed", cls, description="Default DeepSpeed Strategy")
         strategy_registry.register("deepspeed_stage_1", cls, description="DeepSpeed with ZeRO Stage 1 enabled", stage=1)
@@ -591,6 +603,7 @@ def _initialize_engine(
         )
         return deepspeed_engine, deepspeed_optimizer
 
+    @override
     def _setup_distributed(self) -> None:
         if not isinstance(self.accelerator, CUDAAccelerator):
             raise RuntimeError(

src/lightning/fabric/strategies/dp.py

@@ -16,6 +16,7 @@
 import torch
 from torch import Tensor
 from torch.nn import DataParallel, Module
+from typing_extensions import override
 
 from lightning.fabric.accelerators import Accelerator
 from lightning.fabric.plugins.io.checkpoint_io import CheckpointIO
@@ -47,25 +48,31 @@ def __init__(
         )
 
     @property
+    @override
     def root_device(self) -> torch.device:
         assert self.parallel_devices is not None
         return self.parallel_devices[0]
 
     @property
+    @override
     def distributed_sampler_kwargs(self) -> None:
         return None
 
+    @override
     def setup_module(self, module: Module) -> DataParallel:
         """Wraps the given model into a :class:`~torch.nn.DataParallel` module."""
         return DataParallel(module=module, device_ids=self.parallel_devices)
 
+    @override
     def module_to_device(self, module: Module) -> None:
         module.to(self.root_device)
 
+    @override
     def batch_to_device(self, batch: Any, device: Optional[torch.device] = None) -> Any:
         # DataParallel handles the transfer of batch to the device
         return batch
 
+    @override
     def all_reduce(
         self, collection: TReduce, group: Optional[Any] = None, reduce_op: Optional[Union[ReduceOp, str]] = "mean"
     ) -> TReduce:
@@ -75,20 +82,25 @@ def mean(t: Tensor) -> Tensor:
 
         return apply_to_collection(collection, Tensor, mean)
 
+    @override
     def barrier(self, *args: Any, **kwargs: Any) -> None:
         pass
 
+    @override
     def broadcast(self, obj: TBroadcast, src: int = 0) -> TBroadcast:
         return obj
 
+    @override
     def reduce_boolean_decision(self, decision: bool, all: bool = True) -> bool:
         return decision
 
+    @override
     def get_module_state_dict(self, module: Module) -> Dict[str, Union[Any, Tensor]]:
         if isinstance(module, DataParallel):
             module = module.module
         return super().get_module_state_dict(module)
 
+    @override
     def load_module_state_dict(
         self, module: Module, state_dict: Dict[str, Union[Any, Tensor]], strict: bool = True
     ) -> None:
@@ -97,5 +109,6 @@ def load_module_state_dict(
         super().load_module_state_dict(module=module, state_dict=state_dict, strict=strict)
 
     @classmethod
+    @override
     def register_strategies(cls, strategy_registry: _StrategyRegistry) -> None:
         strategy_registry.register("dp", cls, description=cls.__name__)