fix to make ucp load more lenient

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

Author: Schwidola0607

deepspeed/runtime/engine.py

@@ -2899,7 +2899,7 @@ def _get_all_ckpt_names(self, checkpoints_path, tag):
 
         ckpt_files = glob.glob(ckpt_file_pattern)
         ckpt_files.sort()
-        return ckpt_files
+        return ckpt_files, ckpt_file_pattern
 
     def load_checkpoint(self,
                         load_dir,
@@ -2923,7 +2923,7 @@ def load_checkpoint(self,
 
         Returns:
             A tuple of ``load_path`` and ``client_state``.
-            *``load_path``: Path of the loaded checkpoint. ``None`` if loading the checkpoint failed.
+            *``load_path``: Path of the loaded checkpoint. ``None`` if loading the checkpoint failed or loading a HF based UCP
             *``client_state``: State dictionary used for loading required training states in the client code.
 
         Important: under ZeRO3, one cannot load checkpoint with ``engine.load_checkpoint()`` right
@@ -2962,6 +2962,12 @@ def load_checkpoint(self,
                                                          custom_load_fn=custom_load_fn)
 
         load_zero_checkpoint = load_path is not None and (self.zero_optimization() or self.bfloat16_enabled())
+        # import pdb; pdb.set_trace()
+        if self.load_universal_checkpoint():
+            ucp_ckpt_folder = os.path.join(load_dir, tag)
+            # UCP load can ignore '*mp' files or '*model_states.pt' but ucp_ckpt_folder must exist
+            load_zero_checkpoint = os.path.isdir(ucp_ckpt_folder)
+
         if load_zero_checkpoint:
             if (load_optimizer_states and not load_module_only) or self.load_universal_checkpoint():
                 success = self._load_zero_checkpoint(load_dir, tag, load_optimizer_states=load_optimizer_states)
@@ -3002,7 +3008,11 @@ def _load_checkpoint(self,
 
         from deepspeed.runtime.state_dict_factory import SDLoaderFactory
 
-        ckpt_list = self._get_all_ckpt_names(load_dir, tag)
+        ckpt_list, ckpt_file_pattern = self._get_all_ckpt_names(load_dir, tag)
+        if self.load_universal_checkpoint() and len(ckpt_list) == 0:
+            logger.warning(f"Unable to find {ckpt_file_pattern} files in UCP folder {load_dir}")
+            return None, {}
+        
         sd_loader = SDLoaderFactory.get_sd_loader(ckpt_list, checkpoint_engine=self.checkpoint_engine)
 
         is_pipe_parallel = isinstance(self.module, PipelineModule)

deepspeed/runtime/zero/stage3.py

@@ -2777,7 +2777,6 @@ def load_hp_checkpoint_state_from_checkpoint_dir_stage3(self, checkpoint_dir, pa
             # Purge the swapped optimizer state, it was initialized to the freshly created model and not the checkpoint
             self.optimizer_swapper.purge_state()
 
-        if self.swap_optimizer:
             # Touch all parameters to synchronize all buffers
             timer_names = set()
             self._partition_all_parameters()
@@ -2812,7 +2811,7 @@ def _load_global_state_stage3(self, sd):
 
     def load_hp_checkpoint_state(self, folder, key):
         local_rank = dist.get_local_rank()
-
+        
         # Load tensors from files and reshape them to flat vectors
         loaded_state = torch.load(os.path.join(folder, f"{key}.pt"), weights_only=False)
         if isinstance(loaded_state, dict):