memory-efficient shard loading

b9c323bb07b9e258337f8c7290d3ff3a598427b0 · 2024-07-14 23:27:57 -0700 · Alex Cheema

Files touched

M exo/inference/mlx/sharded_utils.py

Diff

commit b9c323bb07b9e258337f8c7290d3ff3a598427b0
Author: Alex Cheema <alexcheema123@gmail.com>
Date:   Sun Jul 14 23:27:57 2024 -0700

    memory-efficient shard loading
---
 exo/inference/mlx/sharded_utils.py | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/exo/inference/mlx/sharded_utils.py b/exo/inference/mlx/sharded_utils.py
index c9a5fb08..e02cf579 100644
--- a/exo/inference/mlx/sharded_utils.py
+++ b/exo/inference/mlx/sharded_utils.py
@@ -107,8 +107,11 @@ def load_model_shard(
         raise FileNotFoundError(f"No safetensors found in {model_path}")
 
     weights = {}
+    all_weights_keys = set()
     for wf in weight_files:
-        weights.update(mx.load(wf))
+        weights_dict = mx.load(wf)
+        all_weights_keys.update(weights_dict.keys())
+        weights.update({k: v for k, v in weights_dict.items() if not k.startswith("model.layers.") or shard.start_layer <= int(k.split('.')[2]) <= shard.end_layer})
 
     model_class, model_args_class = _get_classes(config=config)
 
@@ -123,7 +126,7 @@ def load_model_shard(
         def class_predicate(p, m):
             if not hasattr(m, "to_quantized"):
                 return False
-            return f"{p}.scales" in weights
+            return f"{p}.scales" in all_weights_keys
 
         nn.quantize(
             model,

← 53a5b3fc add uuid requirement · back to Exo · clean debug logs bcd58938 →