[object Object]

← back to Exo

Fixing tinygrad model

bfa3b36be51b433b76f9cbe62bea398dbf88b54f · 2024-12-06 00:24:04 -0800 · Nel Nibcord

Files touched

Diff

commit bfa3b36be51b433b76f9cbe62bea398dbf88b54f
Author: Nel Nibcord <blindcrone@tuta.io>
Date:   Fri Dec 6 00:24:04 2024 -0800

    Fixing tinygrad model
---
 exo/inference/tinygrad/inference.py | 16 ++++++++++++----
 1 file changed, 12 insertions(+), 4 deletions(-)

diff --git a/exo/inference/tinygrad/inference.py b/exo/inference/tinygrad/inference.py
index 7e1f0664..5f7d54d3 100644
--- a/exo/inference/tinygrad/inference.py
+++ b/exo/inference/tinygrad/inference.py
@@ -67,8 +67,12 @@ class TinygradDynamicShardInferenceEngine(InferenceEngine):
     self.session = {}
 
   async def poll_state(self, request_id: str, max_caches=2):
-    if len(self.states) >= self.max_states:
-      self.states.popitem(last=False)
+    if request_id not in self.states:
+      if len(self.states) >= self.max_states:
+        self.states.popitem(last=False)
+      make_prompt_state(
+    else:
+      self.states.move_to_end(request_id)
     state = self.state[request_id]
     return {"start_pos": state.start, "cache": state.cache}
 
@@ -91,8 +95,12 @@ class TinygradDynamicShardInferenceEngine(InferenceEngine):
   async def infer_tensor(self, request_id: str, shard: Shard, input_data: np.ndarray) -> np.ndarray:
     await self.ensure_shard(shard)
     state = self.poll_state(request_id)
-    output_data = await asyncio.get_running_loop().run_in_executor(self.executor, lambda: self.model(Tensor(input_data), **state).realize())
-    self.state[request_id].start += input_data.shape[1]
+    def wrap_infer(data):
+      x = Tensor(data)
+      state = poll_state(request_id, x)
+      self.model(x, **state)
+      self.state[request_id].start += x.shape[1]
+    output_data = await asyncio.get_running_loop().run_in_executor(self.executor, lambda: wrap_infer(input_data).realize())
     return output_data.numpy()
 
   async def evaluate(self, request_id: str, shard: Shard, inputs, targets, lengths, loss=length_masked_ce_loss):

← 37a75d6b Fixing tinygrad model  ·  back to Exo  ·  Fixing tinygrad model 67f5ae25 →