Repository: Lightning-AI/pytorch-lightning
Pull request: Add support for async checkpointing
## What does this PR do?

<!--
Please include a summary of the change and which issue is fixed.
Please also include relevant motivation and context.
List any dependencies that are required for this change.

If we didn't discuss your PR in Github issues there's a high chance it will not be merged.

The following links the related issue to the PR (https://docs.github.com/en/free-pro-team@latest/github/managing-your-work-on-github/linking-a-pull-request-to-an-issue#linking-a-pull-request-to-an-issue-using-a-keyword)
-->

Fixes #11561
### Does your PR introduce any breaking changes? If yes, please list them.

<!-- FILL IN or None -->

## Before submitting

- [x] Was this **discussed/approved** via a GitHub issue? (not for typos and docs)
- [x] Did you read the [contributor guideline](https://github.com/Lightning-AI/lightning/blob/master/.github/CONTRIBUTING.md), **Pull Request** section?
- [x] Did you make sure your **PR does only one thing**, instead of bundling different changes together?
- [ ] Did you make sure to **update the documentation** with your changes? (if necessary)
- [ ] Did you write any **new necessary tests**? (not for typos and docs)
- [x] Did you verify new and **existing tests pass** locally with your changes?
- [x] Did you list all the **breaking changes** introduced by this pull request?
- [x] Did you **update the CHANGELOG**? (not for typos, docs, test updates, or minor internal changes/refactors)

<!-- In the CHANGELOG, separate each item in the unreleased section by a blank line to reduce collisions -->

## PR review

Anyone in the community is welcome to review the PR.
Before you start reviewing, make sure you have read the [review guidelines](https://github.com/Lightning-AI/lightning/wiki/Review-guidelines). In short, see the following bullet-list:

- [x] Is this pull request ready for review? (if not, please submit in draft mode)
- [x] Check that all items from **Before submitting** are resolved
- [x] Make sure the title is self-explanatory and the description concisely explains the PR
- [x] Add labels and milestones (and optionally projects) to the PR so it can be classified

## Did you have fun?

Make sure you had fun coding 🙃


cc @borda @awaelchli @ananthsub @ninginthecloud @rohitgr7 @otaj @akihironitta

Problem statement:
Asynchronous checkpointing
## 🚀 Feature

<!-- A clear and concise description of the feature proposal -->
The idea is to enable asynchronous model checkpointing during training.

### Motivation

I'm training thousands of small identical `gpytorch` models, each one is trained for 1000 epochs on a different 1d time series sample. The training on a GPU takes ~1/4 seconds of time. Saving the model at the end takes about 1.25 seconds, and this blocks the training of the next models and wastes GPU time.

### Pitch

Checkpointing currently blocks the training process and wastes GPU time. If implemented correctly, asynchronous (non-blocking) checkpointing could offer a significant speed up over time.

### Alternatives

Perhaps this doesn't need to be built into PyTorch-Lightning and can be implemented with asyncio.
If so, I would like to see an example.

### Additional context

______________________________________________________________________

cc @borda @awaelchli @ananthsub @ninginthecloud @rohitgr7

Hints from issue discussion:
I believe you would have a second process with a queue responsible to checkpoint the model weights. This could be done directly within a custom ModelCheckpoint.
When you say "saving the model" are you referring precisely to the IO operations to save and delete? If so, we could create a IO plugin that can work asynchronously, it could be adapted from https://github.com/PyTorchLightning/pytorch-lightning/blob/master/pytorch_lightning/plugins/io/torch_plugin.py
Yes, that's exactly what I mean.


Adding asynchronous checkpointing, like it is implemented in #13658, would have the implication that we can't guarantee access to the checkpoints anymore. For example:

```py
trainer.fit(model)
torch.load(trainer.checkpoint_callback.best_model_path)
```
Could lead to either

1) the wrong model being loaded
2) FileNotFoundError

due to the queue being processed in a side thread. This is of course a rare edge case but can't be ignored imo. One has to provide a mechanism to synchronize or lock when accessing the checkpoint this way. 
I guess 
https://github.com/Lightning-AI/lightning/blob/2bdd8e9da838d14ef8d965debe07109462edbf47/src/pytorch_lightning/utilities/cloud_io.py#L49-L52

ensures all the checkpoints are processed before we end the training?
I thought it is supposed to address the use case of @noamsgl? Perhaps the use case here can't be fully addressed.

Nonetheless, we need to be extremely careful about race conditions. The PR does not currently implement locks between checkpoint creation and deletion. This is necessary, otherwise you end up with a main process trying to remove a checkpoint while the thread is trying to overwrite the file. 

Furthermore, there are places where we can't have asynchronous checkpointing. One example is the spawn launcher who saves a temporary checkpoint before it joins in the main process. But with switching the plugin, this would have a global effect on all saving/loading mechanisms, which is not desired in this case. The specific case of spawn here could potentially be solved by passing the data through the queue directly instead of writing a checkpoint file, but that is a different discussion altogether. 
well, it partially addresses the use case here... the only place where the slowdown will happen is when training is done, to ensure all the models are saved before training ends.

> The PR does not currently implement locks between checkpoint creation and deletion. This is necessary, otherwise you end up with a main process trying to remove a checkpoint while the thread is trying to overwrite the file.

that's why added an info message.. this is entirely optional.. user needs to know the side effects.

> Furthermore, there are places where we can't have asynchronous checkpointing. One example is the spawn launcher who saves a temporary checkpoint before it joins in the main process. But with switching the plugin, this would have a global effect on all saving/loading mechanisms, which is not desired in this case. The specific case of spawn here could potentially be solved by passing the data through the queue directly instead of writing a checkpoint file, but that is a different discussion altogether.

interesting.

Commit message: update docs

Patch to review:
diff --git a/docs/source-pytorch/api_references.rst b/docs/source-pytorch/api_references.rst
index 96a061a941b57..db4fc1e2c4cf8 100644
--- a/docs/source-pytorch/api_references.rst
+++ b/docs/source-pytorch/api_references.rst
@@ -210,6 +210,7 @@ io
     :nosignatures:
     :template: classtemplate.rst
 
+    AsyncCheckpointIO
     CheckpointIO
     HPUCheckpointIO
     TorchCheckpointIO
diff --git a/docs/source-pytorch/common/checkpointing_expert.rst b/docs/source-pytorch/common/checkpointing_expert.rst
index c4a948a34cb9d..09af39e3a5969 100644
--- a/docs/source-pytorch/common/checkpointing_expert.rst
+++ b/docs/source-pytorch/common/checkpointing_expert.rst
@@ -45,6 +45,10 @@ Built-in Checkpoint IO Plugins
        respectively, common for most use cases.
    * - :class:`~pytorch_lightning.plugins.io.XLACheckpointIO`
      - CheckpointIO that utilizes :func:`xm.save` to save checkpoints for TPU training strategies.
+   * - :class:`~pytorch_lightning.plugins.io.HPUCheckpointIO`
+     - CheckpointIO to save checkpoints for HPU training strategies.
+   * - :class:`~pytorch_lightning.plugins.io.AsyncCheckpointIO`
+     - ``AsyncCheckpointIO`` enables saving the checkpoints asynchronously using the ``ThreadPoolExecutor``.
 
 
 ***************************
@@ -94,3 +98,41 @@ Custom Checkpoint IO Plugin
 .. note::
 
     Some ``TrainingTypePlugins`` like ``DeepSpeedStrategy`` do not support custom ``CheckpointIO`` as checkpointing logic is not modifiable.
+
+
+**************************
+Asynchronous Checkpointing
+**************************
+
+.. warning::
+
+    This is currently an experimental plugin/feature and API changes are to be expected.
+
+To enable saving the checkpoints asynchronously without blocking your training, you can configure
+:class:`~pytorch_lightning.plugins.io.async_plugin.AsyncCheckpointIO` plugin to ``Trainer``.
+
+It uses its base ``CheckpointIO`` plugin's saving logic to save the checkpoint but performs this operation asynchronously.
+By default, this base ``CheckpointIO`` will be set-up for you and all you need to provide is the ``AsyncCheckpointIO`` instance to the ``Trainer``.
+
+.. code-block:: python
+
+   from pytorch_lightning.plugins.io import AsyncCheckpointIO
+
+
+   async_ckpt_io = AsyncCheckpointIO()
+   trainer = Trainer(plugins=[async_ckpt_io])
+   model = LitModel()
+   trainer.fit(model)
+
+
+If you want the plugin to use your own base ``CheckpointIO`` and want the base to behave asynchronously, pass it as an argument while initializing ``AsyncCheckpointIO``.
+
+.. code-block:: python
+
+   from pytorch_lightning.plugins.io import AsyncCheckpointIO
+
+   base_ckpt_io = MyCustomCheckpointIO()
+   async_ckpt_io = AsyncCheckpointIO(checkpoint_io=base_ckpt_io)
+   trainer = Trainer(plugins=[async_ckpt_io])
+   model = LitModel()
+   trainer.fit(model)
diff --git a/docs/source-pytorch/extensions/plugins.rst b/docs/source-pytorch/extensions/plugins.rst
index 6ea8d42815f46..a0dbefd141464 100644
--- a/docs/source-pytorch/extensions/plugins.rst
+++ b/docs/source-pytorch/extensions/plugins.rst
@@ -87,6 +87,7 @@ Below is a list of built-in plugins for checkpointing.
     :nosignatures:
     :template: classtemplate.rst
 
+    AsyncCheckpointIO
     CheckpointIO
     HPUCheckpointIO
     TorchCheckpointIO
diff --git a/src/pytorch_lightning/CHANGELOG.md b/src/pytorch_lightning/CHANGELOG.md
index b30e6b57203f3..8921f4f1d8534 100644
--- a/src/pytorch_lightning/CHANGELOG.md
+++ b/src/pytorch_lightning/CHANGELOG.md
@@ -105,6 +105,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Added support for DDP Fork ([#13405](https://github.com/PyTorchLightning/pytorch-lightning/pull/13405))
 
 
+- Added support for async checkpointing ([#13658](https://github.com/PyTorchLightning/pytorch-lightning/pull/13658))
+
+
 ### Changed
 
 - Enable validation during overfitting ([#12527](https://github.com/PyTorchLightning/pytorch-lightning/pull/12527))
diff --git a/src/pytorch_lightning/plugins/__init__.py b/src/pytorch_lightning/plugins/__init__.py
index 0f1c4ca85ed5a..afd10c88c951d 100644
--- a/src/pytorch_lightning/plugins/__init__.py
+++ b/src/pytorch_lightning/plugins/__init__.py
@@ -1,6 +1,7 @@
 from typing import Union
 
 from pytorch_lightning.plugins.environments import ClusterEnvironment
+from pytorch_lightning.plugins.io.async_plugin import AsyncCheckpointIO
 from pytorch_lightning.plugins.io.checkpoint_plugin import CheckpointIO
 from pytorch_lightning.plugins.io.hpu_plugin import HPUCheckpointIO
 from pytorch_lightning.plugins.io.torch_plugin import TorchCheckpointIO
@@ -38,6 +39,7 @@
 PLUGIN_INPUT = Union[PLUGIN, str]
 
 __all__ = [
+    "AsyncCheckpointIO",
     "CheckpointIO",
     "TorchCheckpointIO",
     "XLACheckpointIO",
diff --git a/src/pytorch_lightning/plugins/io/__init__.py b/src/pytorch_lightning/plugins/io/__init__.py
index abd196eb2b1e3..19a556bddf29c 100644
--- a/src/pytorch_lightning/plugins/io/__init__.py
+++ b/src/pytorch_lightning/plugins/io/__init__.py
@@ -11,7 +11,10 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from pytorch_lightning.plugins.io.checkpoint_plugin import CheckpointIO  # noqa: F401
-from pytorch_lightning.plugins.io.hpu_plugin import HPUCheckpointIO  # noqa: F401
-from pytorch_lightning.plugins.io.torch_plugin import TorchCheckpointIO  # noqa: F401
-from pytorch_lightning.plugins.io.xla_plugin import XLACheckpointIO  # noqa: F401
+from pytorch_lightning.plugins.io.async_plugin import AsyncCheckpointIO
+from pytorch_lightning.plugins.io.checkpoint_plugin import CheckpointIO
+from pytorch_lightning.plugins.io.hpu_plugin import HPUCheckpointIO
+from pytorch_lightning.plugins.io.torch_plugin import TorchCheckpointIO
+from pytorch_lightning.plugins.io.xla_plugin import XLACheckpointIO
+
+__all__ = ["AsyncCheckpointIO", "CheckpointIO", "HPUCheckpointIO", "TorchCheckpointIO", "XLACheckpointIO"]
diff --git a/src/pytorch_lightning/plugins/io/async_plugin.py b/src/pytorch_lightning/plugins/io/async_plugin.py
new file mode 100644
index 0000000000000..a280eac763f05
--- /dev/null
+++ b/src/pytorch_lightning/plugins/io/async_plugin.py
@@ -0,0 +1,62 @@
+# Copyright The PyTorch Lightning team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from concurrent.futures import ThreadPoolExecutor
+from typing import Any, Dict, Optional
+
+from pytorch_lightning.plugins.io.checkpoint_plugin import CheckpointIO
+from pytorch_lightning.plugins.io.wrapper import _WrappingCheckpointIO
+
+
+class AsyncCheckpointIO(_WrappingCheckpointIO):
+    """``AsyncCheckpointIO`` enables saving the checkpoints asynchronously using the ``ThreadPoolExecutor``.
+
+    .. warning::
+
+        This is currently an experimental plugin/feature and API changes are to be expected.
+
+    Args:
+        checkpoint_io: A checkpoint IO plugin that is used as the basis for async checkpointing.
+        interval: Sleep time between each queue check.
+    """
+
+    def __init__(self, checkpoint_io: Optional["CheckpointIO"] = None, interval: float = 2.0) -> None:
+        super().__init__(checkpoint_io)
+
+        self._executor = ThreadPoolExecutor(max_workers=1)
+        self._error = None
+
+    def save_checkpoint(self, *args: Any, **kwargs: Any) -> None:
+        """Uses the ``ThreadPoolExecutor`` to save the checkpoints using the base ``checkpoint_io``."""
+
+        def _save_checkpoint(*args, **kwargs):
+            try:
+                self.checkpoint_io.save_checkpoint(*args, **kwargs)
+            except Exception as e:
+                self._error = e
+
+        self._executor.submit(_save_checkpoint, *args, **kwargs)
+
+        if self._error:
+            raise self._error
+
+    def remove_checkpoint(self, *args: Any, **kwargs: Any) -> None:
+        super().remove_checkpoint(*args, **kwargs)
+
+    def load_checkpoint(self, *args: Any, **kwargs: Any) -> Dict[str, Any]:
+        return super().load_checkpoint(*args, **kwargs)
+
+    def teardown(self) -> None:
+        """This method is called to close the threads."""
+        self._executor.shutdown(wait=True)
diff --git a/src/pytorch_lightning/plugins/io/checkpoint_plugin.py b/src/pytorch_lightning/plugins/io/checkpoint_plugin.py
index 1425a229963b7..26432614a9373 100644
--- a/src/pytorch_lightning/plugins/io/checkpoint_plugin.py
+++ b/src/pytorch_lightning/plugins/io/checkpoint_plugin.py
@@ -43,12 +43,13 @@ def save_checkpoint(self, checkpoint: Dict[str, Any], path: _PATH, storage_optio
         """
 
     @abstractmethod
-    def load_checkpoint(self, path: _PATH, storage_options: Optional[Any] = None) -> Dict[str, Any]:
+    def load_checkpoint(self, path: _PATH, map_location: Optional[Any] = 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.
+            map_location: a function, :class:`torch.device`, string or a dict specifying how to remap storage
+            locations.
 
         Returns: The loaded checkpoint.
         """
@@ -60,3 +61,7 @@ def remove_checkpoint(self, path: _PATH) -> None:
         Args:
             path: Path to checkpoint
         """
+
+    def teardown(self) -> None:
+        """This method is called to teardown the process."""
+        pass
diff --git a/src/pytorch_lightning/plugins/io/wrapper.py b/src/pytorch_lightning/plugins/io/wrapper.py
new file mode 100644
index 0000000000000..4933237717cdd
--- /dev/null
+++ b/src/pytorch_lightning/plugins/io/wrapper.py
@@ -0,0 +1,68 @@
+# Copyright The PyTorch Lightning team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from typing import Any, Dict, Optional
+
+from pytorch_lightning.plugins.io.checkpoint_plugin import CheckpointIO
+
+
+class _WrappingCheckpointIO(CheckpointIO):
+    """AsyncCheckpointIO enables saving the checkpoints asynchronously.
+
+    .. warning::
+
+        This is currently an experimental plugin/feature and API changes are to be expected.
+
+    Args:
+        checkpoint_io: A checkpoint IO plugin that is used as the basis for async checkpointing.
+        interval: Sleep time between each queue check.
+    """
+
+    def __init__(self, checkpoint_io: Optional["CheckpointIO"] = None) -> None:
+        super().__init__()
+
+        self._checkpoint_io = checkpoint_io
+        self._base_checkpoint_io_configured: bool = False
+
+        if checkpoint_io is not None:
+            if isinstance(checkpoint_io, _WrappingCheckpointIO):
+                self._base_checkpoint_io_configured = checkpoint_io._base_checkpoint_io_configured
+            else:
+                self._base_checkpoint_io_configured = True
+
+    @property
+    def checkpoint_io(self):
+        return self._checkpoint_io
+
+    @checkpoint_io.setter
+    def checkpoint_io(self, checkpoint_io: "CheckpointIO") -> None:
+        assert not isinstance(checkpoint_io, _WrappingCheckpointIO)
+
+        if self._checkpoint_io is None:
+            self._base_checkpoint_io_configured = True
+            self._checkpoint_io = checkpoint_io
+        elif isinstance(self._checkpoint_io, _WrappingCheckpointIO) and not self._base_checkpoint_io_configured:
+            self._base_checkpoint_io_configured = True
+            self._checkpoint_io.checkpoint_io = checkpoint_io
+
+    def save_checkpoint(self, *args: Any, **kwargs: Any) -> None:
+        """Uses the base ``checkpoint_io`` to save the checkpoint."""
+        self.checkpoint_io.save_checkpoint(*args, **kwargs)
+
+    def remove_checkpoint(self, *args, **kwargs) -> None:
+        """Uses the base ``checkpoint_io`` to remove the checkpoint."""
+        self.checkpoint_io.remove_checkpoint(*args, **kwargs)
+
+    def load_checkpoint(self, *args: Any, **kwargs: Any) -> Dict[str, Any]:
+        """Uses the base ``checkpoint_io`` to load the checkpoint."""
+        return self.checkpoint_io.load_checkpoint(*args, **kwargs)
diff --git a/src/pytorch_lightning/strategies/hpu_parallel.py b/src/pytorch_lightning/strategies/hpu_parallel.py
index 591664e93e782..90c0a4e7cf967 100644
--- a/src/pytorch_lightning/strategies/hpu_parallel.py
+++ b/src/pytorch_lightning/strategies/hpu_parallel.py
@@ -23,6 +23,7 @@
 from pytorch_lightning.plugins.environments.cluster_environment import ClusterEnvironment
 from pytorch_lightning.plugins.io.checkpoint_plugin import CheckpointIO
 from pytorch_lightning.plugins.io.hpu_plugin import HPUCheckpointIO
+from pytorch_lightning.plugins.io.wrapper import _WrappingCheckpointIO
 from pytorch_lightning.plugins.precision import PrecisionPlugin
 from pytorch_lightning.strategies.ddp import DDPStrategy
 from pytorch_lightning.utilities.distributed import group as _group
diff --git a/src/pytorch_lightning/strategies/single_hpu.py b/src/pytorch_lightning/strategies/single_hpu.py
index bbba3904f6bc3..f2ab38402140f 100644
--- a/src/pytorch_lightning/strategies/single_hpu.py
+++ b/src/pytorch_lightning/strategies/single_hpu.py
@@ -17,6 +17,7 @@
 import pytorch_lightning as pl
 from pytorch_lightning.plugins.io.checkpoint_plugin import CheckpointIO
 from pytorch_lightning.plugins.io.hpu_plugin import HPUCheckpointIO
+from pytorch_lightning.plugins.io.wrapper import _WrappingCheckpointIO
 from pytorch_lightning.plugins.precision import PrecisionPlugin
 from pytorch_lightning.strategies.single_device import SingleDeviceStrategy
 from pytorch_lightning.utilities import _HPU_AVAILABLE
diff --git a/src/pytorch_lightning/strategies/single_tpu.py b/src/pytorch_lightning/strategies/single_tpu.py
index caf153ace0bb6..890f94993e23c 100644
--- a/src/pytorch_lightning/strategies/single_tpu.py
+++ b/src/pytorch_lightning/strategies/single_tpu.py
@@ -16,6 +16,7 @@
 
 import pytorch_lightning as pl
 from pytorch_lightning.plugins.io.checkpoint_plugin import CheckpointIO
+from pytorch_lightning.plugins.io.wrapper import _WrappingCheckpointIO
 from pytorch_lightning.plugins.io.xla_plugin import XLACheckpointIO
 from pytorch_lightning.plugins.precision import PrecisionPlugin
 from pytorch_lightning.strategies.single_device import SingleDeviceStrategy
diff --git a/src/pytorch_lightning/strategies/strategy.py b/src/pytorch_lightning/strategies/strategy.py
index 01b8f1b793791..91799c528245a 100644
--- a/src/pytorch_lightning/strategies/strategy.py
+++ b/src/pytorch_lightning/strategies/strategy.py
@@ -27,6 +27,7 @@
 from pytorch_lightning.overrides.base import unwrap_lightning_module
 from pytorch_lightning.plugins import TorchCheckpointIO
 from pytorch_lightning.plugins.io.checkpoint_plugin import CheckpointIO
+from pytorch_lightning.plugins.io.wrapper import _WrappingCheckpointIO
 from pytorch_lightning.plugins.precision import PrecisionPlugin
 from pytorch_lightning.strategies.launchers.base import _Launcher
 from pytorch_lightning.trainer.states import TrainerFn
@@ -50,6 +51,7 @@ def __init__(
         checkpoint_io: Optional[CheckpointIO] = None,
         precision_plugin: Optional[PrecisionPlugin] = None,
     ) -> None:
+
         self.accelerator = accelerator
         self._launcher: Optional[_Launcher] = None
         self._model: Optional[Module] = None
@@ -76,6 +78,8 @@ def accelerator(self, accelerator: "pl.accelerators.accelerator.Accelerator") ->
     def checkpoint_io(self) -> CheckpointIO:
         if self._checkpoint_io is None:
             self._checkpoint_io = TorchCheckpointIO()
+        elif isinstance(self._checkpoint_io, _WrappingCheckpointIO):
+            self._checkpoint_io.checkpoint_io = TorchCheckpointIO()
 
         return self._checkpoint_io
 
@@ -448,6 +452,7 @@ def teardown(self) -> None:
             self.lightning_module.cpu()
         self.precision_plugin.teardown()
         self.accelerator.teardown()
+        self.checkpoint_io.teardown()
 
     @classmethod
     def register_strategies(cls, strategy_registry) -> None:
diff --git a/src/pytorch_lightning/strategies/tpu_spawn.py b/src/pytorch_lightning/strategies/tpu_spawn.py
index 1f7effcf75008..a2f04f3090796 100644
--- a/src/pytorch_lightning/strategies/tpu_spawn.py
+++ b/src/pytorch_lightning/strategies/tpu_spawn.py
@@ -24,6 +24,7 @@
 from pytorch_lightning.overrides import LightningDistributedModule
 from pytorch_lightning.plugins.environments import XLAEnvironment
 from pytorch_lightning.plugins.io.checkpoint_plugin import CheckpointIO
+from pytorch_lightning.plugins.io.wrapper import _WrappingCheckpointIO
 from pytorch_lightning.plugins.io.xla_plugin import XLACheckpointIO
 from pytorch_lightning.plugins.precision import PrecisionPlugin
 from pytorch_lightning.strategies.ddp_spawn import DDPSpawnStrategy
diff --git a/src/pytorch_lightning/utilities/cloud_io.py b/src/pytorch_lightning/utilities/cloud_io.py
index 9055ff50c8d0d..81482a8ab24f9 100644
--- a/src/pytorch_lightning/utilities/cloud_io.py
+++ b/src/pytorch_lightning/utilities/cloud_io.py
@@ -62,7 +62,6 @@ def atomic_save(checkpoint: Dict[str, Any], filepath: Union[str, Path]) -> None:
         filepath: The path to which the checkpoint will be saved.
             This points to the file that the checkpoint will be stored in.
     """
-
     bytesbuffer = io.BytesIO()
     torch.save(checkpoint, bytesbuffer)
     with fsspec.open(filepath, "wb") as f:
diff --git a/tests/tests_pytorch/plugins/test_checkpoint_io_plugin.py b/tests/tests_pytorch/plugins/test_checkpoint_io_plugin.py
index 651ab1cc4f49f..0a201d982c15b 100644
--- a/tests/tests_pytorch/plugins/test_checkpoint_io_plugin.py
+++ b/tests/tests_pytorch/plugins/test_checkpoint_io_plugin.py
@@ -12,15 +12,18 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import os
+from pathlib import Path
 from typing import Any, Dict, Optional
-from unittest.mock import MagicMock
+from unittest.mock import MagicMock, Mock
 
 import torch
 
 from pytorch_lightning import Trainer
 from pytorch_lightning.callbacks import ModelCheckpoint
 from pytorch_lightning.demos.boring_classes import BoringModel
-from pytorch_lightning.plugins import CheckpointIO
+from pytorch_lightning.plugins.io.async_plugin import AsyncCheckpointIO
+from pytorch_lightning.plugins.io.checkpoint_plugin import CheckpointIO
+from pytorch_lightning.plugins.io.torch_plugin import TorchCheckpointIO
 from pytorch_lightning.strategies import SingleDeviceStrategy
 from pytorch_lightning.utilities.types import _PATH
 
@@ -49,9 +52,16 @@ def test_checkpoint_plugin_called(tmpdir):
         strategy=SingleDeviceStrategy("cpu", checkpoint_io=checkpoint_plugin),
         callbacks=ck,
         max_epochs=2,
+        limit_train_batches=1,
+        limit_val_batches=0,
+        limit_test_batches=1,
     )
     trainer.fit(model)
 
+    ckpt_files = {fn.name for fn in Path(tmpdir).glob("*.ckpt")}
+    assert ckpt_files == {"epoch=1-step=2.ckpt", "last.ckpt"}
+    assert trainer.checkpoint_callback.best_model_path == tmpdir / "epoch=1-step=2.ckpt"
+    assert trainer.checkpoint_callback.last_model_path == tmpdir / "last.ckpt"
     assert checkpoint_plugin.save_checkpoint.call_count == 4
     assert checkpoint_plugin.remove_checkpoint.call_count == 1
 
@@ -68,12 +78,79 @@ def test_checkpoint_plugin_called(tmpdir):
         plugins=[checkpoint_plugin],
         callbacks=ck,
         max_epochs=2,
+        limit_train_batches=1,
+        limit_val_batches=0,
+        limit_test_batches=1,
     )
     trainer.fit(model)
 
+    ckpt_files = {fn.name for fn in Path(tmpdir).glob("*.ckpt")}
+    assert ckpt_files == {"epoch=1-step=2.ckpt", "last.ckpt", "epoch=1-step=2-v1.ckpt", "last-v1.ckpt"}
+    assert trainer.checkpoint_callback.best_model_path == tmpdir / "epoch=1-step=2-v1.ckpt"
+    assert trainer.checkpoint_callback.last_model_path == tmpdir / "last-v1.ckpt"
     assert checkpoint_plugin.save_checkpoint.call_count == 4
     assert checkpoint_plugin.remove_checkpoint.call_count == 1
 
     trainer.test(model, ckpt_path=ck.last_model_path)
     checkpoint_plugin.load_checkpoint.assert_called_once()
     checkpoint_plugin.load_checkpoint.assert_called_with(tmpdir / "last-v1.ckpt")
+
+
+def test_async_checkpoint_plugin(tmpdir):
+    """Ensure that the custom checkpoint IO plugin and torch checkpoint IO plugin is called when async saving and
+    loading."""
+
+    checkpoint_plugin = AsyncCheckpointIO()
+
+    checkpoint_plugin.save_checkpoint = Mock(wraps=checkpoint_plugin.save_checkpoint)
+    checkpoint_plugin.remove_checkpoint = Mock(wraps=checkpoint_plugin.remove_checkpoint)
+
+    class CustomBoringModel(BoringModel):
+        def on_fit_start(self):
+            base_ckpt_io = self.trainer.strategy.checkpoint_io.checkpoint_io
+            base_ckpt_io.save_checkpoint = Mock(wraps=base_ckpt_io.save_checkpoint)
+            base_ckpt_io.remove_checkpoint = Mock(wraps=base_ckpt_io.remove_checkpoint)
+
+    ck = ModelCheckpoint(dirpath=tmpdir, save_top_k=2, monitor="step", mode="max")
+
+    model = CustomBoringModel()
+    trainer = Trainer(
+        default_root_dir=tmpdir,
+        plugins=[checkpoint_plugin],
+        callbacks=ck,
+        max_epochs=3,
+        limit_train_batches=1,
+        limit_val_batches=0,
+        enable_progress_bar=False,
+        enable_model_summary=False,
+    )
+    trainer.fit(model)
+
+    assert checkpoint_plugin.save_checkpoint.call_count == 3
+    assert checkpoint_plugin.remove_checkpoint.call_count == 1
+
+    ckpt_files = {fn.name for fn in Path(tmpdir).glob("*.ckpt")}
+    assert ckpt_files == {"epoch=1-step=2.ckpt", "epoch=2-step=3.ckpt"}
+
+    base_ckpt_io = trainer.strategy.checkpoint_io.checkpoint_io
+    assert base_ckpt_io.save_checkpoint.call_count == 3
+    assert base_ckpt_io.remove_checkpoint.call_count == 1
+
+
+def test_multi_wrapped_checkpoint_io_initialization():
+    base_ckpt_io = TorchCheckpointIO()
+    wrap_ckpt = AsyncCheckpointIO(base_ckpt_io)
+    ckpt_io = AsyncCheckpointIO(wrap_ckpt)
+    assert ckpt_io.checkpoint_io is wrap_ckpt
+    assert ckpt_io.checkpoint_io.checkpoint_io is base_ckpt_io
+    assert ckpt_io._base_checkpoint_io_configured is True
+    assert ckpt_io.checkpoint_io._base_checkpoint_io_configured is True
+
+    wrap_ckpt = AsyncCheckpointIO()
+    ckpt_io = AsyncCheckpointIO(wrap_ckpt)
+    trainer = Trainer(accelerator="cpu", plugins=[ckpt_io])
+    trainer.strategy.checkpoint_io
+    assert ckpt_io.checkpoint_io is wrap_ckpt
+    assert isinstance(ckpt_io.checkpoint_io.checkpoint_io, TorchCheckpointIO)
+    assert ckpt_io._base_checkpoint_io_configured is True
+    assert ckpt_io.checkpoint_io._base_checkpoint_io_configured is True


Code review comments:
diff_hunk: @@ -0,0 +1,68 @@
+# Copyright The PyTorch Lightning team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from typing import Any, Dict, Optional
+
+from pytorch_lightning.plugins.io.checkpoint_plugin import CheckpointIO
+
+
+class _WrappingCheckpointIO(CheckpointIO):
+    """AsyncCheckpointIO enables saving the checkpoints asynchronously.
path: src/pytorch_lightning/plugins/io/wrapper.py
text: @user1:
This is copy-pasted

diff_hunk: @@ -0,0 +1,62 @@
+# Copyright The PyTorch Lightning team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from concurrent.futures import ThreadPoolExecutor
+from typing import Any, Dict, Optional
+
+from pytorch_lightning.plugins.io.checkpoint_plugin import CheckpointIO
+from pytorch_lightning.plugins.io.wrapper import _WrappingCheckpointIO
+
+
+class AsyncCheckpointIO(_WrappingCheckpointIO):
+    """``AsyncCheckpointIO`` enables saving the checkpoints asynchronously using the ``ThreadPoolExecutor``.
+
+    .. warning::
+
+        This is currently an experimental plugin/feature and API changes are to be expected.
+
+    Args:
+        checkpoint_io: A checkpoint IO plugin that is used as the basis for async checkpointing.
+        interval: Sleep time between each queue check.
+    """
+
+    def __init__(self, checkpoint_io: Optional["CheckpointIO"] = None, interval: float = 2.0) -> None:
+        super().__init__(checkpoint_io)
+
+        self._executor = ThreadPoolExecutor(max_workers=1)
+        self._error = None
path: src/pytorch_lightning/plugins/io/async_plugin.py
text: @user1:
Doesn't this need typing?

diff_hunk: @@ -0,0 +1,62 @@
+# Copyright The PyTorch Lightning team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from concurrent.futures import ThreadPoolExecutor
+from typing import Any, Dict, Optional
+
+from pytorch_lightning.plugins.io.checkpoint_plugin import CheckpointIO
+from pytorch_lightning.plugins.io.wrapper import _WrappingCheckpointIO
+
+
+class AsyncCheckpointIO(_WrappingCheckpointIO):
+    """``AsyncCheckpointIO`` enables saving the checkpoints asynchronously using the ``ThreadPoolExecutor``.
+
+    .. warning::
+
+        This is currently an experimental plugin/feature and API changes are to be expected.
+
+    Args:
+        checkpoint_io: A checkpoint IO plugin that is used as the basis for async checkpointing.
+        interval: Sleep time between each queue check.
+    """
+
+    def __init__(self, checkpoint_io: Optional["CheckpointIO"] = None, interval: float = 2.0) -> None:
+        super().__init__(checkpoint_io)
+
+        self._executor = ThreadPoolExecutor(max_workers=1)
path: src/pytorch_lightning/plugins/io/async_plugin.py
text: @user1:
Should we expose `max_workers` in the `__init__`?

diff_hunk: @@ -60,3 +61,7 @@ def remove_checkpoint(self, path: _PATH) -> None:
         Args:
             path: Path to checkpoint
         """
+
+    def teardown(self) -> None:
+        """This method is called to teardown the process."""
+        pass
path: src/pytorch_lightning/plugins/io/checkpoint_plugin.py
text: @user1:
Not necessary

```suggestion
```

diff_hunk: @@ -0,0 +1,62 @@
+# Copyright The PyTorch Lightning team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from concurrent.futures import ThreadPoolExecutor
+from typing import Any, Dict, Optional
+
+from pytorch_lightning.plugins.io.checkpoint_plugin import CheckpointIO
+from pytorch_lightning.plugins.io.wrapper import _WrappingCheckpointIO
+
+
+class AsyncCheckpointIO(_WrappingCheckpointIO):
+    """``AsyncCheckpointIO`` enables saving the checkpoints asynchronously using the ``ThreadPoolExecutor``.
+
+    .. warning::
+
+        This is currently an experimental plugin/feature and API changes are to be expected.
+
+    Args:
+        checkpoint_io: A checkpoint IO plugin that is used as the basis for async checkpointing.
+        interval: Sleep time between each queue check.
+    """
+
+    def __init__(self, checkpoint_io: Optional["CheckpointIO"] = None, interval: float = 2.0) -> None:
+        super().__init__(checkpoint_io)
+
+        self._executor = ThreadPoolExecutor(max_workers=1)
+        self._error = None
+
+    def save_checkpoint(self, *args: Any, **kwargs: Any) -> None:
+        """Uses the ``ThreadPoolExecutor`` to save the checkpoints using the base ``checkpoint_io``."""
+
+        def _save_checkpoint(*args, **kwargs):
+            try:
+                self.checkpoint_io.save_checkpoint(*args, **kwargs)
+            except Exception as e:
+                self._error = e
+
+        self._executor.submit(_save_checkpoint, *args, **kwargs)
+
+        if self._error:
+            raise self._error
+
+    def remove_checkpoint(self, *args: Any, **kwargs: Any) -> None:
+        super().remove_checkpoint(*args, **kwargs)
+
+    def load_checkpoint(self, *args: Any, **kwargs: Any) -> Dict[str, Any]:
+        return super().load_checkpoint(*args, **kwargs)
path: src/pytorch_lightning/plugins/io/async_plugin.py
text: @user1:
Why do you override these?