fix legacy model loading

eafed8e16e3c4e2a7b0a3078279ffda4c42f3e43 · 2024-08-02 14:18:31 +0100 · Alex Cheema

Files touched

M exo/inference/mlx/models/llama.py
M exo/inference/mlx/sharded_utils.py
M exo/inference/test_inference_engine.py

Diff

commit eafed8e16e3c4e2a7b0a3078279ffda4c42f3e43
Author: Alex Cheema <alexcheema123@gmail.com>
Date:   Fri Aug 2 14:18:31 2024 +0100

    fix legacy model loading
---
 exo/inference/mlx/models/llama.py      | 6 +++++-
 exo/inference/mlx/sharded_utils.py     | 7 ++++++-
 exo/inference/test_inference_engine.py | 2 +-
 3 files changed, 12 insertions(+), 3 deletions(-)

diff --git a/exo/inference/mlx/models/llama.py b/exo/inference/mlx/models/llama.py
index fa07be66..29320cec 100644
--- a/exo/inference/mlx/models/llama.py
+++ b/exo/inference/mlx/models/llama.py
@@ -106,7 +106,11 @@ class Model(nn.Module):
           shard_state_dict[key] = value
       elif self.args.shard.is_first_layer() and key.startswith('model.embed_tokens'):
         shard_state_dict[key] = value
-      elif self.args.shard.is_last_layer() and (key.startswith('model.norm') or key.startswith('lm_head')):
+      elif (self.args.shard.is_last_layer() and self.args.tie_word_embeddings) and key.startswith('model.embed_tokens'):
+        shard_state_dict[key] = value
+      elif (self.args.shard.is_last_layer() and not self.args.tie_word_embeddings) and key.startswith('lm_head'):
+        shard_state_dict[key] = value
+      elif self.args.shard.is_last_layer() and (key.startswith('model.norm')):
         shard_state_dict[key] = value
 
     return shard_state_dict
diff --git a/exo/inference/mlx/sharded_utils.py b/exo/inference/mlx/sharded_utils.py
index 316e5cd0..115e8d86 100644
--- a/exo/inference/mlx/sharded_utils.py
+++ b/exo/inference/mlx/sharded_utils.py
@@ -144,10 +144,15 @@ def load_model_shard(
     weights = model.sanitize(weights)
 
   if (quantization := config.get("quantization", None)) is not None:
+    # Handle legacy models which may not have everything quantized
+    def class_predicate(p, m):
+        if not hasattr(m, "to_quantized"):
+            return False
+        return f"{p}.scales" in weights
     nn.quantize(
       model,
       **quantization,
-      class_predicate=None,
+      class_predicate=class_predicate,
     )
 
   model.load_weights(list(weights.items()), strict=True)
diff --git a/exo/inference/test_inference_engine.py b/exo/inference/test_inference_engine.py
index 1c4707dc..b5021330 100644
--- a/exo/inference/test_inference_engine.py
+++ b/exo/inference/test_inference_engine.py
@@ -44,7 +44,7 @@ asyncio.run(
   test_inference_engine(
     MLXDynamicShardInferenceEngine(),
     MLXDynamicShardInferenceEngine(),
-    "mlx-community/Meta-Llama-3.1-8B-Instruct-4bit",
+    "mlx-community/Meta-Llama-3-8B-Instruct-4bit",
   )
 )

← 08e8cacf run integration test for each inference engine · back to Exo · run tinygrad with llama-3-8b 0fd6bd91 →