Update to Apache-2.0 License, move debloating to checkpointing engine

Signed-off-by: amaurya <[email protected]>

Author: amaurya

deepspeed/datastates/__init__.py

@@ -1,6 +1,6 @@
 # Copyright (c) Microsoft Corporation.
 # SPDX-License-Identifier: Apache-2.0
 
-# MIT License Copyright (c) UChicago Argonne LLC, operator of Argonne National Laboratory.
+# Apache-2.0 License Copyright (c) UChicago Argonne LLC, operator of Argonne National Laboratory.
 
 # DeepSpeed Team

deepspeed/datastates/config.py

@@ -1,7 +1,7 @@
 # Copyright (c) Microsoft Corporation.
 # SPDX-License-Identifier: Apache-2.0
 
-# MIT License Copyright (c) UChicago Argonne LLC, operator of Argonne National Laboratory.
+# Apache-2.0 License Copyright (c) UChicago Argonne LLC, operator of Argonne National Laboratory.
 
 # DeepSpeed Team
 

deepspeed/runtime/checkpoint_engine/checkpoint_engine.py

@@ -32,8 +32,3 @@ def commit(self, tag):
     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):
-        # Check if the checkpoint engine preserves storage sharing
-        # (set to false if cloning is required to get actual tensor sizes)
-        return False

deepspeed/runtime/checkpoint_engine/datastates_checkpoint_engine.py

@@ -1,7 +1,7 @@
 # Copyright (c) Microsoft Corporation.
 # SPDX-License-Identifier: Apache-2.0
 
-# MIT License Copyright (c) UChicago Argonne LLC, operator of Argonne National Laboratory.
+# Apache-2.0 License Copyright (c) UChicago Argonne LLC, operator of Argonne National Laboratory.
 
 # DeepSpeed Team
 
@@ -32,6 +32,3 @@ def commit(self, tag):
 
     def wait(self):
         return self.ckpt_engine.wait()
-
-    def preserves_storage_sharing(self):
-        return True

deepspeed/runtime/checkpoint_engine/nebula_checkpoint_engine.py

@@ -11,6 +11,7 @@
     CheckpointEngine
 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):
@@ -48,7 +49,8 @@ def save(self, state_dict, path: str):
         tag = _get_tag_from_path(path)
         partition_name = os.path.basename(path)
         logger.info(f"[Nebula] Saving {partition_name} under tag {tag}...")
-        self.checkpoint.save(partition_name, state_dict)
+        debloated_state_dict = clone_tensors_for_torch_save(state_dict)
+        self.checkpoint.save(partition_name, debloated_state_dict)
         logger.info(f"[Nebula] Saved {partition_name} under tag {tag}.")
         return None
 

deepspeed/runtime/checkpoint_engine/torch_checkpoint_engine.py

@@ -7,6 +7,7 @@
 from deepspeed.utils import logger, log_dist
 from deepspeed.runtime.checkpoint_engine.checkpoint_engine import \
     CheckpointEngine
+from deepspeed.checkpoint.utils import clone_tensors_for_torch_save
 
 
 class TorchCheckpointEngine(CheckpointEngine):
@@ -19,7 +20,8 @@ def create(self, tag):
 
     def save(self, state_dict, path: str):
         logger.info(f"[Torch] Saving {path}...")
-        torch.save(state_dict, path)
+        debloated_state_dict = clone_tensors_for_torch_save(state_dict)
+        torch.save(debloated_state_dict, path)
         logger.info(f"[Torch] Saved {path}.")
         return None
 

deepspeed/runtime/pipe/module.py

@@ -20,7 +20,6 @@
 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):
@@ -629,12 +628,7 @@ def save_state_dict(self, save_dir, checkpoint_engine, exclude_frozen_params=Fal
                 for n in self._get_frozen_parameter_names(layer):
                     del orig_state_dict[n]
 
-            if checkpoint_engine.preserves_storage_sharing():
-                final_state_dict = orig_state_dict
-            else:
-                final_state_dict = clone_tensors_for_torch_save(orig_state_dict)
-
-            checkpoint_engine.save(final_state_dict, model_ckpt_path)
+            checkpoint_engine.save(orig_state_dict, model_ckpt_path)
 
     def load_state_dir(self, load_dir, checkpoint_engine, strict=True):
         for idx, layer in enumerate(self.forward_funcs):