add llama-3.1-70b-bf16 model option

2948a834488118a26a1646f5ca40d5a4679eef64 · 2024-09-05 16:16:18 +0100 · Alex Cheema

Files touched

M exo/inference/tokenizers.py
M exo/models.py

Diff

commit 2948a834488118a26a1646f5ca40d5a4679eef64
Author: Alex Cheema <alexcheema123@gmail.com>
Date:   Thu Sep 5 16:16:18 2024 +0100

    add llama-3.1-70b-bf16 model option
---
 exo/inference/tokenizers.py | 2 +-
 exo/models.py               | 4 ++++
 2 files changed, 5 insertions(+), 1 deletion(-)

diff --git a/exo/inference/tokenizers.py b/exo/inference/tokenizers.py
index 7e9a5c8e..bf4b5a6c 100644
--- a/exo/inference/tokenizers.py
+++ b/exo/inference/tokenizers.py
@@ -11,7 +11,7 @@ async def resolve_tokenizer(model_id: str):
   local_path = await get_local_snapshot_dir(model_id)
   if DEBUG >= 2: print(f"Checking if local path exists to load tokenizer from local {local_path=}")
   try:
-    if await aios.path.exists(local_path):
+    if local_path and await aios.path.exists(local_path):
       if DEBUG >= 2: print(f"Resolving tokenizer for {model_id=} from {local_path=}")
       return await _resolve_tokenizer(local_path)
   except:
diff --git a/exo/models.py b/exo/models.py
index 1c3816ac..605689b2 100644
--- a/exo/models.py
+++ b/exo/models.py
@@ -10,6 +10,10 @@ model_base_shards = {
     "MLXDynamicShardInferenceEngine": Shard(model_id="mlx-community/Meta-Llama-3.1-70B-Instruct-4bit", start_layer=0, end_layer=0, n_layers=80),
     "TinygradDynamicShardInferenceEngine": Shard(model_id="NousResearch/Meta-Llama-3.1-70B-Instruct", start_layer=0, end_layer=0, n_layers=80),
   },
+  "llama-3.1-70b-bf16": {
+    "MLXDynamicShardInferenceEngine": Shard(model_id="mlx-community/Meta-Llama-3.1-70B-Instruct-bf16", start_layer=0, end_layer=0, n_layers=80),
+    "TinygradDynamicShardInferenceEngine": Shard(model_id="NousResearch/Meta-Llama-3.1-70B-Instruct", start_layer=0, end_layer=0, n_layers=80),
+  },
   "llama-3.1-405b": {"MLXDynamicShardInferenceEngine": Shard(model_id="mlx-community/Meta-Llama-3.1-405B-4bit", start_layer=0, end_layer=0, n_layers=126),},
   "llama-3-8b": {
     "MLXDynamicShardInferenceEngine": Shard(model_id="mlx-community/Meta-Llama-3-8B-Instruct-4bit", start_layer=0, end_layer=0, n_layers=32),

← 11dd952d use set for shard specific patterns · back to Exo · add tinychat option for llama-3.1-70b-bf16 a1a0ffac →