Update DataStates checkpoint engine with decoupled functionality

Author: amaurya

deepspeed/datastates/README.md

@@ -1,3 +1,12 @@
 # DataStates-LLM checkpointing engine.
 
-This feature is not enabled by default. To enable, set the following options in ds_config.json and download [DataStates-LLM checkpointing library](https://github.com/DataStates/datastates-llm/). A detailed tutorial is available [here](../../docs/_tutorials/datastates-async-checkpointing.md).
+This feature is not enabled by default. To enable, set the following options in ds_config.json and download the [DataStates-LLM checkpointing library](https://github.com/DataStates/datastates-llm/). A detailed tutorial is available [here](../../docs/_tutorials/datastates-async-checkpointing.md).
+
+```
+{
+    ... other deepspeed config options,
+    "datastates_ckpt": {
+        "host_cache_size": 16
+	}
+}
+```

deepspeed/datastates/config.py

@@ -6,16 +6,16 @@
 # DeepSpeed Team
 
 from deepspeed.runtime.config_utils import DeepSpeedConfigObject
+import copy
+
+DATASTATES_CHECKPOINTING = "datastates_ckpt"
+DATASTATES_CHECKPOINTING_ENABLED = False
 
 
 class DeepSpeedDataStatesConfig(DeepSpeedConfigObject):
 
     def __init__(self, param_dict):
         super(DeepSpeedDataStatesConfig, self).__init__()
 
-        self.enabled = None
-        self.config = {}
-
-        if "datastates_ckpt" in param_dict.keys():
-            self.enabled = True
-            self.config = param_dict["datastates_ckpt"]
+        self.enabled = param_dict.get(DATASTATES_CHECKPOINTING, DATASTATES_CHECKPOINTING_ENABLED) is not False
+        self.config = copy.deepcopy(param_dict.get(DATASTATES_CHECKPOINTING, None))

deepspeed/runtime/checkpoint_engine/README.md

@@ -39,13 +39,12 @@ class CheckpointEngine(object):
 
 ### Asynchronous Lazy Checkpointing using DataStates-LLM
 
-DataStates-LLM is an asynchronous checkpointing approach optimized for LLM pre-training and can be obtained at https://github.com/DataStates/datastates-llm. A detailed tutorial is available [here](../../../docs/_tutorials/datastates-async-checkpointing.md). To enable datastates-llm checkpointing, specify the `host_cache_size` (in gigabytes) which reserves pinned host memory for asynchronous checkpoint flushing, and `parser_threads` to parse multiple checkpoint file requests in parallel using the following lines in config.json supplied during the launch:
+DataStates-LLM is an asynchronous checkpointing approach optimized for LLM pre-training and can be obtained at https://github.com/DataStates/datastates-llm. A detailed tutorial is available [here](../../../docs/_tutorials/datastates-async-checkpointing.md). To enable datastates-llm checkpointing, specify the `host_cache_size` (in gigabytes) which reserves pinned host memory for asynchronous checkpoint flushing using the following lines in config.json supplied during the launch:
 ```
 {
     ... other deepspeed config options,
     "datastates_ckpt": {
-        "host_cache_size": 16,
-        "parser_threads": 8
+        "host_cache_size": 16
 	}
 }
 ```

deepspeed/runtime/checkpoint_engine/checkpoint_engine.py

@@ -59,6 +59,5 @@ def get_commit_info(self):
     def cleanup(self):
         pass
 
-    def wait(self):
-        # To wait in asynchronous checkpoint engines (e.g. DataStates-LLM) for the previous snapshot to finish
-        pass
+    def preserves_storage_sharing(self):
+        return True

deepspeed/runtime/checkpoint_engine/datastates_checkpoint_engine.py

@@ -5,20 +5,25 @@
 
 # DeepSpeed Team
 
-from deepspeed.utils import log_dist
 from deepspeed.runtime.checkpoint_engine.checkpoint_engine import \
-    CheckpointEngine
-from datastates.llm import Checkpointing
+    CheckpointEngine, CheckpointCommitInfo
+from datastates import CheckpointEngine as DataStatesEngine
+
+ENGINE_NAME = "DataStatesCheckpointEngine"
 
 
 class DataStatesCheckpointEngine(CheckpointEngine):
 
     def __init__(self, deepspeed_config, rank):
         super().__init__(deepspeed_config)
-        self.ckpt_engine = Checkpointing(deepspeed_config, rank)
+        self.commit_info = None
+        self.ckpt_engine = DataStatesEngine(deepspeed_config, rank)
+
+    def __del__(self):
+        self.cleanup()
 
-    def create(self, tag):
-        log_dist(f"[DataStates] Checkpoint {tag} is about to be saved!", ranks=[0])
+    def create(self, info: CheckpointCommitInfo):
+        self.commit_info = info
         return None
 
     def save(self, state_dict, path: str):
@@ -27,8 +32,18 @@ def save(self, state_dict, path: str):
     def load(self, path: str, map_location=None):
         return self.ckpt_engine.load(path, map_location)
 
-    def commit(self, tag):
-        return self.ckpt_engine.commit(tag)
+    def commit(self, info: CheckpointCommitInfo):
+        assert info == self.commit_info
+        self.ckpt_engine.wait()
+        return self.ckpt_engine.commit(info.tag)
+
+    def cleanup(self):
+        self.commit(self.commit_info)
+        self.ckpt_engine.wait(True)
+        del self.ckpt_engine
+
+    def is_decoupled(self):
+        return True
 
-    def wait(self):
-        return self.ckpt_engine.wait()
+    def preserves_storage_sharing(self):
+        return False

deepspeed/runtime/checkpoint_engine/nebula_checkpoint_engine.py

@@ -11,7 +11,6 @@
     CheckpointEngine, CheckpointCommitInfo
 from deepspeed.utils import logger, log_dist
 from deepspeed.nebula.constants import *
-from deepspeed.checkpoint.utils import clone_tensors_for_torch_save
 
 
 def _get_tag_from_path(path):

deepspeed/runtime/checkpoint_engine/utils.py

@@ -6,6 +6,7 @@
 from deepspeed.runtime.model_checkpointing.constants import *
 from deepspeed.runtime.model_checkpointing.utils import create_data_parallel_writer_config
 from deepspeed.utils import logger
+from deepspeed import comm as dist
 
 from .decoupled_checkpoint_engine import DecoupledCheckpointEngine
 from .fast_checkpoint_engine import FastCheckpointEngine
@@ -35,4 +36,14 @@ def create_checkpoint_engine(config_params, groups, zero_stage, has_moe_layers,
             else:
                 return NebulaCheckpointEngine(config_params=config_params.nebula_config)
 
+        if config_params.datastates_config.enabled:
+            try:
+                from deepspeed.runtime.checkpoint_engine.datastates_checkpoint_engine import DataStatesCheckpointEngine
+                return DataStatesCheckpointEngine(deepspeed_config=config_params, rank=dist.get_rank())
+            except ImportError as err:
+                logger.error(
+                    f"No datastates engine found! Install from https://github.com/DataStates/datastates-llm. Will fall back to torch.save. Details: {err}"
+                )
+                return TorchCheckpointEngine(config_params)
+
     return TorchCheckpointEngine(config_params)

deepspeed/runtime/engine.py

@@ -1140,16 +1140,6 @@ def _configure_checkpointing(self):
                                                           has_moe_layers=self.has_moe_layers,
                                                           optimize_dp_state=optimize_dp_state)
 
-        if self._config is not None and self._config.datastates_config.enabled:
-            try:
-                from deepspeed.runtime.checkpoint_engine.datastates_checkpoint_engine import DataStatesCheckpointEngine
-                self.checkpoint_engine = DataStatesCheckpointEngine(deepspeed_config=self._config,
-                                                                    rank=dist.get_rank())
-            except ImportError as err:
-                raise Exception(
-                    f"The datastates-llm checkpoint engine was not found! Will fall back to torch.save. Details: {err}"
-                )
-
         dp_rank = groups._get_sequence_data_parallel_rank()
         rank = self.local_rank if self.use_node_local_storage() else dp_rank
 
@@ -2420,11 +2410,6 @@ def _take_model_step(self, lr_kwargs, block_eigenvalue={}):
                 master_params = amp.master_params(self.optimizer)
                 clip_grad_norm_(parameters=master_params, max_norm=self.gradient_clipping(), mpu=self.mpu)
 
-        try:
-            self.checkpoint_engine.wait()
-        except Exception as exc:
-            logger.error(f"Error during optimizer wait step: {exc}")
-
         self.optimizer.step()
 
         if hasattr(self.optimizer, '_global_grad_norm'):
@@ -3610,7 +3595,9 @@ def _save_moe_checkpoint(self, save_dir, tag, client_state={}, exclude_frozen_pa
                     moe_save_path = self._get_expert_ckpt_name(save_dir, moe_layer_id, global_expert_id, tag, self.mpu)
                     if self.random_ltd_enabled():
                         expert_state_dict = remove_random_ltd_state_dict(expert_state_dict)
-                    saveable_state_dict = clone_tensors_for_torch_save(expert_state_dict)
+                    saveable_state_dict = expert_state_dict
+                    if self.checkpoint_engine.preserves_storage_sharing():
+                        saveable_state_dict = clone_tensors_for_torch_save(expert_state_dict)
                     self.checkpoint_engine.save(saveable_state_dict, moe_save_path)
                 moe_layer_id += 1
 
@@ -3632,7 +3619,9 @@ def _save_moe_checkpoint(self, save_dir, tag, client_state={}, exclude_frozen_pa
         }
         # TODO: why use BufferedWriter not the path
         file_path = self._get_optimizer_ckpt_name(save_dir, tag, expp_rank)
-        saveable_state_dict = clone_tensors_for_torch_save(optimizer_state)
+        saveable_state_dict = optimizer_state
+        if self.checkpoint_engine.preserves_storage_sharing():
+            saveable_state_dict = clone_tensors_for_torch_save(optimizer_state)
         self.checkpoint_engine.save(saveable_state_dict, file_path)
 
         # Load flow uses below saved file for model parameters, RNG and more
@@ -3672,7 +3661,9 @@ def _save_moe_checkpoint(self, save_dir, tag, client_state={}, exclude_frozen_pa
             }
             state.update(client_state)
             logger.info(f'Saving model checkpoint: {save_path}')
-            saveable_state_dict = clone_tensors_for_torch_save(state)
+            saveable_state_dict = state
+            if self.checkpoint_engine.preserves_storage_sharing():
+                saveable_state_dict = clone_tensors_for_torch_save(state)
             self.checkpoint_engine.save(saveable_state_dict, save_path)
 
     def _create_checkpoint_file(self, save_dir, tag, zero_checkpoint):

deepspeed/runtime/pipe/module.py

@@ -20,6 +20,7 @@
 from .topology import PipeDataParallelTopology, PipelineParallelGrid
 from deepspeed.runtime.state_dict_factory import SDLoaderFactory
 from deepspeed.accelerator import get_accelerator
+from deepspeed.checkpoint.utils import clone_tensors_for_torch_save
 
 
 class PipelineError(Exception):
@@ -620,6 +621,7 @@ def save_state_dict(self, save_dir, checkpoint_engine, exclude_frozen_params=Fal
         layer_list = self.forward_funcs[start:end]
 
         checkpoint_engine.makedirs(save_dir, exist_ok=True)
+        should_clone = checkpoint_engine.preserves_storage_sharing()
         for idx, layer in enumerate(layer_list):
             model_ckpt_path = self.ckpt_layer_path(save_dir, start + idx)
             if not hasattr(layer, 'state_dict'):
@@ -629,7 +631,9 @@ def save_state_dict(self, save_dir, checkpoint_engine, exclude_frozen_params=Fal
             if exclude_frozen_params:
                 for n in self._get_frozen_parameter_names(layer):
                     del orig_state_dict[n]
-            final_state_dict = clone_tensors_for_torch_save(orig_state_dict)
+            final_state_dict = orig_state_dict
+            if should_clone:
+                final_state_dict = clone_tensors_for_torch_save(orig_state_dict)
             checkpoint_engine.save(state_dict=final_state_dict, path=model_ckpt_path)
 
     def load_state_dir(self, load_dir, checkpoint_engine, strict=True):