Applied patch idea from https://github.com/exo-explore/exo/issues/458

89815b16b25f35c2f376547130aaadd51d1eb914 · 2024-12-06 22:26:12 -0600 · Will Bickford

Files touched

M exo/inference/tinygrad/models/llama.py

Diff

commit 89815b16b25f35c2f376547130aaadd51d1eb914
Author: Will Bickford <wbic16@gmail.com>
Date:   Fri Dec 6 22:26:12 2024 -0600

    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()}

← 1d5a7c63 pretty name for llama 3.3 70b · back to Exo · add special case for USB adapter over ethernet 6d09b4ae →