remove unused layers

2fdda5174dee5ba7fab930f62b0905389d156255 · 2024-11-27 15:48:31 +0000 · Rory Clear

Files touched

M exo/inference/tinygrad/inference.py

Diff

commit 2fdda5174dee5ba7fab930f62b0905389d156255
Author: Rory Clear <roryclear@Rorys-MacBook-Air.local>
Date:   Wed Nov 27 15:48:31 2024 +0000

    remove unused layers
---
 exo/inference/tinygrad/inference.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/exo/inference/tinygrad/inference.py b/exo/inference/tinygrad/inference.py
index afa9c5df..ad5deabe 100644
--- a/exo/inference/tinygrad/inference.py
+++ b/exo/inference/tinygrad/inference.py
@@ -13,6 +13,7 @@ from exo.download.shard_download import ShardDownloader
 from concurrent.futures import ThreadPoolExecutor
 from .stateful_model import StatefulModel
 import asyncio
+import re
 
 Tensor.no_grad = True
 # default settings
@@ -51,7 +52,9 @@ def build_transformer(model_path: Path, shard: Shard, model_size="8B", device=No
     weights = load(str(model_path), shard)
   weights = convert_from_huggingface(weights, model, MODEL_PARAMS[model_size]["args"]["n_heads"], MODEL_PARAMS[model_size]["args"]["n_kv_heads"])
   weights = fix_bf16(weights)
-
+  for k in list(weights):
+    if (n := re.search(r"\.(\d+)\.", k)) and not (shard.start_layer <= int(n.group(1)) <= shard.end_layer):
+        del weights[k]
   with Context(BEAM=0):
     # replace weights in model
     load_state_dict(model, weights, strict=False, consume=False)  # consume=True

← be82ac7d update readme to reflect that mlx and tinygrad are interoper · back to Exo · move to tinygrad_helpers 1d1fa8c6 →