add hot start

Author: gzy19990617

paddlenlp/experimental/transformers/inference_model.py

@@ -17,15 +17,15 @@
 
 import os
 from typing import Any, Dict
-
+import time
 import paddle
 
 from paddlenlp.transformers import AutoConfig, AutoInferenceModelForCausalLM
 from paddlenlp.utils.log import logger
 
 
 class InferenceModel:
-    def __init__(self, predictor_args, model_args, nranks=1, rank=0, load_model_from_ipc=False, cold_start=False):
+    def __init__(self, predictor_args, model_args, nranks=1, rank=0, load_model_from_ipc=False, hot_start=False):
         """
         Initialize the Causal Language Model Loader.
 
@@ -43,10 +43,11 @@ def __init__(self, predictor_args, model_args, nranks=1, rank=0, load_model_from
         self.load_model_from_ipc = load_model_from_ipc
         self.model = self._build_model()
         self.shared_buffer_to = False
-        self.local_test = False
+        self.local_test = True
+        self.hot_start = hot_start
 
         # (TODO:gaoziyuan)当前启动服务后直接加载参数，后续进行热启动
-        if load_model_from_ipc and not cold_start:
+        if load_model_from_ipc and not hot_start:
             self.update_parameters()
 
     def _setup_environment(self):
@@ -76,16 +77,26 @@ def _build_model(self):
             tensor_parallel_rank=self.rank,
             load_model_from_ipc=self.load_model_from_ipc,
         )
-        print("gaoziyuan test load from config:", self.model.state_dict())
+        # print("gaoziyuan test load from config:", self.model.state_dict())
         return self.model
 
-    def clear_parameters(self) -> None:
+    def clear_parameters(self, pid=0) -> None:
         """Clear all model parameters."""
+        start_time = time.time()
         for name, param in self.model.state_dict().items():
             logger.info(f"Clearing model parameter: {name}")
             param._clear_data()
+        clear_time = time.time() - start_time
+        logger.info(f"Parameter clearing completed in {clear_time:.2f} seconds")
 
+        self.verify_parameters_cleared()
         logger.info("Model parameters cleared successfully")
+        if self.hot_start:
+            array = np.zeros([1],dtype=np.int32)
+            shm = SharedMemory(create=False, size=array.nbytes, name=f"model_weights_status.{pid}")
+            value = np.ndarray(array.shape, dtype=array.dtype, buffer=shm.buf)
+            value[0] = -2
+
 
     def get_model(self) -> paddle.nn.Layer:
         """Get the underlying model instance."""
@@ -117,6 +128,7 @@ def load_tensor_from_ipc_meta(ipc_state_dict: Dict[str, Any]) -> Dict[str, paddl
 
     def update_parameters(
         self,
+        pid=0,
     ) -> None:
         """
         Update model parameters from IPC state dictionary.
