Thread local generation stream (#1090)

Author: angeloskath

mlx_lm/generate.py

@@ -223,7 +223,7 @@ def setup_arg_parser():
 
 
 # A stream on the default device just for generation
-generation_stream = mx.new_stream(mx.default_device())
+generation_stream = mx.new_thread_local_stream(mx.default_device())
 
 
 @contextlib.contextmanager
@@ -1497,6 +1497,7 @@ class BatchGenerator:
     def __init__(
         self,
         model: nn.Module,
+        *,
         max_tokens: int = 128,
         stop_tokens: Optional[Sequence[Sequence[int]]] = None,
         sampler: Optional[Callable[[mx.array], mx.array]] = None,
@@ -1507,6 +1508,7 @@ def __init__(
         prefill_batch_size: int = 8,
         prefill_step_size: int = 2048,
         max_kv_size: Optional[int] = None,
+        stream=None,
     ):
         self.model = model
         self.max_tokens = max_tokens
@@ -1518,6 +1520,8 @@ def __init__(
         self.completion_batch_size = max(completion_batch_size, prefill_batch_size)
         self.max_kv_size = max_kv_size
 
+        self._stream = stream or generation_stream
+
         self._default_state_machine = SequenceStateMachine(
             {"normal": [(seq, None) for seq in stop_tokens]} if stop_tokens else {},
             initial="normal",
@@ -1544,9 +1548,13 @@ def __init__(
         else:
             self._old_wired_limit = None
 
+    @property
+    def stream(self):
+        return self._stream
+
     def close(self):
         if self._old_wired_limit is not None:
-            mx.synchronize(generation_stream)
+            mx.synchronize(self._stream)
             mx.set_wired_limit(self._old_wired_limit)
             self._old_wired_limit = None
 
@@ -1843,7 +1851,7 @@ def next(self):
         Returns:
             Tuple of prompt processing responses and generation responses.
         """
-        with mx.stream(generation_stream):
+        with mx.stream(self._stream):
             return self._next()
 
     def next_generated(self):
@@ -1853,7 +1861,7 @@ def next_generated(self):
         Returns:
             List of GenerationBatch.Response objects
         """
-        with mx.stream(generation_stream):
+        with mx.stream(self._stream):
             while True:
                 prompt_responses, generation_responses = self._next()
                 if not generation_responses and prompt_responses:

mlx_lm/server.py

@@ -36,7 +36,6 @@
 from .generate import (
     BatchGenerator,
     SequenceStateMachine,
-    generation_stream,
     stream_generate,
 )
 from .models.cache import (
@@ -279,8 +278,7 @@ def __next__(self):
         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()
+            loop_time = mx.distributed.all_sum(self._time_spent).item()
             avg_loop_time = loop_time / (
                 mx.distributed.init().size() * self._sync_frequency
             )
@@ -308,94 +306,92 @@ def __init__(self, cli_args: argparse.Namespace):
         )
         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:
-            self.default_model_map[self.cli_args.model] = "default_model"
-            self.load(self.cli_args.model, draft_model_path="default_model")
+        # Maps model and adapter paths the actual paths to be used. Used to
+        # map 'default_model' to the provided model by cli argument but could
+        # be used for more in the future.
+        self._model_map = {}
+        self._adapter_map = {}
+        self._draft_model_map = {}
+        self._model_map["default_model"] = self.cli_args.model
+        self._adapter_map["default_model"] = self.cli_args.adapter_path
+        self._draft_model_map["default_model"] = self.cli_args.draft_model
+
+        # Build the tokenizer config for later use in load
+        self._tokenizer_config = {
+            "trust_remote_code": True if cli_args.trust_remote_code else None
+        }
+        if cli_args.chat_template:
+            self._tokenizer_config["chat_template"] = cli_args.chat_template
 
-    # Added in adapter_path to load dynamically
-    def load(self, model_path, adapter_path=None, draft_model_path=None):
-        model_path = self.default_model_map.get(model_path, model_path)
-        if self.model_key == (model_path, adapter_path, draft_model_path):
-            return self.model, self.tokenizer
+    def _load(self, model_path, adapter_path=None, draft_model_path=None):
+        if self.is_distributed and (
+            adapter_path is not None or draft_model_path is not None
+        ):
+            raise ValueError(
+                "Loading with adapters or draft models not supported in distributed mode"
+            )
 
         # Remove the old model if it exists.
+        self.model_key = None
         self.model = None
         self.tokenizer = None
-        self.model_key = None
         self.draft_model = None
 
-        # Building tokenizer_config
-        tokenizer_config = {
-            "trust_remote_code": True if self.cli_args.trust_remote_code else None
-        }
-        if self.cli_args.chat_template:
-            tokenizer_config["chat_template"] = self.cli_args.chat_template
-
-        if model_path == "default_model":
-            if self.cli_args.model is None:
-                raise ValueError(
-                    "A model path has to be given as a CLI "
-                    "argument or in the HTTP request"
-                )
-            adapter_path = adapter_path or self.cli_args.adapter_path
-            # 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,
-                )
+        # Load the model and tokenizer
+        if self.is_distributed:
+            model, tokenizer = sharded_load(
+                model_path,
+                pipeline_group=self.pipeline_group,
+                tensor_group=self.tensor_group,
+                tokenizer_config=self._tokenizer_config,
+            )
         else:
-            # 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,
-                )
+            model, tokenizer = load(
+                model_path,
+                adapter_path=adapter_path,
+                tokenizer_config=self._tokenizer_config,
+            )
 
+        # Use the default chat template if needed
         if self.cli_args.use_default_chat_template:
             if tokenizer.chat_template is None:
                 tokenizer.chat_template = tokenizer.default_chat_template
 
-        self.model_key = (model_path, adapter_path, draft_model_path)
-        self.model = model
-        self.tokenizer = tokenizer
-
-        def validate_draft_tokenizer(draft_tokenizer):
-            # Check if tokenizers are compatible
+        # Load the draft model for speculative decoding
+        draft_model = None
+        if draft_model_path is not None:
+            draft_model, draft_tokenizer = load(draft_model_path)
             if draft_tokenizer.vocab_size != tokenizer.vocab_size:
                 logging.warning(
                     "Draft model tokenizer does not match model tokenizer. "
                     "Speculative decoding may not work as expected."
                 )
 
-        # Load draft model if specified
-        if (
-            draft_model_path == "default_model"
-            and self.cli_args.draft_model is not None
-        ):
-            self.draft_model, draft_tokenizer = load(self.cli_args.draft_model)
-            validate_draft_tokenizer(draft_tokenizer)
+        # Compute batchability
+        is_batchable = draft_model is None
+        is_batchable = is_batchable and all(
+            hasattr(c, "merge") for c in make_prompt_cache(model)
+        )
+
+        # Update the member variables
+        self.model_key = (model_path, adapter_path, draft_model_path)
+        self.model = model
+        self.tokenizer = tokenizer
+        self.draft_model = draft_model
+        self.is_batchable = is_batchable
 
-        elif draft_model_path is not None and draft_model_path != "default_model":
-            self.draft_model, draft_tokenizer = load(draft_model_path)
-            validate_draft_tokenizer(draft_tokenizer)
+    def load_default(self):
+        if self._model_map["default_model"] is not None:
+            self.load("default_model", None, "default_model")
 
-        if self.draft_model is None:
-            self.is_batchable = all(
-                hasattr(c, "merge") for c in make_prompt_cache(self.model)
-            )
+    def load(self, model_path, adapter_path=None, draft_model_path=None):
+        model_path = self._model_map.get(model_path, model_path)
+        adapter_path = self._adapter_map.get(model_path, adapter_path)
+        draft_model_path = self._draft_model_map.get(draft_model_path, draft_model_path)
+
+        model_key = (model_path, adapter_path, draft_model_path)
+        if self.model_key != model_key:
+            self._load(*model_key)
 
         return self.model, self.tokenizer
 
@@ -489,22 +485,21 @@ 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
-                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
-                data = mx.zeros(size, dtype=mx.uint8)
-                data = mx.distributed.all_sum(data)
-                return pickle.loads(data)
+        if self._rank == 0:
+            if obj is None:
+                mx.eval(mx.distributed.all_sum(0))
+                return None
+            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
+            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:
@@ -691,6 +686,14 @@ def _is_batchable(self, args):
         return self.model_provider.is_batchable and args.seed is None
 
     def _generate(self):
+        # Local thread stream that we 'll pass to the BatchGenerator to make
+        # sure that all generation runs in the same stream as the
+        # synchronization messages.
+        generation_stream = mx.default_stream(mx.default_device())
+
+        # Load the default model if it is given
+        self.model_provider.load_default()
+
         current_model = None
         current_sampling = None
         current_tokenizer = None
@@ -820,6 +823,7 @@ def get_next_request(timeout=None):
                         completion_batch_size=self.cli_args.decode_concurrency,
                         prefill_batch_size=self.cli_args.prompt_concurrency,
                         prefill_step_size=self.cli_args.prefill_step_size,
+                        stream=generation_stream,
                     )
                     unprocessed_requests.append((rqueue, request, args))
                     continue
@@ -909,12 +913,11 @@ def get_next_request(timeout=None):
 
                 uids_to_remove = self._share_object(uids_to_remove)
                 if uids_to_remove:
-                    with mx.stream(generation_stream):
-                        batch_generator.remove(uids_to_remove)
-                        for uid in uids_to_remove:
-                            # It may have already been removed during
-                            # generation
-                            batch_results.pop(uid, None)
+                    batch_generator.remove(uids_to_remove)
+                    for uid in uids_to_remove:
+                        # It may have already been removed during
+                        # generation
+                        batch_results.pop(uid, None)
 
     def _serve_single(self, request):
         rqueue, request, args = request

mlx_lm/utils.py

@@ -507,6 +507,8 @@ def sharded_load(
     pipeline_group: Optional[mx.distributed.Group] = None,
     tensor_group: Optional[mx.distributed.Group] = None,
     return_config: bool = False,
+    *,
+    tokenizer_config: Optional[Dict[str, Any]] = None,
 ):
     # Get model path with everything but weight safetensors
     model_path = _download(
@@ -571,7 +573,7 @@ def sharded_load(
     # Load and shard the model, and load the weights
     tokenizer = load_tokenizer(
         model_path,
-        {"trust_remote_code": True},
+        tokenizer_config or {"trust_remote_code": True},
         eos_token_ids=config.get("eos_token_id", None),
     )
     model, _ = load_model(model_path, lazy=True, strict=False)

setup.py

@@ -10,7 +10,7 @@
 
 from _version import __version__
 
-MIN_MLX_VERSION = "0.30.4"
+MIN_MLX_VERSION = "0.31.2"
 
 setup(
     name="mlx-lm",

tests/test_gguf.py

@@ -35,6 +35,9 @@ def test_convert_to_gguf(
         mock_tokenizer.get_vocab.return_value = {"<pad>": 0, "hello": 1, "world": 2}
         mock_tokenizer.all_special_tokens = ["<pad>"]
         mock_tokenizer.all_special_ids = [0]
+        mock_tokenizer.bos_token_id = None
+        mock_tokenizer.eos_token_id = None
+        mock_tokenizer.unk_token_id = None
         mock_from_pretrained.return_value = mock_tokenizer
 
         model_path = Path(self.test_dir)

tests/test_server.py

@@ -68,6 +68,9 @@ def load(self, model, adapter=None, draft_model=None):
         assert model in ["default_model", "chat_model"]
         return self.model, self.tokenizer
 
+    def load_default(self):
+        return self.load("default_model", None, "default_model")
+
 
 class MockCache:
     def __init__(self, value, is_trimmable: bool = True):