diff --git a/packages/node/meshnet_node/model_backend.py b/packages/node/meshnet_node/model_backend.py index b04000d..823a4ef 100644 --- a/packages/node/meshnet_node/model_backend.py +++ b/packages/node/meshnet_node/model_backend.py @@ -169,40 +169,97 @@ class TorchModelShard: token_id = int(self.torch.argmax(logits[:, -1, :], dim=-1)[0].item()) return self.tokenizer.decode([token_id], skip_special_tokens=True) - def generate_text(self, prompt: str, max_new_tokens: int = 16) -> str: - """Generate text locally when this process owns the full model.""" + def generate_text( + self, + messages: list[dict], + max_new_tokens: int = 256, + temperature: float = 1.0, + top_p: float = 1.0, + ) -> str: + """Autoregressive generation using HF generate() — single-node (head+tail) mode.""" if not self.is_head or not self.is_tail: raise ModelBackendError("local generation requires a full head+tail shard") - if hasattr(self.tokenizer, "apply_chat_template"): - try: - encoded = self.tokenizer.apply_chat_template( - [{"role": "user", "content": prompt}], - add_generation_prompt=True, - return_tensors="pt", - return_dict=True, - ) - except Exception: - encoded = self.tokenizer(prompt, return_tensors="pt") - else: - encoded = self.tokenizer(prompt, return_tensors="pt") + encoded = self._encode_messages(messages) input_ids = encoded["input_ids"].to(self.device) attention_mask = encoded.get("attention_mask") if attention_mask is not None: attention_mask = attention_mask.to(self.device) - pad_token_id = getattr(self.tokenizer, "pad_token_id", None) - if pad_token_id is None: - pad_token_id = getattr(self.tokenizer, "eos_token_id", None) + pad_token_id = getattr(self.tokenizer, "pad_token_id", None) or getattr(self.tokenizer, "eos_token_id", None) + do_sample = temperature != 1.0 or top_p != 1.0 with self.torch.inference_mode(): generated = self.model.generate( input_ids=input_ids, attention_mask=attention_mask, max_new_tokens=max(1, int(max_new_tokens)), - do_sample=False, + do_sample=do_sample, + temperature=temperature if do_sample else None, + top_p=top_p if do_sample else None, pad_token_id=pad_token_id, ) new_tokens = generated[0, input_ids.shape[-1]:] return self.tokenizer.decode(new_tokens, skip_special_tokens=True).strip() + def generate_text_streaming( + self, + messages: list[dict], + max_new_tokens: int = 256, + temperature: float = 1.0, + top_p: float = 1.0, + ): + """Yield decoded token strings one at a time using HF TextIteratorStreamer.""" + if not self.is_head or not self.is_tail: + raise ModelBackendError("streaming generation requires a full head+tail shard") + import threading + try: + from transformers import TextIteratorStreamer # type: ignore[import] + except ImportError: + yield self.generate_text(messages, max_new_tokens, temperature, top_p) + return + + encoded = self._encode_messages(messages) + input_ids = encoded["input_ids"].to(self.device) + attention_mask = encoded.get("attention_mask") + if attention_mask is not None: + attention_mask = attention_mask.to(self.device) + pad_token_id = getattr(self.tokenizer, "pad_token_id", None) or getattr(self.tokenizer, "eos_token_id", None) + do_sample = temperature != 1.0 or top_p != 1.0 + + streamer = TextIteratorStreamer(self.tokenizer, skip_prompt=True, skip_special_tokens=True) + gen_kwargs = dict( + input_ids=input_ids, + attention_mask=attention_mask, + max_new_tokens=max(1, int(max_new_tokens)), + do_sample=do_sample, + temperature=temperature if do_sample else None, + top_p=top_p if do_sample else None, + pad_token_id=pad_token_id, + streamer=streamer, + ) + t = threading.Thread(target=self.model.generate, kwargs=gen_kwargs, daemon=True) + t.start() + for token_text in streamer: + yield token_text + t.join() + + def _encode_messages(self, messages: list[dict]) -> dict: + """Format messages with chat template (if available) and tokenize.""" + if hasattr(self.tokenizer, "apply_chat_template"): + try: + prompt_str = self.tokenizer.apply_chat_template( + messages, + add_generation_prompt=True, + tokenize=False, + ) + return dict(self.tokenizer(prompt_str, return_tensors="pt")) + except Exception: + pass + prompt = " ".join( + str(m.get("content", "")) + for m in messages + if isinstance(m, dict) and m.get("role") == "user" + ) + return dict(self.tokenizer(prompt, return_tensors="pt")) + def _run_layers(self, hidden_states: Any, attention_mask: Any, position_ids: Any) -> Any: position_embeddings = _rotary_position_embeddings( self.model, diff --git a/packages/node/meshnet_node/torch_server.py b/packages/node/meshnet_node/torch_server.py index c287930..ca07797 100644 --- a/packages/node/meshnet_node/torch_server.py +++ b/packages/node/meshnet_node/torch_server.py @@ -213,8 +213,31 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler): if body is None: return messages = body.get("messages", []) + if not isinstance(messages, list): + messages = [] stream = bool(body.get("stream", False)) - model = str(body.get("model", "")) + model_name = str(body.get("model", "")) + max_tokens = int(body.get("max_tokens") or body.get("max_new_tokens") or 256) + temperature = float(body.get("temperature") or 1.0) + top_p = float(body.get("top_p") or 1.0) + + # Fast path: this node owns the complete model — use HF generate() with KV cache. + # Avoids the single-token-per-forward-pass limitation of the distributed path. + if server.backend.is_head and server.backend.is_tail: + try: + if stream: + self._stream_openai_response( + server.backend.generate_text_streaming(messages, max_tokens, temperature, top_p), + model_name, + ) + else: + text = server.backend.generate_text(messages, max_tokens, temperature, top_p) + self._send_openai_response(text, model_name, False) + except Exception as exc: + self._send_json(500, {"error": f"generation failed: {exc}"}) + return + + # Distributed path: encode prompt at the head, forward activations along the route. prompt = " ".join( str(m.get("content", "")) for m in messages @@ -225,9 +248,9 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler): except Exception as exc: self._send_json(500, {"error": f"encode_prompt failed: {exc}"}) return - remaining_route = self._get_remaining_route(model) + remaining_route = self._get_remaining_route(model_name) result_text = self._run_downstream_pipeline(payload, remaining_route) - self._send_openai_response(result_text, model, stream) + self._send_openai_response(result_text, model_name, stream) def _get_remaining_route(self, model: str) -> list[str]: server: _TorchHTTPServer = self.server # type: ignore[assignment] @@ -246,7 +269,8 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler): def _run_downstream_pipeline(self, payload: object, route: list[str]) -> str: server: _TorchHTTPServer = self.server # type: ignore[assignment] if not route: - # Single-node mode: decode tail locally if we're the tail + # Partial shard at tail: decode the activation from the previous node. + # Full single-node (head+tail) is handled before entering this method. if server.backend.is_tail: try: tensor = server.backend.torch.frombuffer( @@ -256,7 +280,7 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler): return server.backend.decode_tail(tensor) except Exception as exc: return f"decode error: {exc}" - return "" + return "no downstream route available for non-tail shard" session = str(uuid.uuid4()) shape = payload.shape # type: ignore[union-attr] @@ -309,6 +333,40 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler): current_pos = resp_headers.get("x-meshnet-position-ids") return "" + def _stream_openai_response(self, token_iter, model: str) -> None: + """Stream tokens from an iterator as SSE chunks.""" + chunk_id = "chatcmpl-node" + created = int(time.time()) + self.send_response(200) + self.send_header("Content-Type", "text/event-stream; charset=utf-8") + self.send_header("Cache-Control", "no-cache") + self.end_headers() + + def _emit(data: str) -> None: + self.wfile.write(f"data: {data}\n\n".encode()) + self.wfile.flush() + + _emit(json.dumps({ + "id": chunk_id, "object": "chat.completion.chunk", "created": created, + "model": model, + "choices": [{"index": 0, "delta": {"role": "assistant", "content": ""}, "finish_reason": None}], + })) + for token_text in token_iter: + if not token_text: + continue + _emit(json.dumps({ + "id": chunk_id, "object": "chat.completion.chunk", "created": created, + "model": model, + "choices": [{"index": 0, "delta": {"content": token_text}, "finish_reason": None}], + })) + _emit(json.dumps({ + "id": chunk_id, "object": "chat.completion.chunk", "created": created, + "model": model, + "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}], + })) + self.wfile.write(b"data: [DONE]\n\n") + self.wfile.flush() + def _send_openai_response(self, text: str, model: str, stream: bool) -> None: chunk_id = "chatcmpl-node" created = int(time.time())