load mlx model shard on mlx thread so it doesnt block

6662d5668cdf77afcd50ea4c7acca6ad5db79d80 · 2025-01-28 18:49:19 +0000 · Alex Cheema

Files touched

M exo/inference/mlx/sharded_inference_engine.py

Diff

commit 6662d5668cdf77afcd50ea4c7acca6ad5db79d80
Author: Alex Cheema <alexcheema123@gmail.com>
Date:   Tue Jan 28 18:49:19 2025 +0000

    load mlx model shard on mlx thread so it doesnt block
---
 exo/inference/mlx/sharded_inference_engine.py | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/exo/inference/mlx/sharded_inference_engine.py b/exo/inference/mlx/sharded_inference_engine.py
index 9a0d28a2..54073c55 100644
--- a/exo/inference/mlx/sharded_inference_engine.py
+++ b/exo/inference/mlx/sharded_inference_engine.py
@@ -4,7 +4,7 @@ import mlx.nn as nn
 from mlx_lm.sample_utils import top_p_sampling, make_sampler
 import mlx.optimizers as optim
 from ..inference_engine import InferenceEngine
-from .sharded_utils import load_shard, get_image_from_str
+from .sharded_utils import load_shard, load_model_shard, resolve_tokenizer
 from .losses import loss_fns
 from ..shard import Shard
 from typing import Dict, Optional, Tuple
@@ -157,7 +157,11 @@ class MLXDynamicShardInferenceEngine(InferenceEngine):
       return
     model_path = await self.shard_downloader.ensure_shard(shard, self.__class__.__name__)
     if self.shard != shard:
-      model_shard, self.tokenizer = await load_shard(model_path, shard)
+      model_shard = await asyncio.get_running_loop().run_in_executor(self._mlx_thread, lambda: load_model_shard(model_path, shard, lazy=False))
+      if hasattr(model_shard, "tokenizer"):
+        self.tokenizer = model_shard.tokenizer
+      else:
+        self.tokenizer = await resolve_tokenizer(model_path)
       self.shard = shard
       self.model = model_shard
       self.caches = OrderedDict()

← 7c649085 fix eta/speed for resuming an existing download, using the s · back to Exo · make sure mlx stuff is on separate thread non blocking 4a5b80a9 →