a generic test for every inference engine

ca6095c04d37f879ff4e13ee9c086d0a55db202c · 2024-07-13 18:25:26 -0700 · Alex Cheema

Files touched

M inference/mlx/models/sharded_llama.py
A inference/test_inference_engine.py

Diff

commit ca6095c04d37f879ff4e13ee9c086d0a55db202c
Author: Alex Cheema <alexcheema123@gmail.com>
Date:   Sat Jul 13 18:25:26 2024 -0700

    a generic test for every inference engine
---
 inference/mlx/models/sharded_llama.py |  1 -
 inference/test_inference_engine.py    | 21 +++++++++++++++++++++
 2 files changed, 21 insertions(+), 1 deletion(-)

diff --git a/inference/mlx/models/sharded_llama.py b/inference/mlx/models/sharded_llama.py
index c2c851b7..267d177e 100644
--- a/inference/mlx/models/sharded_llama.py
+++ b/inference/mlx/models/sharded_llama.py
@@ -241,4 +241,3 @@ class Model(nn.Module):
     @property
     def n_kv_heads(self):
         return self.args.num_key_value_heads
-
diff --git a/inference/test_inference_engine.py b/inference/test_inference_engine.py
new file mode 100644
index 00000000..ea406a09
--- /dev/null
+++ b/inference/test_inference_engine.py
@@ -0,0 +1,21 @@
+from inference.mlx.sharded_inference_engine import MLXDynamicShardInferenceEngine
+from inference.inference_engine import InferenceEngine
+from inference.shard import Shard
+import numpy as np
+
+# An inference engine should work the same for any number of Shards, as long as the Shards are continuous.
+async def test_inference_engine(inference_engine: InferenceEngine, model_id: str, input_data: np.array):
+    resp_full, _ = await inference_engine.infer_tensor(shard=Shard(model_id=model_id, start_layer=0, end_layer=1, n_layers=2), input_data=input_data)
+
+    resp1, _ = await inference_engine.infer_tensor(shard=Shard(model_id=model_id, start_layer=0, end_layer=0, n_layers=2), input_data=input_data)
+    resp2, _ = await inference_engine.infer_tensor(shard=Shard(model_id=model_id, start_layer=1, end_layer=1, n_layers=2), input_data=resp1)
+
+    assert np.array_equal(resp_full, resp2)
+
+import asyncio
+
+asyncio.run(test_inference_engine(
+    MLXDynamicShardInferenceEngine(),
+    "mlx-community/Meta-Llama-3-8B-Instruct-4bit",
+    [1234]
+))

← 850b72d3 make StatefulShardedModel callable, add some tests for mlx s · back to Exo · graceful node shutdown 7077652c →