support optm

Author: li126com

internlm/checkpoint/checkpoint_manager.py

@@ -43,6 +43,7 @@
 )
 from .load_funcs import LOAD_FUNC_DICT
 from .utils import process_load_info
+from internlm.checkpoint.vescale.api import save as vescale_save 
 
 logger = get_logger(__file__)
 internlm_accelerator = get_accelerator()
@@ -60,7 +61,7 @@ class CheckpointLoadContent:
     SCHEDULAER = "scheduler"
 
 
-def try_load_internevo_ckpt(ckpt_mm, load_info, train_state: TrainState = None):
+def try_load_internevo_ckpt(ckpt_mm, load_info, train_state: TrainState = None, universal_ckpt=False):
     """Tries to load a checkpoint from the given folder.
 
     Args:
@@ -82,8 +83,13 @@ def try_load_internevo_ckpt(ckpt_mm, load_info, train_state: TrainState = None):
         and the checkpoint manager ckpt_mm and train state objects
     """
     load_content_str, load_ckpt_folder, load_content = process_load_info(load_info)
+    
+    if universal_ckpt:
+        from internlm.checkpoint.vescale.api import load as vescale_load
+        checkpoint_state = {"model": ckpt_mm.model, "optimizer": ckpt_mm.optimizer}
+        vescale_load(load_ckpt_folder, checkpoint_state, broadcast_checkpoint=False)
 
-    if load_content.need_load(CheckpointLoadContent.MODEL):
+    if not universal_ckpt and load_content.need_load(CheckpointLoadContent.MODEL):
         load_model_checkpoint(folder=load_ckpt_folder, model=ckpt_mm.model)
         load_content_str += f"{CheckpointLoadContent.MODEL}, "
 
@@ -93,7 +99,7 @@ def try_load_internevo_ckpt(ckpt_mm, load_info, train_state: TrainState = None):
             load_context(load_ckpt_folder, train_state)
 
         # load optimizer states.
-        if load_content.need_load(CheckpointLoadContent.OPIMIZER):
+        if not universal_ckpt and load_content.need_load(CheckpointLoadContent.OPIMIZER):
             load_optimizer_checkpoint(load_ckpt_folder, ckpt_mm.optimizer)
             load_content_str += f"{CheckpointLoadContent.OPIMIZER}, "
         else:
@@ -110,6 +116,7 @@ def try_load_internevo_ckpt(ckpt_mm, load_info, train_state: TrainState = None):
                     logger.warning("CheckpointManager has no 'lr_scheduler', skip reload lr_scheduler checkpoint!")
 
             if not load_content.need_load(CheckpointLoadContent.OPIMIZER):
+                assert False
                 if ckpt_mm.lr_scheduler and train_state:
                     gpc.config.only_load_lr = True
                     load_optimizer_checkpoint(load_ckpt_folder, ckpt_mm.optimizer)
@@ -419,10 +426,11 @@ def is_now_to_save_ckpt(self, train_state, force=False) -> (bool, CheckpointSave
     def try_save_checkpoint(self, train_state, force=False):
         if not self.enable_save_ckpt:
             return False
-
+        
         save_ckpts, save_type, now_break = self.is_now_to_save_ckpt(train_state, force=force)
 
         if save_ckpts:
+            begin = time.time()
             # Wait for the previous round of asynchronous upload storage to complete.
             self.storage_manager.wait()
             if save_type == CheckpointSaveType.SNAPSHOT_CHECKPOINT:
@@ -440,6 +448,7 @@ def try_save_checkpoint(self, train_state, force=False):
                 train_state=train_state,
                 model_config=self.model_config,
                 model_config_file=self.model_config_file,
+                universal_ckpt=gpc.config.ckpt.universal_ckpt,
             )
 
             if (
@@ -460,6 +469,8 @@ def try_save_checkpoint(self, train_state, force=False):
                     f"Finish to convert internevo2hf checkpoint from {save_ckpt_folder} to {save_hf_ckpt_folder}."
                 )
             torch.distributed.barrier()
+            end = time.time() - begin
+            print(f"finsh save time {gpc.get_global_rank()}: {end}", flush=True)
 
         return now_break
 
@@ -576,12 +587,19 @@ def try_resume_training(self, train_state: TrainState, current_time=""):
                     f"dp={gpc.get_local_rank(ParallelMode.DATA)}==========="
                 )
         else:
+            begin = time.time()
             load_path = self.load_ckpt_info["path"]
             load_content = self.load_ckpt_info["content"]
             load_type = self.load_ckpt_info["ckpt_type"]
+            universal_ckpt = gpc.config.ckpt.universal_ckpt
+            kwargs = {}
+            
+            if universal_ckpt:
+                assert load_type == "internevo", "Only internevo ckpt support universal ckpt."
+                kwargs = {"universal_ckpt": universal_ckpt}
 
             load_func = CheckpointLoadMethod.get_ckpt_load_type_func(load_type)
-            load_content_str = load_func(self, self.load_ckpt_info, train_state)
+            load_content_str = load_func(self, self.load_ckpt_info, train_state, **kwargs)
 
             # If we only load model weight, we need rewrite zero optim's fp32 buffer.
             if (
@@ -598,6 +616,8 @@ def try_resume_training(self, train_state: TrainState, current_time=""):
                 )
                 if load_content_str:
                     logger.info(f"===========Load contents are: {load_content_str}")
+            end = time.time() - begin
+            print(f"finsh load time {gpc.get_global_rank()}: {end}", flush=True)
 
     @llm_timeout(func_name="save_checkpoint")
     def save_checkpoint(
@@ -609,6 +629,7 @@ def save_checkpoint(
         train_state: TrainState,
         model_config: Dict = None,
         model_config_file: str = None,
+        universal_ckpt=False,
     ):
         """
         Save checkpoint to the given folder path.
@@ -621,13 +642,23 @@ def save_checkpoint(
         if gpc.is_rank_for_log():
             logger.info(f"Saving checkpoint to `{folder}` at batch count:{train_state.step_count}...")
 
-        timer("save-model").start()
-        save_model_checkpoint(folder=folder, model=model)
-        timer("save-model").stop()
+        if not universal_ckpt:
+            print(f"save ckpt: base", flush=True)
+            timer("save-model").start()
+            save_model_checkpoint(folder=folder, model=model)
+            timer("save-model").stop()
 
-        timer("save-optimizer").start()
-        save_optimizer_checkpoint(optim=optimizer, state_path=folder)
-        timer("save-optimizer").stop()
+            timer("save-optimizer").start()
+            save_optimizer_checkpoint(optim=optimizer, state_path=folder)
+            timer("save-optimizer").stop()
+        else:
+            print(f"save ckpt: universal", flush=True)
+            vescale_save(
+                path=folder,
+                checkpoint_state={"model": model, "optimizer": optimizer},
+                async_checkpoint=False,
+            )
+            
 
         if (
             hasattr(train_state, "data_state_dict")