feat: add real PyTorch model backend

This commit is contained in:
Dobromir Popov
2026-06-29 15:54:40 +03:00
parent c358798627
commit 2690d9b9ba
8 changed files with 877 additions and 37 deletions

View File

@@ -20,6 +20,18 @@ def main() -> None:
start_cmd.add_argument(
"--model", default="stub-model", help="Model preset to request from tracker"
)
start_cmd.add_argument(
"--model-id",
help="HuggingFace model id for the real PyTorch backend",
)
start_cmd.add_argument("--shard-start", type=int, help="First layer index for an explicit shard")
start_cmd.add_argument("--shard-end", type=int, help="Exclusive layer end index for an explicit shard")
start_cmd.add_argument(
"--quantization",
choices=["bfloat16", "int8", "nf4"],
default="bfloat16",
help="Weight quantization for the real PyTorch backend",
)
start_cmd.add_argument(
"--host", default="0.0.0.0", help="Interface to bind to"
)
@@ -33,13 +45,21 @@ def main() -> None:
if args.command == "start":
from meshnet_node.startup import run_startup
node = run_startup(
tracker_url=args.tracker,
port=args.port,
model=args.model,
host=args.host,
advertise_host=args.advertise_host,
)
try:
node = run_startup(
tracker_url=args.tracker,
port=args.port,
model=args.model,
model_id=args.model_id,
shard_start=args.shard_start,
shard_end=args.shard_end,
quantization=args.quantization,
host=args.host,
advertise_host=args.advertise_host,
)
except Exception as exc:
print(f"ERROR: {exc}", file=sys.stderr, flush=True)
sys.exit(1)
try:
while True:
time.sleep(1)

View File

