float 16

exo/inference/tinygrad/models/llama.py

@@ -252,8 +252,6 @@ def permute(v: Tensor, n_heads: int):
 
 
 def fix_bf16(weights: Dict[Any, Tensor]):
-    for k, v in weights.items():
-        print(f"Key: {k}, Device: {v.device}, Dtype: {v.dtype}")
 
     if Device.DEFAULT == "CLANG":
         return {
@@ -263,3 +261,14 @@ def fix_bf16(weights: Dict[Any, Tensor]):
         k: v.to(dtypes.float16).to(v.device) if v.dtype == dtypes.bfloat16 else v for k, v in weights.items()
     }
 
+def fix_bf16(weights):
+    converted = {}
+    for k, v in weights.items():
+        if v.dtype == dtypes.bfloat16:
+            curr_device = v.device
+            converted_tensor = v.to(dtypes.float16)
+            device_str = getattr(curr_device, 'name', str(curr_device))
+            converted[k] = converted_tensor.to(device_str)
+        else:
+            converted[k] = v
+    return converted

exo/models.py

@@ -6,8 +6,7 @@
     "MLXDynamicShardInferenceEngine": Shard(model_id="mlx-community/Llama-3.2-1B-Instruct-4bit", start_layer=0, end_layer=0, n_layers=16),
   },
   "llama-3.2-3b": {
-    "MLXDynamicShardInferenceEngine": Shard(model_id="mlx-community/Llama-3.2-3B-Instruct-4bit", start_layer=0, end_layer=0, n_layers=28),
-    "TinygradDynamicShardInferenceEngine":Shard(model_id="unsloth/Llama-3.2-3B-Instruct", start_layer=0, end_layer=0, n_layers=28)
+    "MLXDynamicShardInferenceEngine": Shard(model_id="mlx-community/Llama-3.2-3B-Instruct-4bit", start_layer=0, end_layer=0, n_layers=28)
   },
   "llama-3.1-8b": {
     "MLXDynamicShardInferenceEngine": Shard(model_id="mlx-community/Meta-Llama-3.1-8B-Instruct-4bit", start_layer=0, end_layer=0, n_layers=32),