update tinygrad

9f6c688d62448fa9f9f7666918fad118b5c3a555 · 2025-01-29 18:06:38 +0000 · Alex Cheema

Files touched

M exo/inference/tinygrad/inference.py
M setup.py

Diff

commit 9f6c688d62448fa9f9f7666918fad118b5c3a555
Author: Alex Cheema <alexcheema123@gmail.com>
Date:   Wed Jan 29 18:06:38 2025 +0000

    update tinygrad
---
 exo/inference/tinygrad/inference.py | 6 +++---
 setup.py                            | 2 +-
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/exo/inference/tinygrad/inference.py b/exo/inference/tinygrad/inference.py
index 6543f0b8..49b52105 100644
--- a/exo/inference/tinygrad/inference.py
+++ b/exo/inference/tinygrad/inference.py
@@ -79,8 +79,8 @@ class TinygradDynamicShardInferenceEngine(InferenceEngine):
     return {"start_pos": state.start, "cache": state.cache}
 
   async def sample(self, x: np.ndarray, temp=TEMPERATURE, top_p: float = 0.0) -> np.ndarray:
-    logits = x[:, -1, :]
     def sample_wrapper():
+      logits = x[:, -1, :]
       return sample_logits(Tensor(logits).flatten(), temp, 0, 0.8, top_p, 0.0).realize().numpy().astype(int)
     return await asyncio.get_running_loop().run_in_executor(self.executor, sample_wrapper)
 
@@ -112,9 +112,9 @@ class TinygradDynamicShardInferenceEngine(InferenceEngine):
       state = self.poll_state(h, request_id)
       out = self.model.forward(h, **state)
       self.states[request_id].start += x.shape[1]
-      return out.realize()
+      return out.numpy()
     output_data = await asyncio.get_running_loop().run_in_executor(self.executor, wrap_infer)
-    return output_data.numpy(), inference_state
+    return output_data, inference_state
 
   async def evaluate(self, request_id: str, shard: Shard, inputs, targets, lengths, loss=length_masked_ce_loss):
     def step(x, y, l):
diff --git a/setup.py b/setup.py
index 54e78776..a158d443 100644
--- a/setup.py
+++ b/setup.py
@@ -29,7 +29,7 @@ install_requires = [
   "transformers==4.46.3",
   "uuid==1.30",
   "uvloop==0.21.0",
-  "tinygrad @ git+https://github.com/tinygrad/tinygrad.git@3b26e51fcebfc6576f4e0f99693e6f1406d61d79",
+  "tinygrad @ git+https://github.com/tinygrad/tinygrad.git@ec120ce6b9ce8e4ff4b5692566a683ef240e8bc8",
 ]
 
 extras_require = {

← 4887be51 parallelise model loading · back to Exo · make a singleton thread pool executor for tinygrad since we 31b56e86 →