Add boilerplate for sharing MM/deepstack embeddings across TP model

Author: turboderp

exllamav3/model/model_tp.py

@@ -12,6 +12,7 @@
 from .model_tp_fn import *
 import uuid
 from ..util import log_tp, global_t0
+from ..tokenizer.mm_embedding import send_embeddings
 
 cleanupper = Cleanupper()
 
@@ -341,7 +342,6 @@ def unload_tp(self):
         cleanupper.unregister_atexit(self.destroy_tp_context)
 
 
-
     def prepare_inputs_for_tp(self, x: torch.Tensor, params: dict) -> torch.Tensor:
         self.tp_producer.clear()
         # Use ID of Cache object as reference to avoid having to pickle it
@@ -353,11 +353,15 @@ def prepare_inputs_for_tp(self, x: torch.Tensor, params: dict) -> torch.Tensor:
             "cache_seqlens",
             "positions",
             "position_ids",
-            # "indexed_embeddings",
         ]:
             p = params.get(tensor_param)
             if p is not None:
                 params[tensor_param] = self.tp_producer.send(p)
+
+        p = params.get("indexed_embeddings")
+        if p is not None:
+            params["indexed_embeddings"] = send_embeddings(self.tp_producer, p)
+
         return self.tp_producer.send(x)
 
 

exllamav3/model/model_tp_fn.py

@@ -4,6 +4,7 @@
 from ..ext import exllamav3_ext as ext
 from functools import lru_cache
 from .model_tp_backend import TPBackendNCCL, TPBackendNative
+from ..tokenizer.mm_embedding import recv_embeddings
 from ..util import log_tp, set_t0
 
 def init_pg(device: int, active_devices: list[int], output_device: int, backend_args: dict, master: bool = False):
@@ -191,12 +192,15 @@ def mp_model_forward(
         "cache_seqlens",
         "positions",
         "position_ids",
-        # "indexed_embeddings",
     ]:
         p = params.get(tensor_param)
         if p is not None:
             params[tensor_param] = consumer.recv(p, cuda = True)
 
+    p = params.get("indexed_embeddings")
+    if p is not None:
+        params["indexed_embeddings"] = recv_embeddings(consumer, p)
+
     params["backend"] = backend
 
     x = consumer.recv(shared_input)

exllamav3/model/model_tp_shared.py

@@ -6,6 +6,7 @@
 from .model_tp_cuda import cuda_host_register, cuda_host_unregister, CUDA_HOST_REGISTER_PORTABLE
 
 DEFAULT_BUFFER_SIZE = 2 * 1024 ** 3
