add copymap to engine

Author: RunningLeon

lmdeploy/pytorch/engine/cache_engine.py

@@ -243,7 +243,16 @@ def swap_out(self, src_to_dst: Dict[int, int]) -> None:
             src_to_dst (Dict[int, int]): Map between src and dst.
         """
         self._swap(self.full_gpu_cache, self.full_cpu_cache, src_to_dst)
-
+        
+    def copy_to(self, src_to_dst: Dict[int, int], cache_type: str = 'gpu') -> None:
+        """Copy cache.
+        Args:
+            src_to_dst (Dict[int, int]): Map between src and dst.
+            cache_type (str): cache type 'cpu', 'gpu'
+        """
+        target_cache = self.full_gpu_cache if cache_type == 'gpu' else self.full_cpu_cache
+        self._swap(target_cache, target_cache, src_to_dst)
+        
     @classmethod
     def get_cache_block_size(cls,
                              block_size: int,

lmdeploy/pytorch/engine/engine.py

@@ -798,6 +798,7 @@ def __make_dummy_inputs():
                 inputs=ModelInputs.make_dummy(1, is_decoding=not prefill),
                 swap_in_map=dict(),
                 swap_out_map=dict(),
+                copy_map=dict(),
                 loop_count=num_loops,
                 is_dummy=True,
                 sync_long_context=False,
@@ -826,6 +827,7 @@ def __make_dummy_inputs():
         running = scheduler_output.running
         swap_in_map = scheduler_output.swap_in_map
         swap_out_map = scheduler_output.swap_out_map
+        copy_map = scheduler_output.copy_map
 
         if self.should_execute_dummy_batch and len(running) == 0:
             return __make_dummy_inputs()
@@ -847,6 +849,7 @@ def __make_dummy_inputs():
             inputs=inputs,
             swap_in_map=swap_in_map,
             swap_out_map=swap_out_map,
+            copy_map=copy_map,
             loop_count=num_loops,
             all_ids=all_ids,
             guided_input_ids=guided_input_ids,

lmdeploy/pytorch/engine/model_agent.py

@@ -28,7 +28,7 @@ def msg_with_rank(rank: int, msg: str):
     return f'rank[{rank}] - {msg}'
 
 
-def cache_swapping(cache_engine: CacheEngine, swap_in_map: dict, swap_out_map: dict):
+def cache_swapping(cache_engine: CacheEngine, swap_in_map: dict, swap_out_map: dict, copy_map: dict):
     """perform cache swapping."""
     issued_cache_op = False
     if len(swap_in_map) > 0:
@@ -37,7 +37,9 @@ def cache_swapping(cache_engine: CacheEngine, swap_in_map: dict, swap_out_map: d
     if len(swap_out_map) > 0:
         cache_engine.swap_out(swap_out_map)
         issued_cache_op = True
-
+    if len(copy_map) > 0:
+        cache_engine.copy_to(copy_map)
+        issued_cache_op = True
     if issued_cache_op:
         cache_engine.events.wait()
 
@@ -135,7 +137,7 @@ def all_context(self):
     def _forward_impl(self, inputs: ModelInputs, swap_in_map: SwapMap, swap_out_map: SwapMap):
         raise NotImplementedError('NotImplemented.')
 
-    async def async_forward(self, inputs: ModelInputs, swap_in_map: SwapMap, swap_out_map: SwapMap):
+    async def async_forward(self, inputs: ModelInputs, swap_in_map: SwapMap, swap_out_map: SwapMap, copy_map: SwapMap):
         """model forward.
 
         Args:
@@ -200,6 +202,7 @@ async def _async_model_forward(
         inputs: ModelInputs,
         swap_in_map: Dict,
         swap_out_map: Dict,
+        copy_map: Dict,
         return_logits: bool,
         sync_long_context: bool,
     ):
@@ -241,12 +244,12 @@ def get_output(self):
 
         async def __forward(inputs):
             """forward."""
-            nonlocal swap_done, swap_in_map, swap_out_map
+            nonlocal swap_done, swap_in_map, swap_out_map, copy_map
             if swap_done:
-                return await self.async_forward(inputs, swap_in_map=dict(), swap_out_map=dict())
+                return await self.async_forward(inputs, swap_in_map=dict(), swap_out_map=dict(), copy_map=dict())
             else:
                 swap_done = True
-                return await self.async_forward(inputs, swap_in_map=swap_in_map, swap_out_map=swap_out_map)
+                return await self.async_forward(inputs, swap_in_map=swap_in_map, swap_out_map=swap_out_map, copy_map=copy_map)
 
         async def __long_context_single_forward(new_inputs, max_seqlen: int):
             """one large sequence."""
@@ -334,6 +337,7 @@ async def _async_step_background(
         inputs: ModelInputs,
         swap_in_map: Dict,
         swap_out_map: Dict,
+        copy_map: Dict,
         loop_count: int,
         all_ids: torch.Tensor = None,
         guided_input_ids: torch.Tensor = None,
@@ -420,6 +424,7 @@ async def __await_distworker(worker, timeout: float = 0.001):
                 inputs,
                 swap_in_map=swap_in_map,
                 swap_out_map=swap_out_map,
+                copy_map=copy_map,
                 return_logits=return_logits,
                 sync_long_context=sync_long_context,
             )
@@ -467,6 +472,7 @@ async def __await_distworker(worker, timeout: float = 0.001):
             if is_decoding and idx < loop_count - 1:
                 swap_in_map = dict()
                 swap_out_map = dict()
+                copy_map = dict()
                 inputs.model_metas = model_metas
                 __update_inputs(next_token_ids)
 
@@ -637,8 +643,8 @@ def build_cache_engine(self):
 
             self.cache_engine = CacheEngine(self.cache_config, self.model_config, world_size=tp)
 
-    def _forward_impl(self, inputs: ModelInputs, swap_in_map: SwapMap, swap_out_map: SwapMap):
-        cache_swapping(self.cache_engine, swap_in_map=swap_in_map, swap_out_map=swap_out_map)
+    def _forward_impl(self, inputs: ModelInputs, swap_in_map: SwapMap, swap_out_map: SwapMap, copy_map: SwapMap):
+        cache_swapping(self.cache_engine, swap_in_map=swap_in_map, swap_out_map=swap_out_map, copy_map=copy_map)
         output = model_forward(
             self.patched_model,
             inputs,
@@ -647,15 +653,15 @@ def _forward_impl(self, inputs: ModelInputs, swap_in_map: SwapMap, swap_out_map:
         )
         return output
 
-    async def async_forward(self, inputs: ModelInputs, swap_in_map: SwapMap, swap_out_map: SwapMap):
+    async def async_forward(self, inputs: ModelInputs, swap_in_map: SwapMap, swap_out_map: SwapMap, copy_map: SwapMap):
         """model forward.
 
         Args:
             inputs (Dict): The input data comes from _make_inputs.
             swap_in_map (SwapMap): Cache maps to swap in.
             swap_out_map (SwapMap): Cache maps to swap out.
         """
-        output = self._forward_impl(inputs, swap_in_map=swap_in_map, swap_out_map=swap_out_map)
+        output = self._forward_impl(inputs, swap_in_map=swap_in_map, swap_out_map=swap_out_map, copy_map=copy_map)
         await asyncio.sleep(0)
         return output