Introduce CheckpointIO Plugin

CHANGELOG.md

@@ -37,6 +37,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
     * Added `FastForwardSampler` and `CaptureIterableDataset` injection to data loading utilities ([#8366](https://github.com/PyTorchLightning/pytorch-lightning/pull/8366))
 
 
+- Added `CheckpointIOPlugin` to expose checkpoint IO from training type plugin ([#8743](https://github.com/PyTorchLightning/pytorch-lightning/pull/8743))
+
+
 ### Changed
 
 - Parsing of the `gpus` Trainer argument has changed: `gpus="n"` (str) no longer selects the GPU index n and instead selects the first n devices. ([#8770](https://github.com/PyTorchLightning/pytorch-lightning/pull/8770))

pytorch_lightning/plugins/__init__.py

@@ -1,4 +1,6 @@
 from pytorch_lightning.plugins.base_plugin import Plugin
+from pytorch_lightning.plugins.checkpoint.checkpoint import CheckpointIOPlugin
+from pytorch_lightning.plugins.checkpoint.torch import TorchCheckpointIOPlugin
 from pytorch_lightning.plugins.plugins_registry import (  # noqa: F401
     call_training_type_register_plugins,
     TrainingTypePluginsRegistry,
@@ -29,6 +31,8 @@
 from pytorch_lightning.plugins.training_type.training_type_plugin import TrainingTypePlugin
 
 __all__ = [
+    "CheckpointIOPlugin",
+    "TorchCheckpointIOPlugin",
     "ApexMixedPrecisionPlugin",
     "DataParallelPlugin",
     "DDP2Plugin",

pytorch_lightning/plugins/checkpoint/__init__.py

@@ -0,0 +1,2 @@
+from pytorch_lightning.plugins.checkpoint.checkpoint import CheckpointIOPlugin  # noqa: F401
+from pytorch_lightning.plugins.checkpoint.torch import TorchCheckpointIOPlugin  # noqa: F401

pytorch_lightning/plugins/checkpoint/checkpoint.py

@@ -0,0 +1,33 @@
+from abc import ABC, abstractmethod
+from pathlib import Path
+from typing import Any, Callable, Dict, Mapping, Optional, TypeVar, Union
+
+TSaveStorageOptions = TypeVar("TSaveStorageOptions", Mapping, Callable)
+TLoadStorageOptions = TypeVar("TLoadStorageOptions", Mapping, Callable)
+
+
+class CheckpointIOPlugin(ABC):
+    @abstractmethod
+    def save_checkpoint(
+        self, checkpoint: Dict[str, Any], path: Union[str, Path], storage_options: Optional[TSaveStorageOptions] = None
+    ) -> None:
+        """Save model/training states as a checkpoint file through state-dump and file-write.
+
+        Args:
+            checkpoint: dict containing model and trainer state
+            path: write-target path
+            storage_options: Optional parameters when saving the model/training states.
+        """
+
+    @abstractmethod
+    def load_checkpoint(
+        self, path: Union[str, Path], storage_options: Optional[TLoadStorageOptions] = None
+    ) -> Dict[str, Any]:
+        """
+        Load checkpoint from a path when resuming or loading ckpt for test/validate/predict stages.
+        Args:
+            path: Path to checkpoint
+            storage_options: Optional parameters when loading the model/training states.
+
+        Returns: The loaded checkpoint.
+        """

pytorch_lightning/plugins/checkpoint/torch.py

@@ -0,0 +1,27 @@
+from pathlib import Path
+from typing import Any, Callable, Dict, Optional, Union
+
+import pytorch_lightning as pl
+from pytorch_lightning.plugins.checkpoint.checkpoint import CheckpointIOPlugin
+from pytorch_lightning.utilities import rank_zero_warn
+from pytorch_lightning.utilities.cloud_io import atomic_save
+from pytorch_lightning.utilities.cloud_io import load as pl_load
+
+
+class TorchCheckpointIOPlugin(CheckpointIOPlugin):
+    def save_checkpoint(
+        self, checkpoint: Dict[str, Any], path: Union[str, Path], storage_options: Optional[Any] = None
+    ) -> None:
+        try:
+            # write the checkpoint dictionary on the file
+            atomic_save(checkpoint, path)
+        except AttributeError as err:
+            key = pl.LightningModule.CHECKPOINT_HYPER_PARAMS_KEY
+            checkpoint.pop(key, None)
+            rank_zero_warn(f"Warning, `{key}` dropped from checkpoint. An attribute is not picklable: {err}")
+            atomic_save(checkpoint, path)
+
+    def load_checkpoint(
+        self, path: Union[str, Path], storage_options: Optional[Callable] = lambda storage, loc: storage
+    ) -> Dict[str, Any]:
+        return pl_load(path, map_location=storage_options)

pytorch_lightning/plugins/training_type/ddp.py

@@ -31,6 +31,7 @@
 from pytorch_lightning.distributed import LightningDistributed
 from pytorch_lightning.overrides import LightningDistributedModule
 from pytorch_lightning.overrides.distributed import prepare_for_backward
+from pytorch_lightning.plugins.checkpoint.checkpoint import CheckpointIOPlugin
 from pytorch_lightning.plugins.environments.cluster_environment import ClusterEnvironment
 from pytorch_lightning.plugins.training_type.parallel import ParallelPlugin
 from pytorch_lightning.utilities import (
@@ -75,14 +76,19 @@ def __init__(
         self,
         parallel_devices: Optional[List[torch.device]] = None,
         num_nodes: Optional[int] = None,
-        cluster_environment: ClusterEnvironment = None,
+        cluster_environment: Optional[ClusterEnvironment] = None,
+        checkpoint_plugin: Optional[CheckpointIOPlugin] = None,
         sync_batchnorm: Optional[bool] = None,
         ddp_comm_state: Optional[object] = None,
         ddp_comm_hook: Optional[callable] = None,
         ddp_comm_wrapper: Optional[callable] = None,
         **kwargs: Union[Any, Dict[str, Any]],
     ) -> None:
-        super().__init__(parallel_devices=parallel_devices, cluster_environment=cluster_environment)
+        super().__init__(
+            parallel_devices=parallel_devices,
+            cluster_environment=cluster_environment,
+            checkpoint_plugin=checkpoint_plugin,
+        )
         self.interactive_ddp_procs = []
         if num_nodes is not None:
             rank_zero_deprecation(

pytorch_lightning/plugins/training_type/ddp_spawn.py

@@ -24,6 +24,7 @@
 from pytorch_lightning.distributed.dist import LightningDistributed
 from pytorch_lightning.overrides import LightningDistributedModule
 from pytorch_lightning.overrides.distributed import prepare_for_backward
+from pytorch_lightning.plugins.checkpoint.checkpoint import CheckpointIOPlugin
 from pytorch_lightning.plugins.environments.cluster_environment import ClusterEnvironment
 from pytorch_lightning.plugins.training_type.parallel import ParallelPlugin
 from pytorch_lightning.trainer.states import TrainerFn
@@ -62,14 +63,19 @@ def __init__(
         self,
         parallel_devices: Optional[List[torch.device]] = None,
         num_nodes: Optional[int] = None,
-        cluster_environment: ClusterEnvironment = None,
+        cluster_environment: Optional[ClusterEnvironment] = None,
+        checkpoint_plugin: Optional[CheckpointIOPlugin] = None,
         sync_batchnorm: Optional[bool] = None,
         ddp_comm_state: Optional[object] = None,
         ddp_comm_hook: Optional[callable] = None,
         ddp_comm_wrapper: Optional[callable] = None,
         **kwargs: Any,
     ):
-        super().__init__(parallel_devices=parallel_devices, cluster_environment=cluster_environment)
+        super().__init__(
+            parallel_devices=parallel_devices,
+            cluster_environment=cluster_environment,
+            checkpoint_plugin=checkpoint_plugin,
+        )
         if num_nodes is not None:
             rank_zero_deprecation(
                 "Argument `num_nodes` in `DDPSpawnPlugin` is deprecated in v1.4, and will be removed in v1.6. "

pytorch_lightning/plugins/training_type/deepspeed.py

@@ -26,6 +26,7 @@
 import pytorch_lightning as pl
 from pytorch_lightning.callbacks import GradientAccumulationScheduler
 from pytorch_lightning.overrides.base import _LightningModuleWrapperBase
+from pytorch_lightning.plugins.checkpoint.checkpoint import CheckpointIOPlugin
 from pytorch_lightning.plugins.environments.cluster_environment import ClusterEnvironment
 from pytorch_lightning.plugins.training_type.ddp import DDPPlugin
 from pytorch_lightning.trainer.optimizers import _get_default_scheduler_config
@@ -274,8 +275,11 @@ def __init__(
             pin_memory = cpu_offload_use_pin_memory
 
         super().__init__(
-            parallel_devices=parallel_devices, num_nodes=num_nodes, cluster_environment=cluster_environment
+            parallel_devices=parallel_devices,
+            num_nodes=num_nodes,
+            cluster_environment=cluster_environment,
         )
+
         self.config = self._load_config(config)
         if self.config is None:
             # User has not overridden config, set defaults
@@ -679,6 +683,7 @@ def save_checkpoint(self, checkpoint: Dict, filepath: str) -> None:
             filepath: write-target file's path
         """
         if self.zero_stage_3 and self._multi_device and self.is_global_zero:
+            # todo (sean): Add link to docs once docs are merged.
             warning_cache.warn(
                 "When saving the DeepSpeed Stage 3 checkpoint, "
                 "each worker will save a shard of the checkpoint within a directory. "
@@ -818,3 +823,11 @@ def register_plugins(cls, plugin_registry: Dict) -> None:
             offload_params_device="nvme",
             offload_optimizer_device="nvme",
         )
+
+    @property
+    def checkpoint_plugin(self) -> CheckpointIOPlugin:
+        return self._checkpoint_plugin
+
+    @checkpoint_plugin.setter
+    def checkpoint_plugin(self, plugin: CheckpointIOPlugin) -> None:
+        raise MisconfigurationException("DeepSpeed currently does not support custom checkpoint plugins.")

pytorch_lightning/plugins/training_type/dp.py

@@ -17,6 +17,7 @@
 from torch.nn import DataParallel
 
 from pytorch_lightning.overrides.data_parallel import LightningParallelModule
+from pytorch_lightning.plugins.checkpoint.checkpoint import CheckpointIOPlugin
 from pytorch_lightning.plugins.training_type.parallel import ParallelPlugin
 from pytorch_lightning.utilities.apply_func import apply_to_collection
 from pytorch_lightning.utilities.model_helpers import is_overridden
@@ -29,8 +30,14 @@ class DataParallelPlugin(ParallelPlugin):
     device and each gets a split of the data.
     """
 
-    def __init__(self, parallel_devices: Optional[List[torch.device]]):
-        super().__init__(parallel_devices=parallel_devices, cluster_environment=None)
+    def __init__(
+        self,
+        parallel_devices: Optional[List[torch.device]],
+        checkpoint_plugin: Optional[CheckpointIOPlugin] = None,
+    ):
+        super().__init__(
+            parallel_devices=parallel_devices, cluster_environment=None, checkpoint_plugin=checkpoint_plugin
+        )
 
     @property
     def global_rank(self) -> int:

pytorch_lightning/plugins/training_type/fully_sharded.py

@@ -12,11 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import contextlib
-from typing import Any, Dict, Generator, List, Optional, Union
+from typing import Dict, Generator, List, Optional
 
 import torch
-from torch import Tensor
 
+from pytorch_lightning.plugins.checkpoint.checkpoint import CheckpointIOPlugin
 from pytorch_lightning.plugins.environments.cluster_environment import ClusterEnvironment
 from pytorch_lightning.plugins.training_type.ddp import DDPPlugin
 from pytorch_lightning.utilities import _FAIRSCALE_FULLY_SHARDED_AVAILABLE
@@ -40,7 +40,8 @@ def __init__(
         min_num_params: int = 1e8,
         state_dict_to_cpu: bool = True,
         parallel_devices: Optional[List[torch.device]] = None,
-        cluster_environment: ClusterEnvironment = None,
+        cluster_environment: Optional[ClusterEnvironment] = None,
+        checkpoint_plugin: Optional[CheckpointIOPlugin] = None,
     ):
         """
         Plugin for Fully Sharded Data Parallel provided by FairScale.
@@ -89,7 +90,11 @@ def __init__(
                 (Defautl: True).
         """
 
-        super().__init__(parallel_devices=parallel_devices, cluster_environment=cluster_environment)
+        super().__init__(
+            parallel_devices=parallel_devices,
+            cluster_environment=cluster_environment,
+            checkpoint_plugin=checkpoint_plugin,
+        )
         self.cpu_offload = cpu_offload
         self.move_grads_to_cpu = move_grads_to_cpu
         self.flatten_parameters = flatten_parameters
@@ -169,12 +174,6 @@ def model_to_device(self) -> None:
         # ensure we update the device type in the lightning module
         self.lightning_module.to(self.root_device)
 
-    def lightning_module_state_dict(self) -> Dict[str, Union[Any, Tensor]]:
-        # Currently it is same as default TrainingTypePlugin, i.e. return
-        # the full state dict for FSDP, in the future, we will provide sharded
-        # state dict.
-        return super().lightning_module_state_dict()
-
     @property
     def setup_optimizers_in_pre_dispatch(self) -> bool:
         # Setup optimizers after the Fully Sharded Model has been made

pytorch_lightning/plugins/training_type/horovod.py

@@ -20,6 +20,7 @@
 from torch.optim.lr_scheduler import _LRScheduler
 
 from pytorch_lightning.core.optimizer import LightningOptimizer
+from pytorch_lightning.plugins.checkpoint.checkpoint import CheckpointIOPlugin
 from pytorch_lightning.plugins.training_type.parallel import ParallelPlugin
 from pytorch_lightning.utilities import _HOROVOD_AVAILABLE
 from pytorch_lightning.utilities.distributed import distributed_available
@@ -33,8 +34,14 @@
 class HorovodPlugin(ParallelPlugin):
     """Plugin for Horovod distributed training integration."""
 
-    def __init__(self, parallel_devices: Optional[List[torch.device]] = None):
-        super().__init__(parallel_devices=parallel_devices, cluster_environment=None)
+    def __init__(
+        self,
+        parallel_devices: Optional[List[torch.device]] = None,
+        checkpoint_plugin: Optional[CheckpointIOPlugin] = None,
+    ):
+        super().__init__(
+            parallel_devices=parallel_devices, cluster_environment=None, checkpoint_plugin=checkpoint_plugin
+        )
         rank_zero_only.rank = self.global_rank
 
     @property

pytorch_lightning/plugins/training_type/ipu.py

@@ -22,6 +22,7 @@
 import pytorch_lightning as pl
 from pytorch_lightning.callbacks import GradientAccumulationScheduler
 from pytorch_lightning.overrides.base import _LightningModuleWrapperBase
+from pytorch_lightning.plugins.checkpoint.checkpoint import CheckpointIOPlugin
 from pytorch_lightning.plugins.environments.cluster_environment import ClusterEnvironment
 from pytorch_lightning.plugins.training_type.parallel import ParallelPlugin
 from pytorch_lightning.trainer.states import RunningStage
@@ -67,6 +68,7 @@ def __init__(
         autoreport_dir: Optional[str] = None,
         parallel_devices: Optional[List[torch.device]] = None,
         cluster_environment: Optional[ClusterEnvironment] = None,
+        checkpoint_plugin: Optional[CheckpointIOPlugin] = None,
         training_opts: Optional["poptorch.Options"] = None,
         inference_opts: Optional["poptorch.Options"] = None,
     ) -> None:
@@ -83,7 +85,11 @@ def __init__(
             inference_opts: Optional ``poptorch.Options`` to override the default
                 created options for validation/testing and predicting.
         """
-        super().__init__(parallel_devices, cluster_environment)
+        super().__init__(
+            parallel_devices=parallel_devices,
+            cluster_environment=cluster_environment,
+            checkpoint_plugin=checkpoint_plugin,
+        )
         if not _POPTORCH_AVAILABLE or not poptorch.ipuHardwareIsAvailable():
             raise MisconfigurationException(
                 "The IPU Accelerator requires IPU devices to run. "

pytorch_lightning/plugins/training_type/parallel.py

@@ -21,6 +21,7 @@
 
 import pytorch_lightning as pl
 from pytorch_lightning.overrides.base import unwrap_lightning_module
+from pytorch_lightning.plugins.checkpoint.checkpoint import CheckpointIOPlugin
 from pytorch_lightning.plugins.environments.cluster_environment import ClusterEnvironment
 from pytorch_lightning.plugins.training_type.training_type_plugin import TrainingTypePlugin
 from pytorch_lightning.utilities import _XLA_AVAILABLE
@@ -34,8 +35,9 @@ def __init__(
         self,
         parallel_devices: Optional[List[torch.device]] = None,
         cluster_environment: Optional[ClusterEnvironment] = None,
+        checkpoint_plugin: Optional[CheckpointIOPlugin] = None,
     ):
-        super().__init__()
+        super().__init__(checkpoint_plugin)
         self.parallel_devices = parallel_devices
         self.cluster_environment = cluster_environment
 

pytorch_lightning/plugins/training_type/single_device.py

@@ -15,15 +15,20 @@
 
 import torch
 
+from pytorch_lightning.plugins.checkpoint.checkpoint import CheckpointIOPlugin
 from pytorch_lightning.plugins.training_type.training_type_plugin import TrainingTypePlugin
 from pytorch_lightning.utilities import _XLA_AVAILABLE
 
 
 class SingleDevicePlugin(TrainingTypePlugin):
     """Plugin that handles communication on a single device."""
 
-    def __init__(self, device: torch.device):
-        super().__init__()
+    def __init__(
+        self,
+        device: torch.device,
+        checkpoint_plugin: Optional[CheckpointIOPlugin] = None,
+    ):
+        super().__init__(checkpoint_plugin)
         self.device: torch.device = device
         self.global_rank = 0
         self.local_rank = 0