+MAX_CACHE_PER_PROCESS = 4 * 1024**3
 
 _torch_dtypes = {
     "torch.uint8": torch.uint8,
@@ -37,13 +38,17 @@ def __init__(
         # Pre-touch buffer to avoid page faults later
         self.buf[: self.buffer_size: 4096] = 0
 
+        # Cache
+        self.cached_cpu_tensors = {}
+        self.cache_size = 0
+
     def export(self):
         return {
             "shm_name": self.shm_name,
             "buffer_size": self.buffer_size,
         }
 
-    def send(self, tensor: torch.Tensor | None) -> dict:
+    def send(self, tensor: torch.Tensor | None, cache_id: int = None) -> dict:
 
         # None tensor
         if tensor is None:
@@ -74,13 +79,30 @@ def send(self, tensor: torch.Tensor | None) -> dict:
         dst = np.ndarray((nbytes,), dtype = np.uint8, buffer = self.shm.buf, offset = offset)
         np.copyto(dst, src, casting = "no")
 
+        # Cache
+        if nbytes > MAX_CACHE_PER_PROCESS:
+            cache_id = None
+
+        if cache_id is not None:
+            if cache_id in self.cached_cpu_tensors:
+                # print("sending cache ref:", cache_id)
+                return {
+                    "method": "cached",
+                    "cache_id": cache_id,
+                }
+            while self.cache_size + nbytes > MAX_CACHE_PER_PROCESS:
+                self.cached_cpu_tensors.pop(next(iter(self.cached_cpu_tensors)))
+            self.cached_cpu_tensors[cache_id] = tensor
+            # print("caching send:", cache_id)
+
         # Data is now buffered in shared memory space, store metadata and offset
         return {
             "method": "buffer",
             "offset": offset,
             "nbytes": nbytes,
             "dtype": str(tensor.dtype),
             "shape": tuple(tensor.shape),
+            "cache_id": cache_id
         }
 
     def clear(self):
@@ -142,6 +164,11 @@ def get_local_tensor(shm_buf, _buffer_size):
                 cuda_host_register(self.arena.data_ptr(), self.arena.numel(), flags = CUDA_HOST_REGISTER_PORTABLE)
                 self.producer.buf_is_pinned = True
 
+        # Cache
+        self.cached_cpu_tensors = {}
+        self.cache_size = 0
+
+
     def recv(
         self,
         imp: dict,
@@ -158,6 +185,13 @@ def recv(
         if imp["method"] == "none_tensor":
             return None
 
+        # Send was cached
+        cache_id = imp["cache_id"]
+        if imp["method"] == "cached":
+            # print("receiving cached:", cache_id)
+            assert not cuda, "Cannot share cached tensor for CUDA"
+            return self.cached_cpu_tensors[imp["cache_id"]]
+
         # Fallback method
         if imp["method"] == "share_memory":
             tensor = imp["shared_tensor"]
@@ -169,6 +203,12 @@ def recv(
             dtype = _torch_dtypes[imp["dtype"]]
             shape = imp["shape"]
             tensor = self.arena.narrow(0, offset, nbytes).view(dtype).view(shape)
+            if cache_id is not None:
+                # print("caching recv:", cache_id)
+                assert not cuda, "Cannot share cached tensor for CUDA"
+                while self.cache_size + nbytes > MAX_CACHE_PER_PROCESS:
+                    self.cached_cpu_tensors.pop(next(iter(self.cached_cpu_tensors)))
+                self.cached_cpu_tensors[cache_id] = tensor.clone(memory_format = torch.contiguous_format)
 
         # Slice before cloning
         if slice_dim is not None:
@@ -182,7 +222,7 @@ def recv(
                 copy = True,
                 memory_format = torch.contiguous_format
             )
-        else:
+        elif imp["method"] != "share_memory" or not tensor.is_contiguous():
             tensor = tensor.clone(memory_format = torch.contiguous_format)
 
         return tensor

exllamav3/tokenizer/mm_embedding.py

@@ -28,12 +28,13 @@ class MMEmbedding:
 
     def __init__(
         self,
-        embeddings: torch.Tensor,
-        token_string: torch.Tensor,
+        embeddings: torch.Tensor | None = None,
+        token_string: torch.Tensor | None = None,
         text_alias: str | None = None,
         deepstack_embeddings: list[torch.Tensor] | None = None,
         grid_thw: tuple | None = None,
-        mrope_merge_size: int | None = None
+        mrope_merge_size: int | None = None,
+        imp: dict | None = None
     ):
         """
         :param embeddings:
@@ -46,6 +47,21 @@ def __init__(
             Text string to represent this embedding for tokenizing
         """
 
+        if imp:
+            self.metadata = imp["metadata"]
+            self.full_length = imp["full_length"]
+            self.mm_length = imp["mm_length"]
+            self.first_index = imp["first_index"]
+            self.last_index = imp["last_index"]
+            self.text_alias = imp["text_alias"]
+            self.grid_thw = imp["grid_thw"]
+            self.mrope_merge_size = imp["mrope_merge_size"]
+            self.embeddings = imp["embeddings"]
+            self.deepstack_embeddings = imp["deepstack_embeddings"]
+            self.token_string = None
+            self.token_list = None
+            return
+
         global global_allocator
 
         if deepstack_embeddings is not None:
@@ -65,8 +81,45 @@ def __init__(
         self.grid_thw = grid_thw
         self.mrope_merge_size = mrope_merge_size
 
+        # not exported for TP
         r = torch.arange(self.first_index, self.first_index + self.mm_length, dtype = torch.long)
         m = (token_string == -1)
         token_string.masked_scatter_(m, r)
         self.token_string = token_string
-        self.token_list = token_string[0].tolist()
+        self.token_list = token_string[0].tolist()
+
+
+def send_embeddings(producer, ies: list[MMEmbedding]):
+    return {
+        "method": "list",
+        "data": [
+            {
+                "metadata": ie.metadata,
+                "full_length": ie.full_length,
+                "mm_length": ie.mm_length,
+                "first_index": ie.first_index,
+                "last_index": ie.last_index,
+                "text_alias": ie.text_alias,
+                "grid_thw": ie.grid_thw,
+                "mrope_merge_size": ie.mrope_merge_size,
+                "embeddings": producer.send(ie.embeddings, cache_id = id(ie.embeddings)),
+                "deepstack_embeddings": [
+                    producer.send(dse, cache_id = id(dse))
+                    for dse in ie.deepstack_embeddings
+                ] if ie.deepstack_embeddings is not None else None
+            }
+            for ie in ies
+        ]
+    }
+
+
+def recv_embeddings(consumer, recv) -> list[MMEmbedding]:
+    result = []
+    assert recv.get("method") == "list", "Consumer expected list"
+    for imp in recv["data"]:
+        imp["embeddings"] = consumer.recv(imp["embeddings"])
+        imp["deepstack_embeddings"] = [
+            consumer.recv(dse) for dse in imp["deepstack_embeddings"]
+        ] if imp.get("deepstack_embeddings") else None
+        result.append(MMEmbedding(imp = imp))
+    return result