Optimize sleep level=1 for turbomind backend (#4074)

* use pinned buffer to store weights

* add /is_sleeping endpoint

* update pt

* fix comments

Author: irexyc

lmdeploy/pytorch/engine/model_agent.py

@@ -1052,7 +1052,8 @@ def sleep(self, level: int = 1):
         self.cache_engine = None
         self.reset_graph_runner()
         device = 'cpu' if level == 1 else 'meta'
-        self.patched_model.get_model().to(device=device)
+        self.patched_model.get_model().to(device=device, non_blocking=True)
+        torch.cuda.synchronize()
         torch.cuda.empty_cache()
 
     @torch.inference_mode()

lmdeploy/serve/async_engine.py

@@ -332,6 +332,8 @@ def __init__(self,
         else:
             raise ValueError(f'unsupported backend {backend}')
         self.backend_config = self.engine.engine_config
+        self.is_sleeping = backend_config.empty_init
+        self.sleeping_tags: set[str] = set() if not backend_config.empty_init else {'weights', 'kv_cache'}
         logger.info(f'updated backend_config={self.backend_config}')
 
         # parameters for member functions
@@ -491,6 +493,8 @@ def sleep(self, level: int = 1):
                 discard both the model weights and the kv cache.
         """
         self.engine.sleep(level)
+        self.sleeping_tags = {'weights', 'kv_cache'}
+        self.is_sleeping = True
 
     def wakeup(self, tags: Optional[List[str]] = None):
         """Wake up the model.
@@ -502,11 +506,17 @@ def wakeup(self, tags: Optional[List[str]] = None):
                 wake_up should be called with all tags (or None) before the
                 engine is used again.
         """
+        tags = tags or list(self.sleeping_tags)
+        if any(tag not in self.sleeping_tags for tag in tags):
+            logger.warning(f'some tag in {tags} not in sleeping tags {self.sleeping_tags}')
+            return
         self.engine.wakeup(tags)
         # for TM backend, sleep/wakeup will reset gateway, therefore we need to rebuild instance
-        if self.backend == 'turbomind' and (tags is None or 'kv_cache' in tags):
+        if self.backend == 'turbomind' and 'kv_cache' in tags:
             self.instances = [self.engine.create_instance() for _ in range(self.instance_num)]
             self.free_insts = None
+        self.sleeping_tags = self.sleeping_tags - set(tags)
+        self.is_sleeping = bool(self.sleeping_tags)
 
     def _get_limiter(self):
         if not self.limiter:

lmdeploy/serve/openai/api_server.py

@@ -1115,6 +1115,12 @@ async def wakeup(raw_request: Request = None):
     return Response(status_code=200)
 
 
+@router.get('/is_sleeping', dependencies=[Depends(check_api_key)])
+async def is_sleeping(raw_request: Request = None):
+    is_sleeping = VariableInterface.async_engine.is_sleeping
+    return JSONResponse(content={'is_sleeping': is_sleeping})
+
+
 """ PD Disaggregation API Begin """
 
 

lmdeploy/turbomind/turbomind.py

@@ -213,6 +213,7 @@ def _get_model_params(self):
 
         model_comm = self.model_comm
         tm_params = self._tm_model.tm_params
+        tm_params.clear()
 
         def _get_params(device_id, que):
             rank = self.node_id * self.gpu_count + device_id

src/turbomind/models/llama/LlamaWeight.cc

@@ -116,6 +116,7 @@ void LlamaWeight::release()
     }
 
     decoder_layer_weights.clear();
+    pinned_weights_.clear();
 
     // Wait for deallocations
     core::Context::stream().Sync();
@@ -127,21 +128,22 @@ void LlamaWeight::release()
 
 void LlamaWeight::to_device(const core::Device& device)
 {
-    core::ContextGuard guard = context();
-
-    auto to_device = [&](Tensor& x) -> Tensor {
-        auto tmp = std::exchange(x, empty_like(x, device));
-        Copy(tmp, x);
-        return tmp;
-    };
-
-    std::vector<Tensor> tmp_cpu_tensors;
+    TM_CHECK(device.type == kCPU || device.type == kDEVICE);
+    core::ContextGuard guard{stream_, alloca_, Allocator{kCPUpinned}};
 
     auto tensor_ptr_map = get_parameters();
     for (auto& [name, tensor_ptr] : tensor_ptr_map) {
-        auto tmp_tensor = to_device(*tensor_ptr);
-        if (tmp_tensor.device().type != kDEVICE) {
-            tmp_cpu_tensors.push_back(tmp_tensor);
+        if (device.type == kCPU) {
+            if (pinned_weights_.find(name) == pinned_weights_.end()) {
+                pinned_weights_[name] = empty_like(*tensor_ptr, kCPUpinned);
+                Copy(*tensor_ptr, pinned_weights_[name]);
+            }
+            *tensor_ptr = {};
+        }
+        else {
+            TM_CHECK(pinned_weights_.find(name) != pinned_weights_.end());
+            *tensor_ptr = empty_like(pinned_weights_[name], kDEVICE);
+            Copy(pinned_weights_[name], *tensor_ptr);
         }
     }
     core::Context::stream().Sync();

src/turbomind/models/llama/LlamaWeight.h

@@ -20,6 +20,8 @@
 
 #pragma once
 
+#include <unordered_map>
+
 #include "src/turbomind/core/context.h"
 #include "src/turbomind/models/llama/LlamaDecoderLayerWeight.h"
 #include "src/turbomind/models/llama/LlamaDenseWeight.h"
@@ -75,6 +77,8 @@ struct LlamaWeight: core::Module {
     DataType data_type_;
     DataType weight_type_;
 
+    std::unordered_map<std::string, Tensor> pinned_weights_;
+
     int tp_size_;  // this will follow attn tp param
     int tp_rank_;