I think this is more faithful to how it was originally done

aefc0d7c51c97981cdbf888c75f0e7f4677d2c89 · 2024-11-11 01:05:21 -0800 · Nel Nibcord

Files touched

M exo/orchestration/standard_node.py

Diff

commit aefc0d7c51c97981cdbf888c75f0e7f4677d2c89
Author: Nel Nibcord <blindcrone@tuta.io>
Date:   Mon Nov 11 01:05:21 2024 -0800

    I think this is more faithful to how it was originally done
---
 exo/orchestration/standard_node.py | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/exo/orchestration/standard_node.py b/exo/orchestration/standard_node.py
index 23b40719..f9d221b6 100644
--- a/exo/orchestration/standard_node.py
+++ b/exo/orchestration/standard_node.py
@@ -122,10 +122,9 @@ class StandardNode(Node):
     for i in np.reshape(result, (-1, 1, result.shape[-1])):
       self.buffered_logits[request_id][0].append(i)
 
-    inference_state = json.dumps({"start_pos": len(self.buffered_logits[request_id][0])})
-
     if shard.is_last_layer():
       result = await self.inference_engine.sample(result)
+      inference_state = json.dumps({"start_pos": len(self.buffered_logits[request_id][0]) + 1})
     
     await self.inference_engine.ensure_shard(shard)
     is_finished = result.size == 1 and result.item() == self.inference_engine.tokenizer.eos_token_id or len(self.buffered_token_output[request_id][0]) >= self.max_generate_tokens
@@ -134,7 +133,6 @@ class StandardNode(Node):
 
     if result.size == 1:  # we got a new token out
       self.buffered_token_output[request_id][0].append(result.item())
-      inference_state = json.dumps({"start_pos": json.loads(inference_state or "{}").get("start_pos", 0) + 1})
       self.trigger_on_token_callbacks(request_id, self.buffered_token_output[request_id][0], is_finished)
     
     if DEBUG >= 2: print(f"[{request_id}] result size: {result.size}, is finished: {is_finished}, buffered tokens: {len(self.buffered_token_output[request_id][0])}")

← c06b5f3b Corrected type annotations · back to Exo · Some stability improvements for tinygrad inference 13572e6a →