@@ -0,0 +1,298 @@
"""HuggingFace/PyTorch shard backend for real node inference."""
from __future__ import annotations
import base64
from dataclasses import dataclass
from typing import Any, Literal
Quantization = Literal["bfloat16", "int8", "nf4"]
class ModelBackendError(RuntimeError):
"""Base class for real model backend startup and execution failures."""
class MissingModelDependencyError(ModelBackendError):
"""Raised when optional model dependencies are not installed."""
class InsufficientVRAMError(ModelBackendError):
"""Raised when a requested shard cannot fit in available CUDA memory."""
@dataclass(frozen=True)
class TensorPayload:
body: bytes
shape: list[int]
attention_mask_header: str | None
position_ids_header: str | None
def validate_quantization(value: str) -> Quantization:
if value not in {"bfloat16", "int8", "nf4"}:
raise ValueError("quantization must be one of: bfloat16, int8, nf4")
return value # type: ignore[return-value]
def build_quantization_config(quantization: Quantization) -> Any | None:
"""Return a transformers BitsAndBytesConfig for quantized weights."""
if quantization == "bfloat16":
return None
try:
import torch
from transformers import BitsAndBytesConfig
except ModuleNotFoundError as exc:
raise MissingModelDependencyError(
"transformers and torch are required for int8/nf4 quantization"
) from exc
if quantization == "int8":
return BitsAndBytesConfig(load_in_8bit=True)
return BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype=torch.bfloat16,
)
class TorchModelShard:
"""Executable subset of a HuggingFace causal language model."""
def __init__(
self,
model_id: str,
shard_start: int,
shard_end: int,
quantization: Quantization = "bfloat16",
) -> None:
if shard_start < 0 or shard_end < 0 or shard_start > shard_end:
raise ValueError("shard_start must be <= shard_end and non-negative")
self.model_id = model_id
self.shard_start = shard_start
self.shard_end = shard_end
self.quantization = quantization
try:
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
except ModuleNotFoundError as exc:
raise MissingModelDependencyError(
"real model backend requires torch, transformers, safetensors, accelerate, and bitsandbytes"
) from exc
self.torch = torch
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
quant_config = build_quantization_config(quantization)
try:
self.model = AutoModelForCausalLM.from_pretrained(
model_id,
quantization_config=quant_config,
device_map="auto" if quant_config is not None else None,
torch_dtype=torch.bfloat16,
low_cpu_mem_usage=True,
use_safetensors=True,
)
if quant_config is None:
self.model.to(self.device)
except Exception as exc:
if _looks_like_oom(exc):
raise InsufficientVRAMError(
f"insufficient VRAM to load {model_id} layers {shard_start}:{shard_end} "
f"with {quantization} quantization; choose a smaller shard or lower quantization"
) from exc
raise
self.model.eval()
self.tokenizer = AutoTokenizer.from_pretrained(model_id)
self.layers = _model_layers(self.model)
self.total_layers = len(self.layers)
if shard_end > self.total_layers:
raise ValueError(
f"shard_end {shard_end} exceeds total layer count {self.total_layers}"
)
self.is_head = shard_start == 0
self.is_tail = shard_end == self.total_layers
self.hidden_size = int(
getattr(self.model.config, "hidden_size", 0)
or getattr(self.model.config, "n_embd", 0)
)
self._embed_tokens = _embed_tokens(self.model) if self.is_head else None
self._position_embeddings = _position_embeddings(self.model)
self._norm = _final_norm(self.model) if self.is_tail else None
self._lm_head = getattr(self.model, "lm_head", None) if self.is_tail else None
def encode_prompt(self, prompt: str) -> TensorPayload:
if not self.is_head or self._embed_tokens is None:
raise ModelBackendError("text prompts can only be accepted by the head shard")
encoded = self.tokenizer(prompt, return_tensors="pt")
input_ids = encoded["input_ids"].to(self.device)
attention_mask = encoded.get("attention_mask")
if attention_mask is None:
attention_mask = self.torch.ones_like(input_ids)
attention_mask = attention_mask.to(self.device)
position_ids = _position_ids(attention_mask, self.torch)
hidden_states = self._embed_tokens(input_ids)
if self._position_embeddings is not None:
hidden_states = hidden_states + self._position_embeddings(position_ids)
hidden_states = self._run_layers(hidden_states, attention_mask, position_ids)
return self._payload(hidden_states, attention_mask, position_ids)
def forward_bytes(
self,
body: bytes,
shape: list[int],
attention_mask_header: str | None,
position_ids_header: str | None,
) -> TensorPayload | str:
hidden_states = _tensor_from_bfloat16_bytes(body, shape, self.torch).to(
self.device
)
attention_mask = _tensor_from_int64_header(
attention_mask_header, self.torch, self.device
)
position_ids = _tensor_from_int64_header(
position_ids_header, self.torch, self.device
)
hidden_states = self._run_layers(hidden_states, attention_mask, position_ids)
if self.is_tail:
return self.decode_tail(hidden_states)
return self._payload(hidden_states, attention_mask, position_ids)
def decode_tail(self, hidden_states: Any) -> str:
if self._norm is not None:
hidden_states = self._norm(hidden_states)
if self._lm_head is None:
raise ModelBackendError("tail shard has no lm_head")
logits = self._lm_head(hidden_states)
token_id = int(self.torch.argmax(logits[:, -1, :], dim=-1)[0].item())
return self.tokenizer.decode([token_id], skip_special_tokens=True)
def _run_layers(self, hidden_states: Any, attention_mask: Any, position_ids: Any) -> Any:
with self.torch.inference_mode():
for layer in self.layers[self.shard_start:self.shard_end]:
hidden_states = _call_layer(layer, hidden_states, attention_mask, position_ids)
return hidden_states.to(self.torch.bfloat16)
def _payload(self, hidden_states: Any, attention_mask: Any, position_ids: Any) -> TensorPayload:
hidden_states = hidden_states.to(self.torch.bfloat16).contiguous()
return TensorPayload(
body=_tensor_to_bytes(hidden_states),
shape=list(hidden_states.shape),
attention_mask_header=_int_tensor_header(attention_mask)
if attention_mask is not None
else None,
position_ids_header=_int_tensor_header(position_ids)
if position_ids is not None
else None,
)
def load_torch_shard(
model_id: str,
shard_start: int,
shard_end: int,
quantization: Quantization = "bfloat16",
) -> TorchModelShard:
return TorchModelShard(model_id, shard_start, shard_end, quantization)
def _model_layers(model: Any) -> Any:
if hasattr(model, "model") and hasattr(model.model, "layers"):
return model.model.layers
if hasattr(model, "transformer") and hasattr(model.transformer, "h"):
return model.transformer.h
raise ModelBackendError(
"unsupported HuggingFace model architecture: no transformer layers found"
)
def _embed_tokens(model: Any) -> Any:
if hasattr(model, "model") and hasattr(model.model, "embed_tokens"):
return model.model.embed_tokens
if hasattr(model, "transformer") and hasattr(model.transformer, "wte"):
return model.transformer.wte
raise ModelBackendError(
"unsupported HuggingFace model architecture: no token embeddings found"
)
def _position_embeddings(model: Any) -> Any | None:
if hasattr(model, "transformer") and hasattr(model.transformer, "wpe"):
return model.transformer.wpe
return None
def _final_norm(model: Any) -> Any | None:
if hasattr(model, "model") and hasattr(model.model, "norm"):
return model.model.norm
if hasattr(model, "transformer") and hasattr(model.transformer, "ln_f"):
return model.transformer.ln_f
return None
def _position_ids(attention_mask: Any, torch: Any) -> Any:
position_ids = attention_mask.long().cumsum(-1) - 1
return position_ids.masked_fill(attention_mask == 0, 0).to(torch.long)
def _call_layer(layer: Any, hidden_states: Any, attention_mask: Any, position_ids: Any) -> Any:
attempts = (
{
"attention_mask": attention_mask,
"position_ids": position_ids,
"use_cache": False,
},
{"attention_mask": attention_mask, "use_cache": False},
{"use_cache": False},
{},
)
last_exc: Exception | None = None
for kwargs in attempts:
filtered = {key: value for key, value in kwargs.items() if value is not None}
try:
output = layer(hidden_states, **filtered)
return output[0] if isinstance(output, tuple) else output
except TypeError as exc:
last_exc = exc
if last_exc is not None:
raise last_exc
return layer(hidden_states)[0]
def _tensor_to_bytes(tensor: Any) -> bytes:
import torch
return tensor.detach().cpu().contiguous().view(torch.uint8).numpy().tobytes()
def _tensor_from_bfloat16_bytes(body: bytes, shape: list[int], torch: Any) -> Any:
tensor = torch.frombuffer(bytearray(body), dtype=torch.bfloat16)
return tensor.reshape(shape)
def _int_tensor_header(tensor: Any) -> str:
data = tensor.detach().cpu().to(tensor.int64).contiguous()
raw = data.numpy().tobytes()
shape = ",".join(str(dim) for dim in data.shape)
encoded = base64.b64encode(raw).decode("ascii")
return f"{shape}:{encoded}"
def _tensor_from_int64_header(value: str | None, torch: Any, device: Any) -> Any | None:
if not value:
return None
shape_text, encoded = value.split(":", 1)
shape = [int(part) for part in shape_text.split(",") if part]
raw = base64.b64decode(encoded.encode("ascii"))
return torch.frombuffer(bytearray(raw), dtype=torch.int64).reshape(shape).to(device)
def _looks_like_oom(exc: BaseException) -> bool:
current: BaseException | None = exc
while current is not None:
text = str(current).lower()
if "out of memory" in text or "cuda error: out of memory" in text:
return True
current = current.__cause__ or current.__context__
return False

