fix ps alloc err & avoid mem fragmentation

Author: hongchao

checkpoint_engine/worker.py

@@ -167,7 +167,7 @@ def update_weights_from_ipc(self, zmq_handles: dict[str, str]):
             self.device = torch.device(f"npu:{self.local_rank}")
         assert self.device is not None
 
-        def _load_weights(weights: _WEIGHTS_TYPE):
+        def _update_weights(weights: _WEIGHTS_TYPE):
             # Load main model weights
             self.model_runner.model.load_weights(weights)
             # Load drafter model weights if MTP/speculative decoding is enabled
@@ -177,7 +177,7 @@ def _load_weights(weights: _WEIGHTS_TYPE):
             ):
                 self.model_runner.drafter.model.load_weights(weights=weights)
 
-        def _post_hook():
+        def _process_weight_after_loading():
             process_weights_after_loading(self.model_runner.model, self.model_config, self.device)
             # Also trigger drafter model's post processing if MTP is enabled
             if (
@@ -188,10 +188,15 @@ def _post_hook():
                     self.model_runner.drafter.model, self.model_config, self.device
                 )
 
+        torch.cuda.empty_cache()
+
         update_weights_from_ipc(
             self._zmq_ctx,
             zmq_handles[self._device_uuid],
             device_id=self.device.index,
-            run=_load_weights,
-            post_hook=_post_hook,
+            run=_update_weights,
+            post_hook=_process_weight_after_loading,
         )
+
+        if getattr(self, "_sampler_warmup", None) is not None:
+            self._sampler_warmup()