fix log

Author: gzy19990617

paddlenlp/experimental/transformers/inference_model.py

@@ -82,38 +82,6 @@ def _build_model(self):
         # print("gaoziyuan test load from config:", self.model.state_dict())
         return self.model
 
-    def clear_parameters(self, pid=0) -> None:
-        """Clear all model parameters."""
-
-        if self.verify_parameters_cleared():
-            logger.info("Parameters already cleared!")
-            array = np.zeros([self.nranks],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[self.rank] = -2
-            return
-
-        start_time = time.time()
-        paddle.device.cuda.empty_cache()
-        self.check_memory_usage("start clear parameters")
-        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")
-        
-        array = np.zeros([self.nranks],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[self.rank] = -2
-        paddle.device.cuda.empty_cache()
-        self.check_memory_usage("clear parameters end")
-        logger.info("send clear sigal !")
-
-
     def get_model(self) -> paddle.nn.Layer:
         """Get the underlying model instance."""
         return self.model
@@ -142,103 +110,131 @@ def load_tensor_from_ipc_meta(ipc_state_dict: Dict[str, Any]) -> Dict[str, paddl
 
         return result
 
-    def check_memory_usage(self, msg=""):
-        """ check_memory_usage """
-        max_memory_allocated_size = paddle.device.cuda.max_memory_allocated()/(1024*1024*1024)
-        max_memory_reserved_size = paddle.device.cuda.max_memory_reserved()/(1024*1024*1024)
-        memory_allocated_size = paddle.device.cuda.memory_allocated()/(1024*1024*1024)
-        memory_reserved_size = paddle.device.cuda.memory_reserved()/(1024*1024*1024)
-        logger.info(msg)
-        logger.warning(f"checking gpu memory usage {msg}:\nmax_memory_allocated_size: {max_memory_allocated_size}GB\nmax_memory_reserved_size: {max_memory_reserved_size}GB\nmemory_allocated_size: {memory_allocated_size}GB\nmemory_reserved_size: {memory_reserved_size}GB")
+    def _log_memory_usage(self, context: str = "") -> None:
+        """Log current GPU memory usage."""
+        max_alloc = paddle.device.cuda.max_memory_allocated() / (1024 ** 3)
+        max_reserved = paddle.device.cuda.max_memory_reserved() / (1024 ** 3)
+        curr_alloc = paddle.device.cuda.memory_allocated() / (1024 ** 3)
+        curr_reserved = paddle.device.cuda.memory_reserved() / (1024 ** 3)
+        
+        logger.info(f"GPU memory usage {context}:")
+        logger.warning(
+            f"max_allocated: {max_alloc:.2f}GB\n"
+            f"max_reserved: {max_reserved:.2f}GB\n"
+            f"current_allocated: {curr_alloc:.2f}GB\n"
+            f"current_reserved: {curr_reserved:.2f}GB"
+        )
 
     def generate(self, **kwargs):
         self.model.generate(**kwargs)
+    
+    def _update_shared_status(self, pid: int, status: int) -> None:
+        """Update shared memory status flag."""
+        array = np.zeros([self.nranks], 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[self.rank] = status
 
-    def update_parameters(
-        self,
-        pid=0,
-    ) -> None:
-        """
-        Update model parameters from IPC state dictionary.
-
-        Args:
-            ipc_state_dict: Dictionary containing new parameters in IPC format
-        """
-        if self.verify_parameters_updated() and not self.first_load:
+    def update_parameters(self, pid: int = 0) -> None:
+        """Update model parameters from IPC state dictionary."""
+        if self.verify_parameters_updated(False) and not self.first_load:
             logger.info("Parameters already updated.")
-            array = np.zeros([self.nranks],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[self.rank] = 2
+            self._update_shared_status(pid, 2)
             return
 
         paddle.device.cuda.empty_cache()
-        self.check_memory_usage("start update parameters")
+        self._log_memory_usage("start update parameters")
+        
         if self.local_test:
             current_device_id = int(os.getenv("FLAGS_selected_gpus"))
             model_path = f"/shared_ipc_meta/model_state.tp0{current_device_id}.pdparams"
-            print("model_apth : ", model_path)
-            state_dict = paddle.load(model_path)
+            logger.info(f"Loading model from: {model_path}")
+            
             set_start = time.time()
-            self.model.set_state_dict(state_dict)
+            self.model.set_state_dict(paddle.load(model_path))
             logger.info(f"set_state_dict completed in {time.time() - set_start:.2f} seconds")
+            
             self.verify_parameters_updated()
-            self.check_memory_usage("update parameters end")
+            self._log_memory_usage("update parameters end")
+            
             if not self.first_load:
                 logger.info("send update signal")
-                array = np.zeros([self.nranks],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[self.rank] = 2
+                self._update_shared_status(pid, 2)
             self.first_load = False
             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)
+        ipc_state_dict_path = f"/shared_ipc_meta/ipc_metas_{current_device_id}"
+        logger.info(f"Loading IPC state dict from: {ipc_state_dict_path}")
+        
         convert_start = time.time()
-        state_dict = self.load_tensor_from_ipc_meta(ipc_state_dict)
+        state_dict = self.load_tensor_from_ipc_meta(paddle.load(ipc_state_dict_path))
         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:
+            logger.info("Updating parameters via shared_buffer_to...")
             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}"
+                    
+                    if update_param.dtype != param.dtype:
+                        raise TypeError(f"Type mismatch for {name}: {param.dtype} vs {update_param.dtype}")
+                    if update_param.shape != param.shape:
+                        raise ValueError(f"Shape mismatch for {name}: {param.shape} vs {update_param.shape}")
+                    
                     param._share_buffer_to(update_param)
-                    logger.info(f"Parameter sharing completed in {time.time() - share_start:.2f} seconds")
+            
+            logger.info(f"Parameter sharing completed in {time.time() - share_start:.2f} seconds")
 
         if not self.first_load:
             logger.info("send update signal")
-            array = np.zeros([self.nranks],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[self.rank] = 2
+            self._update_shared_status(pid, 2)
+        
         self.first_load = False
+        self.verify_parameters_updated()
+        paddle.device.cuda.empty_cache()
+        self._log_memory_usage("update parameters end")
+    
+    def clear_parameters(self, pid: int = 0) -> None:
+        """Clear all model parameters."""
+        if self.verify_parameters_cleared(False):
+            logger.info("Parameters already cleared!")
+            self._update_shared_status(pid, -2)
+            return
+
+        start_time = time.time()
+        paddle.device.cuda.empty_cache()
+        self._log_memory_usage("start clear parameters")
+        
+        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")
 
+        self._update_shared_status(pid, -2)
         paddle.device.cuda.empty_cache()
-        self.check_memory_usage("update parameters end")
+        self._log_memory_usage("clear parameters end")
+        logger.info("send clear signal!")
 
 
-    def verify_parameters_cleared(self) -> bool:
+    def verify_parameters_cleared(self, erro_log:bool = True) -> bool:
         """
         Verify that all model parameters have been cleared.
         
@@ -249,17 +245,19 @@ def verify_parameters_cleared(self) -> bool:
         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!")
+                if erro_log:
+                    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!")
+            if erro_log:
+                logger.error("Some parameters were not properly cleared!")
 
         return all_cleared
 
-    def verify_parameters_updated(self) -> bool:
+    def verify_parameters_updated(self, erro_log:bool = True) -> bool:
         """
         Verify that model parameters match the source state dictionary.
         
@@ -273,12 +271,14 @@ def verify_parameters_updated(self) -> bool:
         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!")
+                if erro_log:
+                    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!")
+            if erro_log
+                logger.error("Some parameters were not properly updated!")
 
         return all_update