Repository: intelligent-machine-learning/dlrover (language: python)

Issue:
Discussion: how to name the APIs to async save/load checkpoints.
We are developing the features to asynchronously save the checkpoint into the storage and provide APIs to users.  Now, there are two schemes to name the APIs. We can select one scheme or offer a new scheme. Please vote your scheme in replies.

Scheme 1:
```
def save_checkpoint_to_memory():
    """ Synchronously save the checkpoint to the CPU memory."""
    pass

def save_checkpoint_to_storage():
    """ Asynchronously save the checkpoint to the storage."""
    pass

def load_checkpoint():
    """ 
    Load the checkpoint from the CPU memory if the state dict is in the memory.
    Otherwise, load the checkpoint from the storage.
    """
    pass
```

Scheme 2:
```
def save_checkpoint(in_memory=True/storage_level="memory_only/disk"):
    """ 
    If to_memory, synchronously save the checkpoint to the CPU memory.
    Otherwise, asynchronously save the checkpoint to the storage.
    """
    pass

def load_checkpoint():
    """ 
    Load the checkpoint from the CPU memory if the state dict is in the memory.
    Otherwise, load the checkpoint from the storage.
    """
    pass

```



Discussion / hints:
I perfer schem2 like the spark rdd operation (https://spark.apache.org/docs/latest/api/python/reference/api/pyspark.RDD.persist.html) :)
I prefer scheme 1
From "code smells": It is best not to use flag/boolean types as function parameters. Flags are not a best practice because they reduce the cohesion of a function. It does more than one thing making it hard to understand what a function does based on the call site.
If I want to modify the checkpoint storage mode, for scheme 1, I can simply replace the function name directly. For scheme 2, I may need to adjust the parameters, which may not be as straightforward.
> I perfer schem2 like the spark rdd operation (https://spark.apache.org/docs/latest/api/python/reference/api/pyspark.RDD.persist.html) :)

The RDD scheme is similar to the scheme 2 and I have updated the scheme 2.
I prefer scheme 2 since it hides details from users. Moreover, we should provide a reasonable default setup when users do not provide any configuration.
scheme 2
+1
I prefer scheme 2 since it looks more like a unified interface. When users do not specify which place to checkpoint their models, we make the choices for users to achieve higher performance **by default**.

Fix patch:
diff --git a/dlrover/trainer/torch/deepspeed/__init__.py b/dlrover/trainer/torch/deepspeed/__init__.py
deleted file mode 100644
index 561b39cf2..000000000
--- a/dlrover/trainer/torch/deepspeed/__init__.py
+++ /dev/null
@@ -1,12 +0,0 @@
-# Copyright 2023 The DLRover Authors. All rights reserved.
-# 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.
diff --git a/dlrover/trainer/torch/flash_checkpoint/checkpointer.py b/dlrover/trainer/torch/flash_checkpoint/checkpointer.py
new file mode 100644
index 000000000..874ef07bc
--- /dev/null
+++ b/dlrover/trainer/torch/flash_checkpoint/checkpointer.py
@@ -0,0 +1,65 @@
+# Copyright 2023 The DLRover Authors. All rights reserved.
+# 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 abc import ABCMeta, abstractmethod
+from enum import Enum, auto
+
+
+class StorageType(Enum):
+    MEMORY = auto()
+    DISK = auto()
+
+
+class Checkpointer(metaclass=ABCMeta):
+    """
+    Checkpointer can save and load PyTorch module states efficiently.
+
+    It begins by duplicating the state dictionary to shared memory,
+    then proceeds to save it to storage asynchronously. This process ensures
+    that the Checkpointer's save operation minimally blocks training time.
+    If the node does not fail, the Checkpointer prioritizes restoring the
+    checkpointed state dictionary directly from the shared memory upon load
+    requests. However, if the node has restarted, the Checkpointer reverts
+    to loading the state dictionary from the designated storage instead.
+    """
+
+    @abstractmethod
+    def save_checkpoint(
+        self, step, state_dict, path, storage_type=StorageType.DISK
+    ):
+        """
+        Save the checkpoint of model, optimizer and sampler.
+
+        Args:
+            step (int): the global iteration step.
+            state_dict (dict): the state dict of model and optimizer to save.
+            path (str): the storage path to save the state dict.
+                Note, the path is used to save the state dict to storage
+                only if the training process fails.
+            storage_type (StorageType): StorageType.MEMROY or StorageType.DISK.
+        """
+        pass
+
+    @abstractmethod
+    def load_checkpoint(self, resuming_path=None):
+        """
+        The manager loads the states from the files in the
+        checkpoint direcotry to the model, optimizer and sampler.
+        Args:
+            resuming_path (str, optinoal): The manager will load checkpoint
+                from the path. If the path is None, the manager will load
+                the state checkpoint from the file with the maximum step.
+        Return:
+            A dict: a state dict.
+        """
+        pass
diff --git a/dlrover/trainer/torch/deepspeed/async_checkpoint.py b/dlrover/trainer/torch/flash_checkpoint/deepspeed.py
similarity index 89%
rename from dlrover/trainer/torch/deepspeed/async_checkpoint.py
rename to dlrover/trainer/torch/flash_checkpoint/deepspeed.py
index a46f952ef..7cae9abaf 100644
--- a/dlrover/trainer/torch/deepspeed/async_checkpoint.py
+++ b/dlrover/trainer/torch/flash_checkpoint/deepspeed.py
@@ -28,6 +28,8 @@
     DeepSpeedCheckpointEngine,
 )
 
+from .checkpointer import StorageType
+
 
 class AsyncSaveEngine(CheckpointEngine):
     def __init__(self):
@@ -65,7 +67,7 @@ def commit(self, tag):
         pass
 
 
-class DeepSpeedCheckpointManger(object):
+class DeepSpeedCheckpointer(object):
     """
     The manager can synchronously save the DeepSpeedEngine checkpoint
     to the memory and asynchronously save the checkpointing states
@@ -77,7 +79,7 @@ class DeepSpeedCheckpointManger(object):
 
     Examples::
         >>> engine = deepspeed.initialize(...)
-        >>> ckpt_manager = DeepSpeedCheckpointManger(engine, save_dir)
+        >>> ckpt_manager = DeepSpeedCheckpointer(engine, save_dir)
         >>> if step % 10 == 0:
         >>>     ckpt_manager.save_checkpoint_to_memory(save_dir, tag)
         >>> if step % 100 == 0:
@@ -104,7 +106,26 @@ def __init__(self, engine: DeepSpeedEngine, checkpoint_dir):
         if zero_stage < ZeroStageEnum.weights and self._local_rank == 0:
             self.engine.save_non_zero_checkpoint = True
 
-    def save_checkpoint_to_memory(
+    def save_checkpoint(
+        self,
+        save_dir,
+        tag=None,
+        client_state={},
+        save_latest=True,
+        storage_type=StorageType.DISK,
+    ):
+        if storage_type == StorageType.MEMORY:
+            self._save_checkpoint_to_memory(
+                save_dir, tag, client_state, save_latest
+            )
+        elif storage_type == StorageType.DISK:
+            self._save_checkpoint_to_storage(
+                save_dir, tag, client_state, save_latest
+            )
+        else:
+            raise ValueError(f"No support storage type {storage_type}")
+
+    def _save_checkpoint_to_memory(
         self, save_dir, tag=None, client_state={}, save_latest=True
     ):
         torch_save_func = torch.save
@@ -128,7 +149,7 @@ def save_checkpoint_to_memory(
             except Exception:
                 pass
 
-    def save_checkpoint_to_storage(
+    def _save_checkpoint_to_storage(
         self, save_dir, tag=None, client_state={}, save_latest=True
     ):
         torch_save_func = torch.save
diff --git a/dlrover/trainer/torch/flash_checkpoint/fsdp.py b/dlrover/trainer/torch/flash_checkpoint/fsdp.py
new file mode 100644
index 000000000..97a05a194
--- /dev/null
+++ b/dlrover/trainer/torch/flash_checkpoint/fsdp.py
@@ -0,0 +1,38 @@
+# Copyright 2023 The DLRover Authors. All rights reserved.
+# 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 .checkpointer import Checkpointer, StorageType
+from .fsdp_engine import FsdpCheckpointEngine
+
+
+class FsdpCheckpointer(Checkpointer):
+    def __init__(self, checkpoint_dir: str):
+        self._engine = FsdpCheckpointEngine(checkpoint_dir)
+
+    def save_checkpoint(
+        self, step, state_dict, path, storage_type=StorageType.DISK
+    ):
+        if storage_type == StorageType.MEMORY:
+            self._engine.save_to_memory(step, state_dict, path)
+        elif storage_type == StorageType.DISK:
+            if not path:
+                raise ValueError(
+                    "path cannot be empty if storage type is disk!"
+                )
+            self._engine.save_to_storage(step, state_dict, path)
+        else:
+            raise ValueError(f"No support storage type {storage_type}")
+
+    def load_checkpoint(self, resume_path=""):
+        return self._engine.load(resume_path)
diff --git a/dlrover/trainer/torch/megatron/async_checkpoint.py b/dlrover/trainer/torch/flash_checkpoint/megatron.py
similarity index 62%
rename from dlrover/trainer/torch/megatron/async_checkpoint.py
rename to dlrover/trainer/torch/flash_checkpoint/megatron.py
index d4f20c6c8..334698a99 100644
--- a/dlrover/trainer/torch/megatron/async_checkpoint.py
+++ b/dlrover/trainer/torch/flash_checkpoint/megatron.py
@@ -22,18 +22,16 @@
 
 try:
     from megatron import get_args
-    from megatron.checkpointing import (
-        get_checkpoint_tracker_filename,
-        load_checkpoint,
-        save_checkpoint,
-    )
+    from megatron.checkpointing import get_checkpoint_tracker_filename
+    from megatron.checkpointing import load_checkpoint as megatron_load
+    from megatron.checkpointing import save_checkpoint as megatron_save
 except ImportError:
     logger.warning("Please check the magatron.checkpointing exists.")
 
 from dlrover.python.common.singleton import singleton
-from dlrover.trainer.torch.flash_checkpoint.megatron_engine import (
-    MegatronCheckpointEngine,
-)
+
+from .checkpointer import StorageType
+from .megatron_engine import MegatronCheckpointEngine
 
 
 @singleton
@@ -75,27 +73,12 @@ def update_latest_checkpoint_step(self):
             self._latest_ckpt_iteration = iteration
 
 
-def save_checkpoint_to_storage(
-    iteration, model, optimizer, opt_param_scheduler
-):
-    """
-    Asynchronously save the the checkpointing state dict into the storage.
-    The method will not wait for saving the checkpointing to the storage.
-
-    Args:
-        same as the `megatron.checkpointing.load_checkpoint`
-    """
-    args = get_args()
-    saver = MegatronCheckpointManager(args.save)
-    torch_save_func = torch.save
-    torch.save = saver.save
-    save_checkpoint(iteration, model, optimizer, opt_param_scheduler)
-    saver.engine.save_to_storage(iteration, saver.state_dict, saver.path)
-    torch.save = torch_save_func
-
-
-def save_checkpoint_to_memory(
-    iteration, model, optimizer, opt_param_scheduler
+def save_checkpoint(
+    iteration,
+    model,
+    optimizer,
+    opt_param_scheduler,
+    storage_type=StorageType.DISK,
 ):
     """
     Synchronously save the the checkpointing state dict into the CPU memory.
@@ -103,23 +86,34 @@ def save_checkpoint_to_memory(
     Args:
         same as the `megatron.checkpointing.load_checkpoint`
     """
-    args = get_args()
-    saver = MegatronCheckpointManager(args.save)
-    saver.update_latest_checkpoint_step()
-    torch_save_func = torch.save
-    torch.save = saver.save
-    save_checkpoint(iteration, model, optimizer, opt_param_scheduler)
-    saver.engine.save_to_memory(iteration, saver.state_dict, saver.path)
-    torch.save = torch_save_func
-
-    # Megatron save_checkpoint will create the directory with the iteration
-    # and write the iteration into the tracerfile. But async saver only
-    # save the state dict into the CPU memory not the storage. The saver
-    # need to clear the empty checkpoint directory.
-    saver.clear_empty_checkpoint(iteration)
-
-
-def load_latest_checkpoint(
+    if storage_type == StorageType.MEMORY:
+        args = get_args()
+        saver = MegatronCheckpointManager(args.save)
+        saver.update_latest_checkpoint_step()
+        torch_save_func = torch.save
+        torch.save = saver.save
+        megatron_save(iteration, model, optimizer, opt_param_scheduler)
+        saver.engine.save_to_memory(iteration, saver.state_dict, saver.path)
+        torch.save = torch_save_func
+
+        # Megatron save_checkpoint will create the directory with the iteration
+        # and write the iteration into the tracerfile. But async saver only
+        # save the state dict into the CPU memory not the storage. The saver
+        # need to clear the empty checkpoint directory.
+        saver.clear_empty_checkpoint(iteration)
+    elif storage_type == StorageType.DISK:
+        args = get_args()
+        saver = MegatronCheckpointManager(args.save)
+        torch_save_func = torch.save
+        torch.save = saver.save
+        megatron_save(iteration, model, optimizer, opt_param_scheduler)
+        saver.engine.save_to_storage(iteration, saver.state_dict, saver.path)
+        torch.save = torch_save_func
+    else:
+        raise ValueError(f"No support storage type {storage_type}")
+
+
+def load_checkpoint(
     model, optimizer, opt_param_scheduler, load_arg="load", strict=True
 ):
     """Load the checkpointing state dict. The method firstly
@@ -131,5 +125,5 @@ def load_latest_checkpoint(
     saver = MegatronCheckpointManager(args.save)
     torch_load_func = torch.load
     torch.load = saver.load
-    load_checkpoint(model, optimizer, opt_param_scheduler, load_arg, strict)
+    megatron_load(model, optimizer, opt_param_scheduler, load_arg, strict)
     torch.load = torch_load_func
diff --git a/dlrover/trainer/torch/megatron/__init__.py b/dlrover/trainer/torch/megatron/__init__.py
deleted file mode 100644
index 561b39cf2..000000000
--- a/dlrover/trainer/torch/megatron/__init__.py
+++ /dev/null
@@ -1,12 +0,0 @@
-# Copyright 2023 The DLRover Authors. All rights reserved.
-# 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.


Test patch:
diff --git a/dlrover/trainer/tests/torch/deepspeed_ckpt_test.py b/dlrover/trainer/tests/torch/deepspeed_ckpt_test.py
index 68b518a3c..ddabb37a3 100644
--- a/dlrover/trainer/tests/torch/deepspeed_ckpt_test.py
+++ b/dlrover/trainer/tests/torch/deepspeed_ckpt_test.py
@@ -22,8 +22,9 @@
 import torch.optim as optim
 
 from dlrover.python.elastic_agent.torch.ckpt_saver import AsyncCheckpointSaver
-from dlrover.trainer.torch.deepspeed.async_checkpoint import (
-    DeepSpeedCheckpointManger,
+from dlrover.trainer.torch.flash_checkpoint.deepspeed import (
+    DeepSpeedCheckpointer,
+    StorageType,
 )
 
 
@@ -103,15 +104,17 @@ def test_save_load(self):
         step = 100
         with tempfile.TemporaryDirectory() as tmpdirname:
             engine = MockDeepSpeedEngine(model, optimizer)
-            ckpt_manager = DeepSpeedCheckpointManger(engine, tmpdirname)
-            ckpt_manager.save_checkpoint_to_memory(tmpdirname, str(step))
-            shm_handler = ckpt_manager._async_save_engine._shm_handler
+            checkpointer = DeepSpeedCheckpointer(engine, tmpdirname)
+            checkpointer.save_checkpoint(
+                tmpdirname, str(step), storage_type=StorageType.MEMORY
+            )
+            shm_handler = checkpointer._async_save_engine._shm_handler
             self.assertFalse(shm_handler.no_checkpint_state())
             self.assertEqual(
-                ckpt_manager._async_save_engine._shm_handler._buffer_size, 9640
+                checkpointer._async_save_engine._shm_handler._buffer_size, 9640
             )
             tensor_meta = (
-                ckpt_manager._async_save_engine._shm_handler.metadata.get()
+                checkpointer._async_save_engine._shm_handler.metadata.get()
             )
             ds_ckpt_config = tensor_meta["_DLORVER_CKPT_CONFIG"]
             self.assertEqual(ds_ckpt_config.step, str(step))
@@ -123,7 +126,9 @@ def test_save_load(self):
             tracer_file = os.path.join(tmpdirname, "latest")
             self.assertFalse(os.path.exists(tracer_file))
 
-            ckpt_manager.save_checkpoint_to_storage(tmpdirname, str(step))
+            checkpointer.save_checkpoint(
+                tmpdirname, str(step), storage_type=StorageType.DISK
+            )
             # Wait asynchronously saving.
             while True:
                 if "100" in os.listdir(tmpdirname):
@@ -137,4 +142,9 @@ def test_save_load(self):
 
             files = os.listdir(tmpdirname + "/100")
             self.assertEqual(files, ["optim_states.pt", "model_states.pt"])
-            ckpt_manager.load_checkpoint(tmpdirname, str(step))
+            checkpointer.load_checkpoint(tmpdirname, str(step))
+
+            with self.assertRaises(ValueError):
+                checkpointer.save_checkpoint(
+                    tmpdirname, str(step), storage_type=2
+                )
diff --git a/dlrover/trainer/tests/torch/fsdp_ckpt_test.py b/dlrover/trainer/tests/torch/fsdp_ckpt_test.py
index ce0309d86..f753ffef9 100644
--- a/dlrover/trainer/tests/torch/fsdp_ckpt_test.py
+++ b/dlrover/trainer/tests/torch/fsdp_ckpt_test.py
@@ -52,6 +52,10 @@
     AsyncCheckpointSaver,
     SharedMemoryHandler,
 )
+from dlrover.trainer.torch.flash_checkpoint.fsdp import (
+    FsdpCheckpointer,
+    StorageType,
+)
 from dlrover.trainer.torch.flash_checkpoint.fsdp_engine import (
     FileReader,
     FsdpCheckpointEngine,
@@ -351,3 +355,40 @@ def test_fsdp_engine(self):
             self.assertListEqual(files, [".metadata", "__0_0.distcp"])
             reader = engine.load(path)
             self.assertTrue(isinstance(reader, SharedMemoryReader))
+
+    def test_fsdp_checkpointer(self):
+        step = 100
+        state_dict = {
+            _OPTIMIZER_KEY: {"learning_rate": 0.1},
+        }
+        with tempfile.TemporaryDirectory() as tmpdir:
+            tmpdir = Path(tmpdir)
+            checkpointer = FsdpCheckpointer(tmpdir)
+            path = tmpdir / str(step)
+            engine = checkpointer._engine
+            checkpointer.save_checkpoint(
+                step, state_dict, path=path, storage_type=StorageType.MEMORY
+            )
+            checkpointer.save_checkpoint(
+                step, state_dict, path=path, storage_type=StorageType.DISK
+            )
+            self.assertEqual(engine._cached_step, 100)
+            time.sleep(1)
+            files = sorted(os.listdir(tmpdir))
+            self.assertListEqual(
+                files,
+                [
+                    "._dlrover_ckpt_stage",
+                    "100",
+                    "latest_checkpointed_iteration.txt",
+                ],
+            )
+            files = sorted(os.listdir(path))
+            self.assertListEqual(files, [".metadata", "__0_0.distcp"])
+            reader = checkpointer.load_checkpoint(path)
+            self.assertTrue(isinstance(reader, SharedMemoryReader))
+
+            with self.assertRaises(ValueError):
+                checkpointer.save_checkpoint(
+                    step, state_dict, path=path, storage_type=2
+                )
diff --git a/dlrover/trainer/tests/torch/megatron_ckpt_test.py b/dlrover/trainer/tests/torch/megatron_ckpt_test.py
index cd589b0c2..bdb91dfe1 100644
--- a/dlrover/trainer/tests/torch/megatron_ckpt_test.py
+++ b/dlrover/trainer/tests/torch/megatron_ckpt_test.py
@@ -24,12 +24,12 @@
 
 from dlrover.python.common.constants import CheckpointConstant
 from dlrover.python.elastic_agent.torch.ckpt_saver import AsyncCheckpointSaver
-from dlrover.trainer.torch.megatron import async_checkpoint
-from dlrover.trainer.torch.megatron.async_checkpoint import (
+from dlrover.trainer.torch.flash_checkpoint import megatron
+from dlrover.trainer.torch.flash_checkpoint.checkpointer import StorageType
+from dlrover.trainer.torch.flash_checkpoint.megatron import (
     MegatronCheckpointManager,
-    load_latest_checkpoint,
-    save_checkpoint_to_memory,
-    save_checkpoint_to_storage,
+    load_checkpoint,
+    save_checkpoint,
 )
 
 
@@ -58,7 +58,7 @@ def setUp(self):
         AsyncCheckpointSaver._saver_instance = None
         AsyncCheckpointSaver.start_async_saving_ckpt()
         mock_func = mock_get_tracker_filename
-        async_checkpoint.get_checkpoint_tracker_filename = mock_func
+        megatron.get_checkpoint_tracker_filename = mock_func
 
     def tearDown(self) -> None:
         if AsyncCheckpointSaver._saver_instance:
@@ -80,7 +80,9 @@ def get_args():
                 args.save = tempfile
                 return args
 
-            def save_checkpoint(iteration, model, optimizer, opt_scheduler):
+            def mock_save_checkpoint(
+                iteration, model, optimizer, opt_scheduler
+            ):
                 state_dict = {
                     "iteration": iteration,
                     "model_states": model.state_dict(),
@@ -91,7 +93,7 @@ def save_checkpoint(iteration, model, optimizer, opt_scheduler):
                 )
                 torch.save(state_dict, path)
 
-            def load_checkpoint(
+            def mock_load_checkpoint(
                 model,
                 optimizer,
                 opt_param_scheduler,
@@ -103,12 +105,14 @@ def load_checkpoint(
                 model.load_state_dict(state_dict["model_states"])
                 optimizer.load_state_dict(state_dict["optim_states"])
 
-            async_checkpoint.save_checkpoint = save_checkpoint
-            async_checkpoint.load_checkpoint = load_checkpoint
-            async_checkpoint.get_args = get_args
+            megatron.megatron_save = mock_save_checkpoint
+            megatron.megatron_load = mock_load_checkpoint
+            megatron.get_args = get_args
 
             ckpt_manager = MegatronCheckpointManager(tmpdirname)
-            save_checkpoint_to_memory(10, model, optimizer, None)
+            save_checkpoint(
+                10, model, optimizer, None, storage_type=StorageType.MEMORY
+            )
             self.assertFalse(
                 ckpt_manager.engine._shm_handler.no_checkpint_state()
             )
@@ -129,7 +133,9 @@ def load_checkpoint(
                 restored_step = int(f.read())
             self.assertEqual(restored_step, 10)
 
-            save_checkpoint_to_storage(20, model, optimizer, None)
+            save_checkpoint(
+                20, model, optimizer, None, storage_type=StorageType.DISK
+            )
 
             # Wait asynchronously saving.
             while True:
@@ -143,4 +149,7 @@ def load_checkpoint(
             self.assertEqual(restored_step, 20)
             files = os.listdir(tmpdirname + "/20")
             self.assertEqual(files, ["checkpoint.pt"])
-            load_latest_checkpoint(model, optimizer, None)
+            load_checkpoint(model, optimizer, None)
+
+            with self.assertRaises(ValueError):
+                save_checkpoint(20, model, optimizer, None, storage_type=2)