Support distributed inference in the server (ml-explore#741)

Author: angeloskath

mlx_lm/generate.py

@@ -1014,10 +1014,16 @@ def insert(
         )
         return uids
 
-    def remove(self, uids: List[int]):
+    def remove(self, uids: List[int], return_prompt_caches: bool = False):
+        caches = {}
         uids = set(uids)
         if self.active_batch is not None:
             batch = self.active_batch
+            if return_prompt_caches:
+                for e, uid in enumerate(batch.uids):
+                    if uid not in uids:
+                        continue
+                    caches[uid] = batch.extract_cache(e)
             keep_idx = [e for e, uid in enumerate(batch.uids) if uid not in uids]
             if len(keep_idx) > 0:
                 batch.filter(keep_idx)
@@ -1028,6 +1034,9 @@ def remove(self, uids: List[int]):
             if self.unprocessed_prompts[i][0] in uids:
                 self.unprocessed_prompts.pop(i)
 
+        if return_prompt_caches:
+            return caches
+
     def _process_prompts(self, prompts):
         uids, inputs, max_tokens, caches, samplers, logits_processors = zip(*prompts)
 

mlx_lm/server.py

@@ -4,6 +4,7 @@
 import copy
 import json
 import logging
+import pickle
 import platform
 import socket
 import time
@@ -15,7 +16,7 @@
 from pathlib import Path
 from queue import Empty as QueueEmpty
 from queue import Queue
-from threading import Condition, Lock, Thread
+from threading import Thread
 from typing import (
     Any,
     Callable,
@@ -33,14 +34,14 @@
 from huggingface_hub import scan_cache_dir
 
 from ._version import __version__
-from .generate import BatchGenerator, stream_generate
+from .generate import BatchGenerator, generation_stream, stream_generate
 from .models.cache import (
     can_trim_prompt_cache,
     make_prompt_cache,
     trim_prompt_cache,
 )
 from .sample_utils import make_logits_processors, make_sampler
-from .utils import load
+from .utils import load, sharded_load
 
 
 def get_system_fingerprint():
@@ -384,6 +385,46 @@ class Response:
     top_tokens: Optional[Tuple[int, float]]
 
 
+class TimeBudget:
+    def __init__(self, budget=0.5, iterations=25, sync_frequency=10):
+        self._is_distributed = mx.distributed.init().size() > 1
+        self._budget = budget
+        self._iterations = iterations
+        self._sync_frequency = sync_frequency
+
+        self._start = None
+        self._current_iterations = None
+        self._loops = 0
+        self._time_spent = 0
+
+    def __iter__(self):
+        self._start = time.time()
+        self._current_iterations = 0
+        return self
+
+    def __next__(self):
+        if not self._is_distributed:
+            if time.time() - self._start > self._budget:
+                raise StopIteration()
+            return None
+
+        self._current_iterations += 1
+        if self._current_iterations > self._iterations:
+            self._loops += 1
+            self._time_spent += time.time() - self._start
+            if self._loops % self._sync_frequency == 0:
+                with mx.stream(generation_stream):
+                    loop_time = mx.distributed.all_sum(self._time_spent).item()
+                avg_loop_time = loop_time / (
+                    mx.distributed.init().size() * self._sync_frequency
+                )
+                factor = self._budget / avg_loop_time
+                self._iterations = max(round(self._iterations * factor), 1)
+                self._loops = 0
+                self._time_spent = 0
+            raise StopIteration()
+
+
 class ModelProvider:
     def __init__(self, cli_args: argparse.Namespace):
         """Load models on demand and persist them across the whole process."""
@@ -394,6 +435,13 @@ def __init__(self, cli_args: argparse.Namespace):
         self.draft_model = None
         self.is_batchable = False
 
+        group = mx.distributed.init()
+        self.pipeline_group = group if group.size() > 1 and cli_args.pipeline else None
+        self.tensor_group = (
+            group if group.size() > 1 and not cli_args.pipeline else None
+        )
+        self.is_distributed = group.size() > 1
+
         # Preload the default model if it is provided
         self.default_model_map = {}
         if self.cli_args.model is not None:
@@ -426,15 +474,29 @@ def load(self, model_path, adapter_path=None, draft_model_path=None):
                     "argument or in the HTTP request"
                 )
             adapter_path = adapter_path or self.cli_args.adapter_path
-            model, tokenizer = load(
-                self.cli_args.model,
-                adapter_path=adapter_path,
-                tokenizer_config=tokenizer_config,
-            )
+            # TODO: Generalize distributed load
+            if self.is_distributed:
+                model, tokenizer = sharded_load(
+                    self.cli_args.model, self.pipeline_group, self.tensor_group
+                )
+            else:
+                model, tokenizer = load(
+                    self.cli_args.model,
+                    adapter_path=adapter_path,
+                    tokenizer_config=tokenizer_config,
+                )
         else:
-            model, tokenizer = load(
-                model_path, adapter_path=adapter_path, tokenizer_config=tokenizer_config
-            )
+            # TODO: Generalize distributed load
+            if self.is_distributed:
+                model, tokenizer = sharded_load(
+                    model_path, self.pipeline_group, self.tensor_group
+                )
+            else:
+                model, tokenizer = load(
+                    model_path,
+                    adapter_path=adapter_path,
+                    tokenizer_config=tokenizer_config,
+                )
 
         if self.cli_args.use_default_chat_template:
             if tokenizer.chat_template is None:
@@ -501,6 +563,9 @@ def __init__(self, model_provider: ModelProvider, prompt_cache: LRUPromptCache):
         self.prompt_cache = prompt_cache
         self.requests = Queue()
 
+        self._time_budget = TimeBudget()
+        self._is_distributed = mx.distributed.init().size() > 1
+        self._rank = mx.distributed.init().rank()
         self._stop = False
         self._generation_thread = Thread(target=self._generate)
         self._generation_thread.start()
@@ -509,6 +574,57 @@ def stop_and_join(self):
         self._stop = True
         self._generation_thread.join()
 
+    def join(self):
+        self._generation_thread.join()
+
+    def _next_request(self, timeout=None):
+        request = None
+        if not self._is_distributed or self._rank == 0:
+            try:
+                if timeout is not None:
+                    request = self.requests.get(timeout=timeout)
+                else:
+                    request = self.requests.get_nowait()
+            except QueueEmpty:
+                pass
+
+        return self._share_request(request)
+
+    def _share_object(self, obj):
+        if not self._is_distributed:
+            return obj
+
+        with mx.stream(generation_stream):
+            if self._rank == 0:
+                if obj is None:
+                    mx.eval(mx.distributed.all_sum(0))
+                    return None
+                else:
+                    data = mx.array(pickle.dumps(obj))
+                    mx.eval(mx.distributed.all_sum(data.size))
+                    mx.eval(mx.distributed.all_sum(data))
+                    return obj
+            else:
+                size = mx.distributed.all_sum(0).item()
+                if size == 0:
+                    return None
+                else:
+                    data = mx.zeros(size, dtype=mx.uint8)
+                    data = mx.distributed.all_sum(data)
+                    return pickle.loads(data)
+
+    def _share_request(self, request):
+        if not self._is_distributed:
+            return request
+
+        shareable = request[1:] if request is not None else None
+        shareable = self._share_object(shareable)
+        if shareable is None:
+            return None
+
+        rq = request[0] if request is not None else Queue()
+        return rq, *shareable
+
     def _tokenize(self, tokenizer, request):
         if request.request_type == "chat":
             messages = request.messages
@@ -559,19 +675,17 @@ def get_next_request(timeout=None):
             if unprocessed_requests:
                 return unprocessed_requests.pop()
             else:
-                try:
-                    if timeout is not None:
-                        return self.requests.get(timeout=timeout)
-                    else:
-                        return self.requests.get_nowait()
-                except QueueEmpty:
-                    return None
+                return self._next_request(timeout)
 
         def progress_callback(info):
             for uid, processed, total in info:
                 if uid in batch_results:
                     batch_results[uid]["rqueue"].put((min(processed, total), total))
 
+        if self._is_distributed:
+            seed = mx.distributed.all_sum(mx.random.state[0]).view(mx.uint64).item()
+            mx.random.seed(seed)
+
         while not self._stop:
             request = None
             if not drain_batch:
@@ -657,6 +771,8 @@ def progress_callback(info):
                     batch_generator = BatchGenerator(
                         model,
                         stop_tokens=tokenizer.eos_token_ids,
+                        completion_batch_size=self.cli_args.decode_concurrency,
+                        prefill_batch_size=self.cli_args.prompt_concurrency,
                         prompt_progress_callback=progress_callback,
                     )
                     unprocessed_requests.append((rqueue, request, args))
@@ -683,12 +799,7 @@ def progress_callback(info):
                     continue
 
                 uids_to_remove = []
-                time_budget = 0.5
-                start = time.time()
-                while True:
-                    if time.time() - start > time_budget:
-                        break
-
+                for _ in self._time_budget:
                     responses = batch_generator.next()
                     if not responses:
                         break
@@ -730,8 +841,20 @@ def progress_callback(info):
                         if result["ctx"]._should_stop:
                             uids_to_remove.append(r.uid)
 
-                    if uids_to_remove:
-                        batch_generator.remove(uids_to_remove)
+                uids_to_remove = self._share_object(uids_to_remove)
+                if uids_to_remove:
+                    with mx.stream(generation_stream):
+                        caches = batch_generator.remove(
+                            uids_to_remove, return_prompt_caches=True
+                        )
+                        for uid, prompt_cache in caches.items():
+                            if uid not in batch_results:
+                                continue
+                            result = batch_results[uid]
+                            self.prompt_cache.insert_cache(
+                                current_model_key, result["cache_key"], prompt_cache
+                            )
+                            del batch_results[uid]
 
     def _serve_single(self, request):
         rqueue, request, args = request
@@ -822,6 +945,8 @@ def progress(tokens_processed, tokens_total):
                 cache_key.append(gen.token)
 
                 if ctx._should_stop:
+                    if self._is_distributed:
+                        raise NotImplementedError()
                     break
 
             rqueue.put(None)
@@ -1522,15 +1647,14 @@ def probably_mlx_lm(repo):
         self.wfile.flush()
 
 
-def run(
+def _run_http_server(
     host: str,
     port: int,
-    model_provider: ModelProvider,
+    response_generator,
     server_class=ThreadingHTTPServer,
     handler_class=APIHandler,
 ):
     server_address = (host, port)
-    response_generator = ResponseGenerator(model_provider, LRUPromptCache())
     infos = socket.getaddrinfo(
         *server_address, type=socket.SOCK_STREAM, flags=socket.AI_PASSIVE
     )
@@ -1556,6 +1680,21 @@ def run(
         response_generator.stop_and_join()
 
 
+def run(
+    host: str,
+    port: int,
+    model_provider: ModelProvider,
+    server_class=ThreadingHTTPServer,
+    handler_class=APIHandler,
+):
+    group = mx.distributed.init()
+    response_generator = ResponseGenerator(model_provider, LRUPromptCache())
+    if group.rank() == 0:
+        _run_http_server(host, port, response_generator)
+    else:
+        response_generator.join()
+
+
 def main():
     parser = argparse.ArgumentParser(description="MLX Http Server.")
     parser.add_argument(
@@ -1652,6 +1791,23 @@ def main():
         help="""A JSON formatted string of arguments for the tokenizer's apply_chat_template, e.g. '{"enable_thinking":false}'""",
         default="{}",
     )
+    parser.add_argument(
+        "--decode-concurrency",
+        type=int,
+        default=32,
+        help="When a request is batchable then decode that many requests in parallel",
+    )
+    parser.add_argument(
+        "--prompt-concurrency",
+        type=int,
+        default=8,
+        help="When a request is batchable then process that many prompts in parallel",
+    )
+    parser.add_argument(
+        "--pipeline",
+        action="store_true",
+        help="Use pipelining instead of tensor parallelism",
+    )
     args = parser.parse_args()
     if mx.metal.is_available():
         wired_limit = mx.device_info()["max_recommended_working_set_size"]

tests/test_server.py

@@ -41,6 +41,8 @@ def __init__(self, with_draft=False):
                 "max_tokens": 512,
                 "chat_template_args": {},
                 "model": None,
+                "decode_concurrency": 32,
+                "prompt_concurrency": 8,
             },
         )