From 89815b16b25f35c2f376547130aaadd51d1eb914 Mon Sep 17 00:00:00 2001
From: Will Bickford <wbic16@gmail.com>
Date: Fri, 6 Dec 2024 22:26:12 -0600
Subject: [PATCH 1/2] Applied patch idea from
 https://github.com/exo-explore/exo/issues/458

---
 exo/inference/tinygrad/models/llama.py | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/exo/inference/tinygrad/models/llama.py b/exo/inference/tinygrad/models/llama.py
index 52ed3072..f8a0794c 100644
--- a/exo/inference/tinygrad/models/llama.py
+++ b/exo/inference/tinygrad/models/llama.py
@@ -275,6 +275,12 @@ def convert_from_huggingface(weights: Dict[str, Tensor], model: Transformer, n_h
 
 
 def fix_bf16(weights: Dict[Any, Tensor]):
+  if Device.DEFAULT == "CLANG":
+    # TODO: without casting to float16, 70B llama OOM on tinybox.
+    return {
+      k: (v.llvm_bf16_cast(dtypes.float32).to(v.device) if v.dtype == dtypes.bfloat16 else v) 
+      for k, v in weights.items()
+    }
   if getenv("SUPPORT_BF16", 1):
     # TODO: without casting to float16, 70B llama OOM on tinybox.
     return {k: v.cast(dtypes.float16) if v.dtype == dtypes.bfloat16 else v for k, v in weights.items()}

From a0bada3b2af443b3007d1d4e82e64aafb87f7807 Mon Sep 17 00:00:00 2001
From: Alex Cheema <alexcheema123@gmail.com>
Date: Thu, 12 Dec 2024 17:13:34 +0000
Subject: [PATCH 2/2] add llama-3.2-1b-8bit, llama-3.2-3b-8bit,
 llama-3.2-3b-bf16

---
 exo/models.py | 24 ++++++++++++++++++++++++
 1 file changed, 24 insertions(+)

diff --git a/exo/models.py b/exo/models.py
index 0d4c6e3e..0f984d48 100644
--- a/exo/models.py
+++ b/exo/models.py
@@ -17,7 +17,28 @@ model_cards = {
       "TinygradDynamicShardInferenceEngine": "unsloth/Llama-3.2-1B-Instruct",
     },
   },
+  "llama-3.2-1b-8bit": {
+    "layers": 16,
+    "repo": {
+      "MLXDynamicShardInferenceEngine": "mlx-community/Llama-3.2-1B-Instruct-8bit",
+      "TinygradDynamicShardInferenceEngine": "unsloth/Llama-3.2-1B-Instruct",
+    },
+  },
   "llama-3.2-3b": {
+    "layers": 28,
+    "repo": {
+       "MLXDynamicShardInferenceEngine": "mlx-community/Llama-3.2-3B-Instruct-4bit",
+       "TinygradDynamicShardInferenceEngine": "unsloth/Llama-3.2-3B-Instruct",
+    },
+  },
+  "llama-3.2-3b-8bit": {
+    "layers": 28,
+    "repo": {
+       "MLXDynamicShardInferenceEngine": "mlx-community/Llama-3.2-3B-Instruct-8bit",
+       "TinygradDynamicShardInferenceEngine": "unsloth/Llama-3.2-3B-Instruct",
+    },
+  },
+  "llama-3.2-3b-bf16": {
     "layers": 28,
     "repo": {
        "MLXDynamicShardInferenceEngine": "mlx-community/Llama-3.2-3B-Instruct",
@@ -94,7 +115,10 @@ model_cards = {
 pretty_name = {
   "llama-3.3-70b": "Llama 3.3 70B",
   "llama-3.2-1b": "Llama 3.2 1B",
+  "llama-3.2-1b-8bit": "Llama 3.2 1B (8-bit)",
   "llama-3.2-3b": "Llama 3.2 3B",
+  "llama-3.2-3b-8bit": "Llama 3.2 3B (8-bit)",
+  "llama-3.2-3b-bf16": "Llama 3.2 3B (BF16)",
   "llama-3.1-8b": "Llama 3.1 8B",
   "llama-3.1-70b": "Llama 3.1 70B",
   "llama-3.1-70b-bf16": "Llama 3.1 70B (BF16)",