← back to Exo
Removed inference state entirely
8b71d57da733ebd58bacea97033bb3446b33f5f7 · 2024-11-12 21:38:55 -0800 · Nel Nibcord
For tinygrad, this means caching start_pos in StatefulModel
Files touched
M exo/inference/dummy_inference_engine.pyM exo/inference/inference_engine.pyM exo/inference/mlx/sharded_inference_engine.pyM exo/inference/tinygrad/inference.pyM exo/inference/tinygrad/models/llama.pyM exo/inference/tinygrad/stateful_model.pyM exo/networking/grpc/grpc_peer_handle.pyM exo/networking/grpc/grpc_server.pyM exo/networking/grpc/node_service.protoM exo/networking/grpc/node_service_pb2.pyM exo/networking/peer_handle.pyM exo/orchestration/node.pyM exo/orchestration/standard_node.py
Diff
commit 8b71d57da733ebd58bacea97033bb3446b33f5f7
Author: Nel Nibcord <blindcrone@tuta.io>
Date: Tue Nov 12 21:38:55 2024 -0800
Removed inference state entirely
For tinygrad, this means caching start_pos in StatefulModel
---
exo/inference/dummy_inference_engine.py | 2 +-
exo/inference/inference_engine.py | 6 +--
exo/inference/mlx/sharded_inference_engine.py | 2 +-
exo/inference/tinygrad/inference.py | 5 +-
exo/inference/tinygrad/models/llama.py | 11 ++--
exo/inference/tinygrad/stateful_model.py | 32 +++++++-----
exo/networking/grpc/grpc_peer_handle.py | 6 +--
exo/networking/grpc/grpc_server.py | 3 +-
exo/networking/grpc/node_service.proto | 2 -
exo/networking/grpc/node_service_pb2.py | 74 +++++++++++++--------------
exo/networking/peer_handle.py | 4 +-
exo/orchestration/node.py | 4 +-
exo/orchestration/standard_node.py | 35 +++++--------
13 files changed, 91 insertions(+), 95 deletions(-)
diff --git a/exo/inference/dummy_inference_engine.py b/exo/inference/dummy_inference_engine.py
index be0cc94e..2712b7d2 100644
--- a/exo/inference/dummy_inference_engine.py
+++ b/exo/inference/dummy_inference_engine.py
@@ -28,7 +28,7 @@ class DummyInferenceEngine(InferenceEngine):
async def decode(self, shard: Shard, tokens: np.ndarray) -> str:
return ' '.join([random_string(np.random.randint(1, 34)) for token in tokens])
- async def infer_tensor(self, request_id: str, shard: Shard, input_data: np.ndarray, inference_state: Optional[str] = None) -> np.ndarray:
+ async def infer_tensor(self, request_id: str, shard: Shard, input_data: np.ndarray) -> np.ndarray:
await self.ensure_shard(shard)
sequence_length = input_data.shape[0 if self.shard.is_first_layer() else 1]
output = np.random.random(size=(1, sequence_length, self.vocab_size if self.shard.is_last_layer() else self.hidden_size))
diff --git a/exo/inference/inference_engine.py b/exo/inference/inference_engine.py
index 4349889c..f30c9707 100644
--- a/exo/inference/inference_engine.py
+++ b/exo/inference/inference_engine.py
@@ -21,12 +21,12 @@ class InferenceEngine(ABC):
pass
@abstractmethod
- async def infer_tensor(self, request_id: str, shard: Shard, input_data: np.ndarray, inference_state: Optional[str] = None) -> np.ndarray:
+ async def infer_tensor(self, request_id: str, shard: Shard, input_data: np.ndarray) -> np.ndarray:
pass
- async def infer_prompt(self, request_id: str, shard: Shard, prompt: str, inference_state: Optional[str] = None) -> np.ndarray:
+ async def infer_prompt(self, request_id: str, shard: Shard, prompt: str) -> np.ndarray:
tokens = await self.encode(shard, prompt)
- output_data = await self.infer_tensor(request_id, shard, tokens, inference_state)
+ output_data = await self.infer_tensor(request_id, shard, tokens)
return output_data
diff --git a/exo/inference/mlx/sharded_inference_engine.py b/exo/inference/mlx/sharded_inference_engine.py
index 354e4ee6..c1052233 100644
--- a/exo/inference/mlx/sharded_inference_engine.py
+++ b/exo/inference/mlx/sharded_inference_engine.py
@@ -53,7 +53,7 @@ class MLXDynamicShardInferenceEngine(InferenceEngine):
tokens = await asyncio.get_running_loop().run_in_executor(self.executor, self.tokenizer.decode, tokens)
return tokens
- async def infer_tensor(self, request_id: str, shard: Shard, input_data: np.ndarray, inference_state: Optional[str] = None) -> np.ndarray:
+ async def infer_tensor(self, request_id: str, shard: Shard, input_data: np.ndarray) -> np.ndarray:
await self.ensure_shard(shard)
output_data: np.ndarray = np.array(await asyncio.get_running_loop().run_in_executor(self.executor, self.model, mx.array(input_data), request_id))
return output_data
diff --git a/exo/inference/tinygrad/inference.py b/exo/inference/tinygrad/inference.py
index afc8f061..0c595b20 100644
--- a/exo/inference/tinygrad/inference.py
+++ b/exo/inference/tinygrad/inference.py
@@ -82,10 +82,9 @@ class TinygradDynamicShardInferenceEngine(InferenceEngine):
tokens = await asyncio.get_running_loop().run_in_executor(self.executor, self.tokenizer.decode, tokens)
return tokens
- async def infer_tensor(self, request_id: str, shard: Shard, input_data: np.ndarray, inference_state: Optional[str] = None) -> np.ndarray:
+ async def infer_tensor(self, request_id: str, shard: Shard, input_data: np.ndarray) -> np.ndarray:
await self.ensure_shard(shard)
- start_pos = json.loads(inference_state or "{}").get("start_pos", 0)
- output_data = await asyncio.get_running_loop().run_in_executor(self.executor, lambda: self.model(Tensor(input_data), start_pos, request_id).realize())
+ output_data = await asyncio.get_running_loop().run_in_executor(self.executor, lambda: self.model(Tensor(input_data), request_id).realize())
return output_data.numpy()
async def ensure_shard(self, shard: Shard):
diff --git a/exo/inference/tinygrad/models/llama.py b/exo/inference/tinygrad/models/llama.py
index 9928ab43..b1087b04 100644
--- a/exo/inference/tinygrad/models/llama.py
+++ b/exo/inference/tinygrad/models/llama.py
@@ -196,10 +196,10 @@ class Transformer:
self.output.weight = self.tok_embeddings.weight
self.max_context = max_context
self.freqs_cis = precompute_freqs_cis(dim // n_heads, self.max_context*2, rope_theta, rope_scaling=rope_scaling).contiguous()
- self.forward_jit = TinyJit(self.forward) if jit else None
+ self.forward_jit = TinyJit(self.forward_base) if jit else None
self.shard = shard
- def forward(self, x: Tensor, start_pos: Union[Variable, int], cache: Optional[List[Tensor]] = None):
+ def forward_base(self, x: Tensor, start_pos: Union[Variable, int], cache: Optional[List[Tensor]] = None):
seqlen = x.shape[1]
freqs_cis = self.freqs_cis.shrink((None, (start_pos, start_pos + seqlen), None, None, None))
mask = Tensor.full((1, 1, seqlen, start_pos + seqlen), float("-100000000"), dtype=x.dtype, device=x.device).triu(start_pos + 1).realize() if seqlen > 1 else None
@@ -225,11 +225,14 @@ class Transformer:
h = inputs
return h
+ def forward(self, x: Tensor, start_pos: Variable, cache: Optional[List[Tensor]] = None):
+ if x.shape[0:2] == (1, 1) and self.forward_jit is not None:
+ return self.forward_jit(x, Variable("start_pos", 0, self.max_context).bind(start_pos), cache=cache)
+ return self.forward_base(x, start_pos, cache=cache)
+
def __call__(self, tokens: Tensor, start_pos: Variable, cache: Optional[List[Tensor]] = None):
# TODO: better way to handle the first call v.s. the rest?
h = self.embed(x)
- if tokens.shape[0:2] == (1, 1) and self.forward_jit is not None:
- return self.forward_jit(h, Variable("start_pos", 0, self.max_context).bind(start_pos), cache=cache)
return self.forward(h, start_pos, cache=cache)
diff --git a/exo/inference/tinygrad/stateful_model.py b/exo/inference/tinygrad/stateful_model.py
index 141f671a..e9178a38 100644
--- a/exo/inference/tinygrad/stateful_model.py
+++ b/exo/inference/tinygrad/stateful_model.py
@@ -1,5 +1,6 @@
from tinygrad import Tensor, Variable
from collections import OrderedDict
+from typing import List
def create_kv_cache(x: Tensor, max_context: int, n_kv_heads: int, head_dim: int):
cache_kv = Tensor.zeros(2, x.shape[0], max_context, n_kv_heads, head_dim, dtype=x.dtype).contiguous().realize()
@@ -8,27 +9,34 @@ def create_kv_cache(x: Tensor, max_context: int, n_kv_heads: int, head_dim: int)
cache_kv.shard_((x.device), axis=3 if getenv("SHARD_KVCACHE") else None).realize()
return cache_kv.realize()
+class ModelState:
+ cache: List[Tensor]
+ start: int
+ def __init__(self, cache: List[Tensor], start: int = 0):
+ self.cache = cache
+ self.start = start
+
class StatefulModel:
- def __init__(self, model, max_caches: int = 2):
+ def __init__(self, model, max_states: int = 2):
super().__init__()
self.model = model
- self.max_caches = max_caches
- self.caches = OrderedDict()
+ self.max_states = max_states
+ self.states = OrderedDict()
def init_cache(self, x: Tensor, request_id: str):
cache = [create_kv_cache(x, self.model.layers[i].attention.max_context, self.model.layers[i].attention.n_kv_heads, self.model.layers[i].attention.head_dim) for i in range(self.model.shard.start_layer, self.model.shard.end_layer + 1)]
- if len(self.caches) >= self.max_caches:
- self.caches.popitem(last=False)
+ if len(self.states) >= self.max_states:
+ self.states.popitem(last=False)
- self.caches[request_id] = cache
+ self.states[request_id] = ModelState(cache)
- def __call__(self, x: Tensor, start_pos: Variable, request_id: str):
+ def __call__(self, x: Tensor, request_id: str):
h = self.model.embed(x)
- if request_id not in self.caches:
+ if request_id not in self.states:
self.init_cache(h, request_id)
else:
- self.caches.move_to_end(request_id)
- if h.shape[0:2] == (1, 1) and self.model.forward_jit is not None:
- return self.model.forward_jit(h, Variable("start_pos", 0, self.model.max_context).bind(start_pos), cache=self.caches[request_id])
- return self.model.forward(h, start_pos, cache=self.caches[request_id])
+ self.states.move_to_end(request_id)
+ out = self.model.forward(h, self.states[request_id].start, cache=self.states[request_id].cache)
+ self.states[request_id].start += h.shape[1]
+ return out
diff --git a/exo/networking/grpc/grpc_peer_handle.py b/exo/networking/grpc/grpc_peer_handle.py
index 11fbb48f..a788b87a 100644
--- a/exo/networking/grpc/grpc_peer_handle.py
+++ b/exo/networking/grpc/grpc_peer_handle.py
@@ -67,7 +67,7 @@ class GRPCPeerHandle(PeerHandle):
traceback.print_exc()
return False
- async def send_prompt(self, shard: Shard, prompt: str, request_id: Optional[str] = None, inference_state: Optional[str] = None) -> Optional[np.array]:
+ async def send_prompt(self, shard: Shard, prompt: str, request_id: Optional[str] = None) -> Optional[np.array]:
request = node_service_pb2.PromptRequest(
prompt=prompt,
shard=node_service_pb2.Shard(
@@ -77,7 +77,6 @@ class GRPCPeerHandle(PeerHandle):
n_layers=shard.n_layers,
),
request_id=request_id,
- inference_state=inference_state,
)
response = await self.stub.SendPrompt(request)
@@ -86,7 +85,7 @@ class GRPCPeerHandle(PeerHandle):
return np.frombuffer(response.tensor_data, dtype=np.dtype(response.dtype)).reshape(response.shape)
- async def send_tensor(self, shard: Shard, tensor: np.ndarray, request_id: Optional[str] = None, inference_state: Optional[str] = None) -> Optional[np.array]:
+ async def send_tensor(self, shard: Shard, tensor: np.ndarray, request_id: Optional[str] = None) -> Optional[np.array]:
request = node_service_pb2.TensorRequest(
shard=node_service_pb2.Shard(
model_id=shard.model_id,
@@ -96,7 +95,6 @@ class GRPCPeerHandle(PeerHandle):
),
tensor=node_service_pb2.Tensor(tensor_data=tensor.tobytes(), shape=tensor.shape, dtype=str(tensor.dtype)),
request_id=request_id,
- inference_state=inference_state,
)
response = await self.stub.SendTensor(request)
diff --git a/exo/networking/grpc/grpc_server.py b/exo/networking/grpc/grpc_server.py
index 3e4eceea..db489475 100644
--- a/exo/networking/grpc/grpc_server.py
+++ b/exo/networking/grpc/grpc_server.py
@@ -64,9 +64,8 @@ class GRPCServer(node_service_pb2_grpc.NodeServiceServicer):
)
tensor = np.frombuffer(request.tensor.tensor_data, dtype=np.dtype(request.tensor.dtype)).reshape(request.tensor.shape)
request_id = request.request_id
- inference_state = request.inference_state
- result = await self.node.process_tensor(shard, tensor, request_id, inference_state)
+ result = await self.node.process_tensor(shard, tensor, request_id)
if DEBUG >= 5: print(f"SendTensor tensor {shard=} {tensor=} {request_id=} result: {result}")
tensor_data = result.tobytes() if result is not None else None
return node_service_pb2.Tensor(tensor_data=tensor_data, shape=result.shape, dtype=str(result.dtype)) if result is not None else node_service_pb2.Tensor()
diff --git a/exo/networking/grpc/node_service.proto b/exo/networking/grpc/node_service.proto
index 03c8f8eb..146976b6 100644
--- a/exo/networking/grpc/node_service.proto
+++ b/exo/networking/grpc/node_service.proto
@@ -23,14 +23,12 @@ message PromptRequest {
Shard shard = 1;
string prompt = 2;
optional string request_id = 3;
- optional string inference_state = 4;
}
message TensorRequest {
Shard shard = 1;
Tensor tensor = 2;
optional string request_id = 3;
- optional string inference_state = 4;
}
message GetInferenceResultRequest {
diff --git a/exo/networking/grpc/node_service_pb2.py b/exo/networking/grpc/node_service_pb2.py
index 0eb22415..ab9e6bac 100644
--- a/exo/networking/grpc/node_service_pb2.py
+++ b/exo/networking/grpc/node_service_pb2.py
@@ -14,7 +14,7 @@ _sym_db = _symbol_database.Default()
-DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x12node_service.proto\x12\x0cnode_service\"S\n\x05Shard\x12\x10\n\x08model_id\x18\x01 \x01(\t\x12\x13\n\x0bstart_layer\x18\x02 \x01(\x05\x12\x11\n\tend_layer\x18\x03 \x01(\x05\x12\x10\n\x08n_layers\x18\x04 \x01(\x05\"\x9d\x01\n\rPromptRequest\x12\"\n\x05shard\x18\x01 \x01(\x0b\x32\x13.node_service.Shard\x12\x0e\n\x06prompt\x18\x02 \x01(\t\x12\x17\n\nrequest_id\x18\x03 \x01(\tH\x00\x88\x01\x01\x12\x1c\n\x0finference_state\x18\x04 \x01(\tH\x01\x88\x01\x01\x42\r\n\x0b_request_idB\x12\n\x10_inference_state\"\xb3\x01\n\rTensorRequest\x12\"\n\x05shard\x18\x01 \x01(\x0b\x32\x13.node_service.Shard\x12$\n\x06tensor\x18\x02 \x01(\x0b\x32\x14.node_service.Tensor\x12\x17\n\nrequest_id\x18\x03 \x01(\tH\x00\x88\x01\x01\x12\x1c\n\x0finference_state\x18\x04 \x01(\tH\x01\x88\x01\x01\x42\r\n\x0b_request_idB\x12\n\x10_inference_state\"/\n\x19GetInferenceResultRequest\x12\x12\n\nrequest_id\x18\x01 \x01(\t\"\\\n\x0fInferenceResult\x12)\n\x06tensor\x18\x01 \x01(\x0b\x32\x14.node_service.TensorH\x00\x88\x01\x01\x12\x13\n\x0bis_finished\x18\x02 \x01(\x08\x42\t\n\x07_tensor\";\n\x06Tensor\x12\x13\n\x0btensor_data\x18\x01 \x01(\x0c\x12\r\n\x05shape\x18\x02 \x03(\x05\x12\r\n\x05\x64type\x18\x03 \x01(\t\"<\n\x16\x43ollectTopologyRequest\x12\x0f\n\x07visited\x18\x01 \x03(\t\x12\x11\n\tmax_depth\x18\x02 \x01(\x05\"\x8e\x02\n\x08Topology\x12\x30\n\x05nodes\x18\x01 \x03(\x0b\x32!.node_service.Topology.NodesEntry\x12\x39\n\npeer_graph\x18\x02 \x03(\x0b\x32%.node_service.Topology.PeerGraphEntry\x1aN\n\nNodesEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12/\n\x05value\x18\x02 \x01(\x0b\x32 .node_service.DeviceCapabilities:\x02\x38\x01\x1a\x45\n\x0ePeerGraphEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\"\n\x05value\x18\x02 \x01(\x0b\x32\x13.node_service.Peers:\x02\x38\x01\"\x19\n\x05Peers\x12\x10\n\x08peer_ids\x18\x01 \x03(\t\"7\n\x0b\x44\x65viceFlops\x12\x0c\n\x04\x66p32\x18\x01 \x01(\x02\x12\x0c\n\x04\x66p16\x18\x02 \x01(\x02\x12\x0c\n\x04int8\x18\x03 \x01(\x02\"k\n\x12\x44\x65viceCapabilities\x12\r\n\x05model\x18\x01 \x01(\t\x12\x0c\n\x04\x63hip\x18\x02 \x01(\t\x12\x0e\n\x06memory\x18\x03 \x01(\x05\x12(\n\x05\x66lops\x18\x04 \x01(\x0b\x32\x19.node_service.DeviceFlops\"L\n\x11SendResultRequest\x12\x12\n\nrequest_id\x18\x01 \x01(\t\x12\x0e\n\x06result\x18\x02 \x03(\x05\x12\x13\n\x0bis_finished\x18\x03 \x01(\x08\"=\n\x17SendOpaqueStatusRequest\x12\x12\n\nrequest_id\x18\x01 \x01(\t\x12\x0e\n\x06status\x18\x02 \x01(\t\"\x14\n\x12HealthCheckRequest\")\n\x13HealthCheckResponse\x12\x12\n\nis_healthy\x18\x01 \x01(\x08\"\x07\n\x05\x45mpty2\xb4\x04\n\x0bNodeService\x12\x41\n\nSendPrompt\x12\x1b.node_service.PromptRequest\x1a\x14.node_service.Tensor\"\x00\x12\x41\n\nSendTensor\x12\x1b.node_service.TensorRequest\x1a\x14.node_service.Tensor\"\x00\x12^\n\x12GetInferenceResult\x12\'.node_service.GetInferenceResultRequest\x1a\x1d.node_service.InferenceResult\"\x00\x12Q\n\x0f\x43ollectTopology\x12$.node_service.CollectTopologyRequest\x1a\x16.node_service.Topology\"\x00\x12\x44\n\nSendResult\x12\x1f.node_service.SendResultRequest\x1a\x13.node_service.Empty\"\x00\x12P\n\x10SendOpaqueStatus\x12%.node_service.SendOpaqueStatusRequest\x1a\x13.node_service.Empty\"\x00\x12T\n\x0bHealthCheck\x12 .node_service.HealthCheckRequest\x1a!.node_service.HealthCheckResponse\"\x00\x62\x06proto3')
+DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x12node_service.proto\x12\x0cnode_service\"S\n\x05Shard\x12\x10\n\x08model_id\x18\x01 \x01(\t\x12\x13\n\x0bstart_layer\x18\x02 \x01(\x05\x12\x11\n\tend_layer\x18\x03 \x01(\x05\x12\x10\n\x08n_layers\x18\x04 \x01(\x05\"k\n\rPromptRequest\x12\"\n\x05shard\x18\x01 \x01(\x0b\x32\x13.node_service.Shard\x12\x0e\n\x06prompt\x18\x02 \x01(\t\x12\x17\n\nrequest_id\x18\x03 \x01(\tH\x00\x88\x01\x01\x42\r\n\x0b_request_id\"\x81\x01\n\rTensorRequest\x12\"\n\x05shard\x18\x01 \x01(\x0b\x32\x13.node_service.Shard\x12$\n\x06tensor\x18\x02 \x01(\x0b\x32\x14.node_service.Tensor\x12\x17\n\nrequest_id\x18\x03 \x01(\tH\x00\x88\x01\x01\x42\r\n\x0b_request_id\"/\n\x19GetInferenceResultRequest\x12\x12\n\nrequest_id\x18\x01 \x01(\t\"\\\n\x0fInferenceResult\x12)\n\x06tensor\x18\x01 \x01(\x0b\x32\x14.node_service.TensorH\x00\x88\x01\x01\x12\x13\n\x0bis_finished\x18\x02 \x01(\x08\x42\t\n\x07_tensor\";\n\x06Tensor\x12\x13\n\x0btensor_data\x18\x01 \x01(\x0c\x12\r\n\x05shape\x18\x02 \x03(\x05\x12\r\n\x05\x64type\x18\x03 \x01(\t\"<\n\x16\x43ollectTopologyRequest\x12\x0f\n\x07visited\x18\x01 \x03(\t\x12\x11\n\tmax_depth\x18\x02 \x01(\x05\"\x8e\x02\n\x08Topology\x12\x30\n\x05nodes\x18\x01 \x03(\x0b\x32!.node_service.Topology.NodesEntry\x12\x39\n\npeer_graph\x18\x02 \x03(\x0b\x32%.node_service.Topology.PeerGraphEntry\x1aN\n\nNodesEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12/\n\x05value\x18\x02 \x01(\x0b\x32 .node_service.DeviceCapabilities:\x02\x38\x01\x1a\x45\n\x0ePeerGraphEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\"\n\x05value\x18\x02 \x01(\x0b\x32\x13.node_service.Peers:\x02\x38\x01\"\x19\n\x05Peers\x12\x10\n\x08peer_ids\x18\x01 \x03(\t\"7\n\x0b\x44\x65viceFlops\x12\x0c\n\x04\x66p32\x18\x01 \x01(\x02\x12\x0c\n\x04\x66p16\x18\x02 \x01(\x02\x12\x0c\n\x04int8\x18\x03 \x01(\x02\"k\n\x12\x44\x65viceCapabilities\x12\r\n\x05model\x18\x01 \x01(\t\x12\x0c\n\x04\x63hip\x18\x02 \x01(\t\x12\x0e\n\x06memory\x18\x03 \x01(\x05\x12(\n\x05\x66lops\x18\x04 \x01(\x0b\x32\x19.node_service.DeviceFlops\"L\n\x11SendResultRequest\x12\x12\n\nrequest_id\x18\x01 \x01(\t\x12\x0e\n\x06result\x18\x02 \x03(\x05\x12\x13\n\x0bis_finished\x18\x03 \x01(\x08\"=\n\x17SendOpaqueStatusRequest\x12\x12\n\nrequest_id\x18\x01 \x01(\t\x12\x0e\n\x06status\x18\x02 \x01(\t\"\x14\n\x12HealthCheckRequest\")\n\x13HealthCheckResponse\x12\x12\n\nis_healthy\x18\x01 \x01(\x08\"\x07\n\x05\x45mpty2\xb4\x04\n\x0bNodeService\x12\x41\n\nSendPrompt\x12\x1b.node_service.PromptRequest\x1a\x14.node_service.Tensor\"\x00\x12\x41\n\nSendTensor\x12\x1b.node_service.TensorRequest\x1a\x14.node_service.Tensor\"\x00\x12^\n\x12GetInferenceResult\x12\'.node_service.GetInferenceResultRequest\x1a\x1d.node_service.InferenceResult\"\x00\x12Q\n\x0f\x43ollectTopology\x12$.node_service.CollectTopologyRequest\x1a\x16.node_service.Topology\"\x00\x12\x44\n\nSendResult\x12\x1f.node_service.SendResultRequest\x1a\x13.node_service.Empty\"\x00\x12P\n\x10SendOpaqueStatus\x12%.node_service.SendOpaqueStatusRequest\x1a\x13.node_service.Empty\"\x00\x12T\n\x0bHealthCheck\x12 .node_service.HealthCheckRequest\x1a!.node_service.HealthCheckResponse\"\x00\x62\x06proto3')
_globals = globals()
_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals)
@@ -27,40 +27,40 @@ if not _descriptor._USE_C_DESCRIPTORS:
_globals['_TOPOLOGY_PEERGRAPHENTRY']._serialized_options = b'8\001'
_globals['_SHARD']._serialized_start=36
_globals['_SHARD']._serialized_end=119
- _globals['_PROMPTREQUEST']._serialized_start=122
- _globals['_PROMPTREQUEST']._serialized_end=279
- _globals['_TENSORREQUEST']._serialized_start=282
- _globals['_TENSORREQUEST']._serialized_end=461
- _globals['_GETINFERENCERESULTREQUEST']._serialized_start=463
- _globals['_GETINFERENCERESULTREQUEST']._serialized_end=510
- _globals['_INFERENCERESULT']._serialized_start=512
- _globals['_INFERENCERESULT']._serialized_end=604
- _globals['_TENSOR']._serialized_start=606
- _globals['_TENSOR']._serialized_end=665
- _globals['_COLLECTTOPOLOGYREQUEST']._serialized_start=667
- _globals['_COLLECTTOPOLOGYREQUEST']._serialized_end=727
- _globals['_TOPOLOGY']._serialized_start=730
- _globals['_TOPOLOGY']._serialized_end=1000
- _globals['_TOPOLOGY_NODESENTRY']._serialized_start=851
- _globals['_TOPOLOGY_NODESENTRY']._serialized_end=929
- _globals['_TOPOLOGY_PEERGRAPHENTRY']._serialized_start=931
- _globals['_TOPOLOGY_PEERGRAPHENTRY']._serialized_end=1000
- _globals['_PEERS']._serialized_start=1002
- _globals['_PEERS']._serialized_end=1027
- _globals['_DEVICEFLOPS']._serialized_start=1029
- _globals['_DEVICEFLOPS']._serialized_end=1084
- _globals['_DEVICECAPABILITIES']._serialized_start=1086
- _globals['_DEVICECAPABILITIES']._serialized_end=1193
- _globals['_SENDRESULTREQUEST']._serialized_start=1195
- _globals['_SENDRESULTREQUEST']._serialized_end=1271
- _globals['_SENDOPAQUESTATUSREQUEST']._serialized_start=1273
- _globals['_SENDOPAQUESTATUSREQUEST']._serialized_end=1334
- _globals['_HEALTHCHECKREQUEST']._serialized_start=1336
- _globals['_HEALTHCHECKREQUEST']._serialized_end=1356
- _globals['_HEALTHCHECKRESPONSE']._serialized_start=1358
- _globals['_HEALTHCHECKRESPONSE']._serialized_end=1399
- _globals['_EMPTY']._serialized_start=1401
- _globals['_EMPTY']._serialized_end=1408
- _globals['_NODESERVICE']._serialized_start=1411
- _globals['_NODESERVICE']._serialized_end=1975
+ _globals['_PROMPTREQUEST']._serialized_start=121
+ _globals['_PROMPTREQUEST']._serialized_end=228
+ _globals['_TENSORREQUEST']._serialized_start=231
+ _globals['_TENSORREQUEST']._serialized_end=360
+ _globals['_GETINFERENCERESULTREQUEST']._serialized_start=362
+ _globals['_GETINFERENCERESULTREQUEST']._serialized_end=409
+ _globals['_INFERENCERESULT']._serialized_start=411
+ _globals['_INFERENCERESULT']._serialized_end=503
+ _globals['_TENSOR']._serialized_start=505
+ _globals['_TENSOR']._serialized_end=564
+ _globals['_COLLECTTOPOLOGYREQUEST']._serialized_start=566
+ _globals['_COLLECTTOPOLOGYREQUEST']._serialized_end=626
+ _globals['_TOPOLOGY']._serialized_start=629
+ _globals['_TOPOLOGY']._serialized_end=899
+ _globals['_TOPOLOGY_NODESENTRY']._serialized_start=750
+ _globals['_TOPOLOGY_NODESENTRY']._serialized_end=828
+ _globals['_TOPOLOGY_PEERGRAPHENTRY']._serialized_start=830
+ _globals['_TOPOLOGY_PEERGRAPHENTRY']._serialized_end=899
+ _globals['_PEERS']._serialized_start=901
+ _globals['_PEERS']._serialized_end=926
+ _globals['_DEVICEFLOPS']._serialized_start=928
+ _globals['_DEVICEFLOPS']._serialized_end=983
+ _globals['_DEVICECAPABILITIES']._serialized_start=985
+ _globals['_DEVICECAPABILITIES']._serialized_end=1092
+ _globals['_SENDRESULTREQUEST']._serialized_start=1094
+ _globals['_SENDRESULTREQUEST']._serialized_end=1170
+ _globals['_SENDOPAQUESTATUSREQUEST']._serialized_start=1172
+ _globals['_SENDOPAQUESTATUSREQUEST']._serialized_end=1233
+ _globals['_HEALTHCHECKREQUEST']._serialized_start=1235
+ _globals['_HEALTHCHECKREQUEST']._serialized_end=1255
+ _globals['_HEALTHCHECKRESPONSE']._serialized_start=1257
+ _globals['_HEALTHCHECKRESPONSE']._serialized_end=1298
+ _globals['_EMPTY']._serialized_start=1300
+ _globals['_EMPTY']._serialized_end=1307
+ _globals['_NODESERVICE']._serialized_start=1310
+ _globals['_NODESERVICE']._serialized_end=1874
# @@protoc_insertion_point(module_scope)
diff --git a/exo/networking/peer_handle.py b/exo/networking/peer_handle.py
index b5bf49b2..e5085f0b 100644
--- a/exo/networking/peer_handle.py
+++ b/exo/networking/peer_handle.py
@@ -36,11 +36,11 @@ class PeerHandle(ABC):
pass
@abstractmethod
- async def send_prompt(self, shard: Shard, prompt: str, request_id: Optional[str] = None, inference_state: Optional[str] = None) -> Optional[np.array]:
+ async def send_prompt(self, shard: Shard, prompt: str, request_id: Optional[str] = None) -> Optional[np.array]:
pass
@abstractmethod
- async def send_tensor(self, shard: Shard, tensor: np.array, request_id: Optional[str] = None, inference_state: Optional[str] = None) -> Optional[np.array]:
+ async def send_tensor(self, shard: Shard, tensor: np.array, request_id: Optional[str] = None) -> Optional[np.array]:
pass
@abstractmethod
diff --git a/exo/orchestration/node.py b/exo/orchestration/node.py
index 93b777d4..2df59bb5 100644
--- a/exo/orchestration/node.py
+++ b/exo/orchestration/node.py
@@ -16,11 +16,11 @@ class Node(ABC):
pass
@abstractmethod
- async def process_prompt(self, shard: Shard, prompt: str, request_id: Optional[str] = None, inference_state: Optional[str] = None) -> Optional[np.ndarray]:
+ async def process_prompt(self, shard: Shard, prompt: str, request_id: Optional[str] = None) -> Optional[np.ndarray]:
pass
@abstractmethod
- async def process_tensor(self, shard: Shard, tensor: np.ndarray, request_id: Optional[str] = None, inference_state: Optional[str] = None) -> Optional[np.ndarray]:
+ async def process_tensor(self, shard: Shard, tensor: np.ndarray, request_id: Optional[str] = None) -> Optional[np.ndarray]:
pass
@abstractmethod
diff --git a/exo/orchestration/standard_node.py b/exo/orchestration/standard_node.py
index c0322746..8d82af08 100644
--- a/exo/orchestration/standard_node.py
+++ b/exo/orchestration/standard_node.py
@@ -111,7 +111,6 @@ class StandardNode(Node):
shard,
result: np.ndarray,
request_id: Optional[str] = None,
- inference_state: Optional[str] = None,
):
if request_id not in self.buffered_token_output:
self.buffered_token_output[request_id] = ([], False)
@@ -123,7 +122,6 @@ class StandardNode(Node):
if shard.is_last_layer():
result = await self.inference_engine.sample(result)
- inference_state = json.dumps({"start_pos": len(self.buffered_logits[request_id]) + 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
@@ -139,7 +137,7 @@ class StandardNode(Node):
if is_finished:
self.buffered_token_output[request_id] = (self.buffered_token_output[request_id][0], True)
else:
- asyncio.create_task(self.forward_to_next_shard(shard, result, request_id, inference_state=inference_state))
+ asyncio.create_task(self.forward_to_next_shard(shard, result, request_id))
return np.array(self.buffered_token_output[request_id][0]) if len(self.buffered_token_output[request_id][0]) > 0 else None
@@ -148,7 +146,6 @@ class StandardNode(Node):
base_shard: Shard,
prompt: str,
request_id: Optional[str] = None,
- inference_state: Optional[str] = None
) -> Optional[np.ndarray]:
shard = self.get_current_shard(base_shard)
asyncio.create_task(
@@ -161,13 +158,12 @@ class StandardNode(Node):
"base_shard": base_shard.to_dict(),
"shard": shard.to_dict(),
"prompt": prompt,
- "inference_state": inference_state,
"request_id": request_id,
}),
)
)
start_time = time.perf_counter_ns()
- resp = await self._process_prompt(base_shard, prompt, request_id, inference_state)
+ resp = await self._process_prompt(base_shard, prompt, request_id)
end_time = time.perf_counter_ns()
elapsed_time_ns = end_time - start_time
asyncio.create_task(
@@ -180,7 +176,6 @@ class StandardNode(Node):
"base_shard": base_shard.to_dict(),
"shard": shard.to_dict(),
"prompt": prompt,
- "inference_state": inference_state,
"request_id": request_id,
"elapsed_time_ns": elapsed_time_ns,
"result_size": resp.size if resp is not None else 0,
@@ -189,7 +184,7 @@ class StandardNode(Node):
)
return resp
- async def _process_prompt(self, base_shard: Shard, prompt: str, request_id: Optional[str] = None, inference_state: Optional[str] = None) -> Optional[np.ndarray]:
+ async def _process_prompt(self, base_shard: Shard, prompt: str, request_id: Optional[str] = None) -> Optional[np.ndarray]:
if request_id is None:
request_id = str(uuid.uuid4())
shard = self.get_current_shard(base_shard)
@@ -197,11 +192,11 @@ class StandardNode(Node):
if DEBUG >= 2: print(f"[{request_id}] process prompt: {base_shard=} {shard=} {prompt=}")
if shard.start_layer != 0:
if DEBUG >= 2: print(f"[{request_id}] forwarding to next shard: {base_shard=} {shard=} {prompt=}")
- await self.forward_to_next_shard(shard, prompt, request_id, inference_state=inference_state)
+ await self.forward_to_next_shard(shard, prompt, request_id)
return None
else:
- result = await self.inference_engine.infer_prompt(request_id, shard, prompt, inference_state=inference_state)
- ret = await self.process_result(shard, result, request_id, inference_state=inference_state)
+ result = await self.inference_engine.infer_prompt(request_id, shard, prompt)
+ ret = await self.process_result(shard, result, request_id)
return result
async def process_tensor(
@@ -209,7 +204,6 @@ class StandardNode(Node):
base_shard: Shard,
tensor: np.ndarray,
request_id: Optional[str] = None,
- inference_state: Optional[str] = None,
) -> Optional[np.ndarray]:
shard = self.get_current_shard(base_shard)
asyncio.create_task(
@@ -224,12 +218,11 @@ class StandardNode(Node):
"tensor_size": tensor.size,
"tensor_shape": tensor.shape,
"request_id": request_id,
- "inference_state": inference_state,
}),
)
)
start_time = time.perf_counter_ns()
- resp = await self._process_tensor(shard, tensor, request_id, inference_state)
+ resp = await self._process_tensor(shard, tensor, request_id)
end_time = time.perf_counter_ns()
elapsed_time_ns = end_time - start_time
asyncio.create_task(
@@ -254,7 +247,6 @@ class StandardNode(Node):
base_shard: Shard,
tensor: np.ndarray,
request_id: Optional[str] = None,
- inference_state: Optional[str] = None,
) -> Optional[np.ndarray]:
if request_id is None:
request_id = str(uuid.uuid4())
@@ -262,8 +254,8 @@ class StandardNode(Node):
if DEBUG >= 1: print(f"[{request_id}] process_tensor: {tensor.size=} {tensor.shape=}")
try:
- result = await self.inference_engine.infer_tensor(request_id, shard, tensor, inference_state=inference_state)
- ret = await self.process_result(shard, result, request_id, inference_state=inference_state)
+ result = await self.inference_engine.infer_tensor(request_id, shard, tensor)
+ ret = await self.process_result(shard, result, request_id)
return ret
except Exception as e:
print(f"Error processing tensor for shard {shard}: {e}")
@@ -275,7 +267,6 @@ class StandardNode(Node):
base_shard: Shard,
tensor_or_prompt: Union[np.ndarray, str],
request_id: str,
- inference_state: Optional[str] = None,
) -> None:
if not self.partitioning_strategy:
if DEBUG >= 1: print("No partitioning strategy found. Skipping forward.")
@@ -290,18 +281,18 @@ class StandardNode(Node):
is_tensor = isinstance(tensor_or_prompt, np.ndarray)
if target_id == self.id:
if is_tensor:
- await self.process_tensor(next_shard, tensor_or_prompt, request_id, inference_state=inference_state)
+ await self.process_tensor(next_shard, tensor_or_prompt, request_id)
else:
- await self.process_prompt(next_shard, tensor_or_prompt, request_id, inference_state=inference_state)
+ await self.process_prompt(next_shard, tensor_or_prompt, request_id)
else:
target_peer = next((p for p in self.peers if p.id() == target_id), None)
if not target_peer:
raise ValueError(f"Peer for {next_partition_index} not found")
if is_tensor:
if DEBUG >= 1: print(f"Sending tensor to {target_peer.id()}: {tensor_or_prompt}")
- await target_peer.send_tensor(next_shard, tensor_or_prompt, request_id=request_id, inference_state=inference_state)
+ await target_peer.send_tensor(next_shard, tensor_or_prompt, request_id=request_id)
else:
- await target_peer.send_prompt(next_shard, tensor_or_prompt, request_id=request_id, inference_state=inference_state)
+ await target_peer.send_prompt(next_shard, tensor_or_prompt, request_id=request_id)
def get_partition_index(self, offset: int = 0):
partitions = self.partitioning_strategy.partition(self.topology)
← f67e18a7 adding option to expand error to see stack trace and clearin
·
back to Exo
·
Updated unit tests 42172b2c →