View File

@@ -14,6 +14,7 @@ from typing import Any
from .downloader import compute_shard_checksum, download_shard
from .hardware import detect_hardware
from .server import StubNodeServer
from .torch_server import TorchNodeServer
from .wallet import load_or_create_wallet
@@ -35,12 +36,16 @@ def run_startup(
tracker_url: str,
port: int = 0,
model: str = "stub-model",
model_id: str | None = None,
shard_start: int | None = None,
shard_end: int | None = None,
quantization: str = "bfloat16",
wallet_path: Path | None = None,
cache_dir: Path | None = None,
host: str = "127.0.0.1",
advertise_host: str | None = None,
contracts: Any | None = None,
) -> StubNodeServer:
) -> StubNodeServer | TorchNodeServer:
"""Execute the full startup sequence and return a running node server.
Steps (all non-interactive):
@@ -79,6 +84,35 @@ def run_startup(
if probationary_line is not None:
print(f" {probationary_line}", flush=True)
if model_id is not None and shard_start is not None and shard_end is not None:
print("Loading real PyTorch model shard...", flush=True)
node = TorchNodeServer(
host=host,
port=port,
model_id=model_id,
shard_start=shard_start,
shard_end=shard_end,
quantization=quantization,
)
actual_port = node.start()
public_host = advertise_host or (socket.getfqdn() if host == "0.0.0.0" else host)
endpoint = f"http://{public_host}:{actual_port}"
print(
f"\n{'=' * 32}\n"
f"meshnet-node ready\n"
f" Wallet: {address}\n"
f" Model ID: {model_id}\n"
f" Shard: layers {shard_start}-{shard_end}\n"
f" Quantization: {quantization}\n"
f" Endpoint: {endpoint}\n"
f" Hardware: {device.upper()}\n"
f"{'=' * 32}",
flush=True,
)
return node
if model_id is not None or shard_start is not None or shard_end is not None:
raise ValueError("--model-id, --shard-start, and --shard-end must be provided together")
# 3. Shard assignment from tracker
print("Querying tracker for shard assignment...", flush=True)
assign_qs = urllib.parse.urlencode({

View File

@@ -0,0 +1,271 @@
"""HTTP server for real PyTorch-backed shard nodes."""
from __future__ import annotations
import http.server
import json
import sys
import threading
from .model_backend import (
InsufficientVRAMError,
MissingModelDependencyError,
Quantization,
TorchModelShard,
validate_quantization,
)
from .server import (
_WIRE_VERSION,
_compress_body,
_decompress_body,
_parse_shape,
_validate_activation_body,
)
class _TorchHTTPServer(http.server.HTTPServer):
def __init__(self, addr, handler, backend: TorchModelShard):
super().__init__(addr, handler)
self.backend = backend
self.received_activations = False
self.forward_chunk_count = 0
class _TorchHandler(http.server.BaseHTTPRequestHandler):
def log_message(self, fmt, *args): # noqa: suppress request logs in tests
pass
def do_POST(self):
if self.path == "/forward":
self._handle_forward()
elif self.path == "/v1/infer":
self._handle_infer()
else:
self.send_response(404)
self.end_headers()
def _handle_infer(self) -> None:
body = self._read_json_body()
if body is None:
return
messages = body.get("messages", [])
prompt = ""
if isinstance(messages, list) and messages:
last = messages[-1]
if isinstance(last, dict):
prompt = str(last.get("content", ""))
server: _TorchHTTPServer = self.server # type: ignore[assignment]
try:
payload = server.backend.encode_prompt(prompt)
if server.backend.is_tail:
text = server.backend.decode_tail(
server.backend.torch.frombuffer(
bytearray(payload.body),
dtype=server.backend.torch.bfloat16,
)
.reshape(payload.shape)
.to(server.backend.device)
)
self._send_json(200, {"text": text})
return
self._send_json(200, {"activations": {"shape": payload.shape, "dtype": "bfloat16"}})
except Exception as exc:
self._send_json(500, {"error": str(exc)})
def _handle_forward(self) -> None:
content_type = self.headers.get("Content-Type", "")
if content_type.startswith("application/json"):
self._handle_prompt_forward()
return
self._handle_binary_forward()
def _handle_prompt_forward(self) -> None:
body = self._read_json_body()
if body is None:
return
prompt = str(body.get("prompt", ""))
server: _TorchHTTPServer = self.server # type: ignore[assignment]
try:
payload = server.backend.encode_prompt(prompt)
except Exception as exc:
self._send_json(400, {"error": str(exc)})
return
self._send_activation(payload)
def _handle_binary_forward(self) -> None:
server: _TorchHTTPServer = self.server # type: ignore[assignment]
try:
shape = _parse_shape(self.headers.get("X-Meshnet-Shape"))
dtype = self.headers.get("X-Meshnet-Dtype", "")
session = self.headers["X-Meshnet-Session"]
chunk_index = self.headers["X-Meshnet-Chunk-Index"]
chunk_total = self.headers["X-Meshnet-Chunk-Total"]
encoding = self.headers.get("X-Meshnet-Encoding")
length = int(self.headers.get("Content-Length", 0))
body = self.rfile.read(length)
raw_body = _decompress_body(body, encoding)
_validate_activation_body(raw_body, shape, dtype)
if dtype != "bfloat16":
raise ValueError("real model backend requires bfloat16 activation input")
chunk_index_value = int(chunk_index)
chunk_total_value = int(chunk_total)
if chunk_total_value <= 0 or not 0 <= chunk_index_value < chunk_total_value:
raise ValueError("invalid chunk index/total")
except (KeyError, ValueError, TypeError):
self.send_response(400)
self.send_header("X-Meshnet-Wire", _WIRE_VERSION)
self.end_headers()
return
server.forward_chunk_count += 1
if int(self.headers.get("X-Meshnet-Hop-Index", "0")) > 0:
server.received_activations = True
try:
result = server.backend.forward_bytes(
raw_body,
shape,
self.headers.get("X-Meshnet-Attn-Mask"),
self.headers.get("X-Meshnet-Position-Ids"),
)
except Exception as exc:
self._send_json(500, {"error": str(exc)})
return
if isinstance(result, str):
self._send_json(200, {"text": result})
return
response_body = _compress_body(result.body, encoding)
self.send_response(200)
self.send_header("Content-Type", "application/octet-stream")
self.send_header("Content-Length", str(len(response_body)))
self.send_header("X-Meshnet-Wire", _WIRE_VERSION)
self.send_header("X-Meshnet-Shape", ",".join(str(dim) for dim in result.shape))
self.send_header("X-Meshnet-Dtype", "bfloat16")
self.send_header("X-Meshnet-Session", session)
self.send_header("X-Meshnet-Chunk-Index", chunk_index)
self.send_header("X-Meshnet-Chunk-Total", chunk_total)
if encoding:
self.send_header("X-Meshnet-Encoding", encoding)
if result.attention_mask_header:
self.send_header("X-Meshnet-Attn-Mask", result.attention_mask_header)
if result.position_ids_header:
self.send_header("X-Meshnet-Position-Ids", result.position_ids_header)
self.end_headers()
self.wfile.write(response_body)
def _send_activation(self, payload) -> None:
body = payload.body
self.send_response(200)
self.send_header("Content-Type", "application/octet-stream")
self.send_header("Content-Length", str(len(body)))
self.send_header("X-Meshnet-Wire", _WIRE_VERSION)
self.send_header("X-Meshnet-Shape", ",".join(str(dim) for dim in payload.shape))
self.send_header("X-Meshnet-Dtype", "bfloat16")
if payload.attention_mask_header:
self.send_header("X-Meshnet-Attn-Mask", payload.attention_mask_header)
if payload.position_ids_header:
self.send_header("X-Meshnet-Position-Ids", payload.position_ids_header)
self.end_headers()
self.wfile.write(body)
def _read_json_body(self) -> dict | None:
length = int(self.headers.get("Content-Length", 0))
try:
body = json.loads(self.rfile.read(length) or b"{}")
except (json.JSONDecodeError, ValueError):
self._send_json(400, {"error": "invalid JSON body"})
return None
if not isinstance(body, dict):
self._send_json(400, {"error": "JSON body must be an object"})
return None
return body
def _send_json(self, status: int, data: dict) -> None:
payload = json.dumps(data).encode()
self.send_response(status)
self.send_header("Content-Type", "application/json")
self.send_header("Content-Length", str(len(payload)))
self.end_headers()
self.wfile.write(payload)
class TorchNodeServer:
"""HTTP server backed by a HuggingFace causal language model shard."""
def __init__(
self,
host: str = "127.0.0.1",
port: int = 0,
model_id: str = "openai-community/gpt2",
shard_start: int = 0,
shard_end: int = 6,
quantization: str = "bfloat16",
backend: TorchModelShard | None = None,
) -> None:
self._host = host
self._requested_port = port
self._backend = backend or _load_backend(
model_id,
shard_start,
shard_end,
quantization,
)
self._server: _TorchHTTPServer | None = None
self._thread: threading.Thread | None = None
self.port: int | None = None
@property
def backend(self) -> TorchModelShard:
return self._backend
@property
def received_activations(self) -> bool:
return self._server.received_activations if self._server is not None else False
@property
def forward_chunk_count(self) -> int:
return self._server.forward_chunk_count if self._server is not None else 0
def start(self) -> int:
if self._server is not None:
raise RuntimeError("TorchNodeServer is already running")
self._server = _TorchHTTPServer(
(self._host, self._requested_port),
_TorchHandler,
self._backend,
)
self.port = self._server.server_address[1]
self._thread = threading.Thread(target=self._server.serve_forever, daemon=True)
self._thread.start()
return self.port
def stop(self) -> None:
if self._server is None:
return
self._server.shutdown()
self._server.server_close()
if self._thread is not None:
self._thread.join(timeout=1)
self._server = None
self._thread = None
self.port = None
def _load_backend(
model_id: str,
shard_start: int,
shard_end: int,
quantization: str,
) -> TorchModelShard:
from .model_backend import load_torch_shard
quant = validate_quantization(quantization)
try:
return load_torch_shard(model_id, shard_start, shard_end, quant)
except MissingModelDependencyError:
raise
except InsufficientVRAMError as exc:
print(f"ERROR: {exc}", file=sys.stderr, flush=True)
raise

View File

@@ -11,6 +11,11 @@ requires-python = ">=3.10"
dependencies = [
"cryptography>=41",
"huggingface-hub>=0.20",
"accelerate>=0.28",
"bitsandbytes>=0.43",
"safetensors>=0.4",
"torch>=2.1",
"transformers>=4.39",
"zstandard>=0.22",
]