@@ -125,34 +137,98 @@ def update_parameters(
             ipc_state_dict: Dictionary containing new parameters in IPC format
         """
         if self.local_test:
-            state_dict = paddle.load("/root/paddlejob/workspace/env_run/output/model_local_qwen")
-            if not self.shared_buffer_to:
-                print("通过set_state_dict更新参数")
-                self.model.set_state_dict(state_dict)
-        else:
-            model_path = "/shared_ipc_meta"
             current_device_id = int(os.getenv("FLAGS_selected_gpus"))
-            ipc_state_dict_path = os.path.join(model_path, f"ipc_metas_{current_device_id}")
-            ipc_state_dict = paddle.load(ipc_state_dict_path)
-            state_dict = self.load_tensor_from_ipc_meta(ipc_state_dict)
-            if not self.shared_buffer_to:
-                print("通过set_state_dict更新参数")
-                self.model.set_state_dict(state_dict)
-                return
-
-        infer_model_state_dict = self.model.state_dict()
-
-        print("通过shared_buffer_to更新参数")
-        for name, param in state_dict.items():
-            if name in infer_model_state_dict:
-                logger.info(f"Updating model parameter: {name}")
-                update_param = infer_model_state_dict[name]
-                assert (
-                    update_param.dtype == param.dtype
-                ), f"Type mismatch for {name}: {param.dtype} vs {update_param.dtype}"
-                assert (
-                    update_param.shape == param.shape
-                ), f"Shape mismatch for {name}: {param.shape} vs {update_param.shape}"
-                param._share_buffer_to(update_param)
-
-        logger.info("Model parameters updated successfully")
+            model_path = f"/shared_ipc_meta/model_state.tp0{current_device_id}.pdparams"
+            print("model_apth : ", model_path)
+            state_dict = paddle.load(model_path)
+            set_start = time.time()
+            self.model.set_state_dict(state_dict)
+            logger.info(f"set_state_dict completed in {time.time() - set_start:.2f} seconds")
+            self.verify_parameters_updated()
+            return
+
+        start_time = time.time()
+        logger.info("Starting parameter update process...")
+
+        model_path = "/shared_ipc_meta"
+        current_device_id = int(os.getenv("FLAGS_selected_gpus"))
+        ipc_state_dict_path = os.path.join(model_path, f"ipc_metas_{current_device_id}")
+        logger.info(f"ipc_state_dict_path is {ipc_state_dict_path}")
+        ipc_state_dict = paddle.load(ipc_state_dict_path)
+        convert_start = time.time()
+        state_dict = self.load_tensor_from_ipc_meta(ipc_state_dict)
+        logger.info(f"IPC meta converted to tensors in {time.time() - convert_start:.2f} seconds")
+        if not self.shared_buffer_to:
+            logger.info("Updating parameters via set_state_dict...")
+            set_start = time.time()
+            self.model.set_state_dict(state_dict)
+            logger.info(f"set_state_dict completed in {time.time() - set_start:.2f} seconds")
+            self.verify_parameters_updated()
+        else:
+            share_start = time.time()
+            logger.info("通过shared_buffer_to更新参数")
+            infer_model_state_dict = self.model.state_dict()
+            for name, param in state_dict.items():
+                if name in infer_model_state_dict:
+                    logger.info(f"Updating model parameter: {name}")
+                    update_param = infer_model_state_dict[name]
+                    assert (
+                        update_param.dtype == param.dtype
+                    ), f"Type mismatch for {name}: {param.dtype} vs {update_param.dtype}"
+                    assert (
+                        update_param.shape == param.shape
+                    ), f"Shape mismatch for {name}: training {param.shape} vs infer {update_param.shape}"
+                    param._share_buffer_to(update_param)
+                    logger.info(f"Parameter sharing completed in {time.time() - share_start:.2f} seconds")
+            
+        if self.hot_start:
+            array = np.zeros([1],dtype=np.int32)
+            shm = SharedMemory(create=False, size=array.nbytes, name=f"model_weights_status.{pid}")
+            value = np.ndarray(array.shape, dtype=array.dtype, buffer=shm.buf)
+            value[0] = 2
+
+    
+    def verify_parameters_cleared(self) -> bool:
+        """
+        Verify that all model parameters have been cleared.
+        
+        Returns:
+            bool: True if all parameters are cleared, False otherwise
+        """
+        logger.info("Verifying parameters are cleared...")
+        all_cleared = True
+        for name, param in self.model.state_dict().items():
+            if param._is_initialized():
+                logger.error(f"Parameter {name} was not properly cleared!")
+                all_cleared = False
+        
+        if all_cleared:
+            logger.info("All parameters verified as cleared successfully")
+        else:
+            logger.error("Some parameters were not properly cleared!")
+        
+        return all_cleared
+
+    def verify_parameters_updated(self) -> bool:
+        """
+        Verify that model parameters match the source state dictionary.
+        
+        Args:
+            source_state_dict: Dictionary containing the expected parameters
+            
+        Returns:
+            bool: True if all parameters match, False otherwise
+        """
+        logger.info("Verifying parameters are cleared...")
+        all_update = True
+        for name, param in self.model.state_dict().items():
+            if not param._is_initialized():
+                logger.error(f"Parameter {name} was not properly cleared!")
+                all_update = False
+        
+        if all_update:
+            logger.info("All parameters verified as updated successfully")
+        else:
+            logger.error("Some parameters were not properly updated!")
+        
+        return all_update

paddlenlp/experimental/transformers/test_inference_model.py

@@ -13,8 +13,7 @@
 # limitations under the License.
 
 # 后面会删掉，仅提供示例
-from inference_utils import ModelArgument, PredictorArgument
-
+from paddlenlp.experimental.transformers.inference_utils import ModelArgument, PredictorArgument
 from paddlenlp.experimental.transformers.inference_model import InferenceModel
 
 predictor_args = PredictorArgument()
@@ -27,9 +26,15 @@
 # 如果需要
 # predictor_args.quant_type = "weight_only_int8"
 
-inference_model = InferenceModel(predictor_args, model_args, load_model_from_ipc=True, cold_start=False)
+inference_model = InferenceModel(predictor_args, model_args, load_model_from_ipc=True, cold_start=True)
 
 model = inference_model.model
+
+print(inference_model.verify_parameters_cleared())
+
+print(inference_model.verify_parameters_updated())
+
+
 # print(model.get_name_mappings_to_training())
 
 # 获取inference model 的 key\shape\type