docs: add ADRs and user stories for real model inference stack (US-011–014)
ADR-0008: binary activation wire format — raw bfloat16 over HTTP, zstd compression, 128-token chunked prefill; replaces base64 JSON (~33% overhead removed). ADR-0009: coverage-first shard assignment and tracker-as-first-layer-node — any node serving layers[0..k] becomes the inference entry point for that model; bin-packing fills all coverage gaps before adding redundancy; tracker issues LOAD_SHARD/DROP_SHARD rebalance directives; nodes declare VRAM + quantization. US-011: binary wire format migration US-012: real PyTorch layer execution (transformers + bitsandbytes, test on GPT-2) US-013: coverage-first tracker bin-packing with VRAM-aware shard assignment US-014: tracker-as-node (tracker node serves first layers + handles client requests) CONTEXT.md: Tracker Node, Coverage Map, Rebalance Directive terms added. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
Dobromir Popov
@@ -0,0 +1,50 @@
|
||||
# US-011 — Binary activation wire format
|
||||
|
||||
Replace the base64-encoded JSON activation payload from US-002 with a binary HTTP body, zstd compression, and chunked prefill. This is a protocol migration that must be applied to all nodes and the gateway before US-012 (real model backend) can be built on a sane foundation.
|
||||
|
||||
## Context
|
||||
|
||||
The current wire format encodes activation tensors as base64 strings inside a JSON dict:
|
||||
```json
|
||||
{"shape": [1, 1, 64], "dtype": "float32", "data": "<base64>", "context": {"prompt": "..."}}
|
||||
```
|
||||
This works for stub nodes (tiny tensors) but is unsuitable for real models:
|
||||
- base64 adds 33% size overhead
|
||||
- A 2048-token prefill at hidden_dim=16384 generates 64MB of activations per boundary — 85MB base64
|
||||
- JSON parsing overhead grows with payload size
|
||||
|
||||
The new format per ADR-0008:
|
||||
```
|
||||
POST /forward HTTP/1.1
|
||||
Content-Type: application/octet-stream
|
||||
X-Meshnet-Shape: 1,128,16384
|
||||
X-Meshnet-Dtype: bfloat16
|
||||
X-Meshnet-Session: <uuid>
|
||||
X-Meshnet-Chunk-Index: 0
|
||||
X-Meshnet-Chunk-Total: 16
|
||||
X-Meshnet-Encoding: zstd
|
||||
Content-Length: <n>
|
||||
|
||||
<raw zstd-compressed little-endian bfloat16 bytes>
|
||||
```
|
||||
|
||||
Stub nodes still emit zeroed tensors — just in binary now. No real model required for this story.
|
||||
|
||||
## Acceptance Criteria
|
||||
|
||||
- `packages/node` `/forward` endpoint reads `X-Meshnet-Shape`, `X-Meshnet-Dtype`, `X-Meshnet-Session`, `X-Meshnet-Chunk-Index`, `X-Meshnet-Chunk-Total`, and `X-Meshnet-Encoding` headers; reads raw body bytes; decompresses if `zstd`; reconstructs tensor as numpy array
|
||||
- Node forward handler returns a binary response with the same header set reflecting output shape
|
||||
- The gateway sends binary chunked activations to the first node; reassembles binary responses from the last node
|
||||
- Chunked prefill: the gateway splits input sequences longer than `MESHNET_CHUNK_TOKENS` (default 128) into N chunks and sends them sequentially through the pipeline; stub nodes pass through each chunk unchanged
|
||||
- An integration test sends a 512-token stub activation (4 chunks of 128) through a two-node pipeline and asserts all 4 chunk responses are received with correct headers
|
||||
- `zstd` Python package added as a dependency to `packages/node` and `packages/gateway`
|
||||
- The old `_make_stub_activations` function in `server.py` is replaced with `_make_stub_binary_activation(shape, dtype) -> bytes`
|
||||
- `python -m pytest` passes from repo root
|
||||
- Commit only this story's changes
|
||||
|
||||
## Implementation Notes
|
||||
|
||||
- Use `numpy` to pack/unpack tensor bytes: `np.frombuffer(body, dtype=np.float16).reshape(shape)` (bfloat16 → numpy uses `np.dtype('bfloat16')` in numpy >= 1.20, or load as uint16 and reinterpret)
|
||||
- zstd: `import zstandard as zstd; cctx = zstd.ZstdCompressor(level=1); compressed = cctx.compress(raw_bytes)`
|
||||
- Chunk index + total allow the receiving node to know whether this is a mid-sequence chunk (no special handling needed for stub; real model in US-012 will need to manage KV cache across chunks)
|
||||
- Wire version header `X-Meshnet-Wire: 2` can be added for debugging mismatches between old and new nodes
|
||||
@@ -0,0 +1,61 @@
|
||||
# US-012 — Real PyTorch model backend
|
||||
|
||||
Replace stub node inference with actual `transformers` layer execution. A node downloads a HuggingFace SafeTensors model, loads `model.layers[start:end]` into GPU/CPU memory at the declared quantization level, and runs real forward passes. The first-layer node additionally runs `embed_tokens`; the last-layer node additionally runs `model.norm` and `lm_head`.
|
||||
|
||||
## Context
|
||||
|
||||
Per ADR-0001 (PyTorch over llama.cpp) and the grilling session decisions:
|
||||
- Model format: HuggingFace SafeTensors (GGUF rejected — no per-layer hidden state API)
|
||||
- Layer extraction: `model.model.layers[start:end]` from a loaded `AutoModelForCausalLM`
|
||||
- Quantization: bitsandbytes NF4 / INT8 / bfloat16, declared by node at registration
|
||||
- Activation dtype at boundaries: always `bfloat16` regardless of weight quantization
|
||||
- First-layer node owns: tokenizer, `model.embed_tokens`, `model.layers[0..k]`
|
||||
- Last-layer node owns: `model.layers[N-k..N]`, `model.norm`, `model.lm_head`
|
||||
|
||||
Start with a small real model for testing: `openai-community/gpt2` (117M, no special license, SafeTensors available). The production target is 700B+ models — the layer extraction pattern is identical regardless of size.
|
||||
|
||||
## Acceptance Criteria
|
||||
|
||||
- `packages/node` can be started with `--model-id openai-community/gpt2 --shard-start 0 --shard-end 6` and loads the corresponding transformer blocks
|
||||
- Node detects whether it is a head shard (`shard_start == 0`) and loads tokenizer + `model.embed_tokens`
|
||||
- Node detects whether it is a tail shard (`shard_end == total_layers`) and loads `model.norm + model.lm_head`
|
||||
- A `/forward` request with a text prompt (head node) or binary bfloat16 tensor (mid/tail node) runs a real forward pass and returns binary bfloat16 activations
|
||||
- A two-node integration test with GPT-2 sharded across two nodes returns a coherent (non-random, deterministic) text completion
|
||||
- Quantization level is read from `--quantization [bfloat16|int8|nf4]` flag; `bfloat16` is default; `int8` and `nf4` use `bitsandbytes.BitsAndBytesConfig`
|
||||
- A node with insufficient VRAM for the assigned shard at the declared quantization prints a clear error and exits (no silent OOM)
|
||||
- `transformers`, `bitsandbytes`, `safetensors`, and `accelerate` added as dependencies to `packages/node`
|
||||
- `python -m pytest` passes (GPT-2 test downloads model once, cached; mark with `@pytest.mark.integration` and skip in CI if no GPU)
|
||||
- Commit only this story's changes
|
||||
|
||||
## Implementation Notes
|
||||
|
||||
```python
|
||||
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
|
||||
import torch
|
||||
|
||||
def load_shard(model_id, start, end, quantization):
|
||||
quant_config = None
|
||||
if quantization == "nf4":
|
||||
quant_config = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_quant_type="nf4",
|
||||
bnb_4bit_compute_dtype=torch.bfloat16)
|
||||
elif quantization == "int8":
|
||||
quant_config = BitsAndBytesConfig(load_in_8bit=True)
|
||||
|
||||
# Load full model metadata without weights, then extract layers
|
||||
model = AutoModelForCausalLM.from_pretrained(
|
||||
model_id,
|
||||
quantization_config=quant_config,
|
||||
device_map="auto",
|
||||
torch_dtype=torch.bfloat16,
|
||||
)
|
||||
layers = model.model.layers[start:end]
|
||||
return model, layers
|
||||
```
|
||||
|
||||
- `device_map="auto"` handles GPU/CPU offload automatically based on available VRAM
|
||||
- For large models, use `low_cpu_mem_usage=True` to avoid doubling peak RAM during load
|
||||
- The forward pass for a mid-shard node: `hidden_states = layer(hidden_states)[0]` (most HF models return a tuple, take index 0)
|
||||
- Head node forward: `input_ids → embed_tokens → position_embed → layers[0..k] → bfloat16 binary out`
|
||||
- Tail node forward: `bfloat16 binary in → layers[N-k..N] → norm → lm_head → argmax/sample → token_id → back to head`
|
||||
- Attention mask and position IDs must be forwarded alongside hidden states for models that need them (pack them in additional headers or a small JSON sidecar in the same HTTP request)
|
||||
- Keep attention/position passing simple for v1: assume left-padded inputs, pass `attention_mask` as a separate binary blob with header `X-Meshnet-Attn-Mask: <base64-uint8>` (small enough for base64 — it's [B, S] not [B, S, D])
|
||||
@@ -0,0 +1,49 @@
|
||||
# US-013 — Coverage-first tracker shard assignment
|
||||
|
||||
Upgrade the tracker's route selection to a coverage-first, speed-weighted bin-packing algorithm. The tracker maintains a live coverage map per model, issues rebalance directives when gaps appear, and assigns shard ranges to nodes based on declared VRAM, quantization, and benchmark throughput.
|
||||
|
||||
## Context
|
||||
|
||||
Per ADR-0009:
|
||||
- **Coverage gap** = any layer range in a model with zero serving nodes → model is unroutable
|
||||
- **Coverage-first**: fill all gaps before adding redundancy to covered ranges
|
||||
- **Speed-weighted bin-packing**: fill fastest nodes to their VRAM limit first; cascade to next-fastest
|
||||
- **Continuous rebalancing**: triggered every 30s and on every node join/leave event
|
||||
- **Rebalance directive**: `LOAD_SHARD` or `DROP_SHARD` instruction sent from tracker to node
|
||||
|
||||
Node registration now includes capability data:
|
||||
```json
|
||||
{
|
||||
"wallet": "...",
|
||||
"endpoint": "http://10.0.0.5:7000",
|
||||
"vram_bytes": 25769803776,
|
||||
"ram_bytes": 137438953472,
|
||||
"quantizations": ["nf4", "int8", "bfloat16"],
|
||||
"benchmark_tokens_per_sec": 12.4,
|
||||
"benchmark_model": "openai-community/gpt2"
|
||||
}
|
||||
```
|
||||
|
||||
Model preset metadata (stored in tracker config) includes `bytes_per_layer` at each quantization level and `total_layers`. The tracker can compute max assignable layers for any node.
|
||||
|
||||
## Acceptance Criteria
|
||||
|
||||
- Tracker accepts `vram_bytes`, `ram_bytes`, `quantizations`, `benchmark_tokens_per_sec` in the node registration payload
|
||||
- `GET /v1/coverage/<model_preset>` returns the coverage map: list of `{start_layer, end_layer, node_count}` for each assigned range
|
||||
- A model is only routable when all layer ranges have `node_count >= 1`
|
||||
- When a new node registers, the tracker assigns it to the highest-priority uncovered range (or expands the most-loaded range if fully covered)
|
||||
- When a node disconnects (heartbeat timeout), the tracker marks affected ranges as reduced-coverage and, if any reach 0, issues `LOAD_SHARD` directives to idle nodes within 30 seconds
|
||||
- The shard assignment respects VRAM: `assigned_layers <= floor((vram_bytes * 0.8) / bytes_per_layer_at_quant)`
|
||||
- Speed weighting: given two idle nodes both capable of covering a gap, the tracker assigns the wider sub-range to the faster node
|
||||
- An integration test: start tracker with a model preset, register three nodes with different VRAM, assert the coverage map shows 100% coverage and the largest VRAM node received the widest shard range
|
||||
- A second integration test: kill the node covering the middle layer range; assert tracker issues a `LOAD_SHARD` directive to an idle node and coverage recovers
|
||||
- `python -m pytest` passes
|
||||
- Commit only this story's changes
|
||||
|
||||
## Implementation Notes
|
||||
|
||||
- Coverage map data structure: `dict[str, list[tuple[int, int, int]]]` keyed by model preset
|
||||
- Bin-packing is run in-process on the tracker (not a separate service); it's triggered by registration/heartbeat events and a 30s periodic timer
|
||||
- Rebalance directives are delivered as responses to node heartbeat POSTs (node polls tracker every 10s anyway) — no new push channel needed
|
||||
- `bytes_per_layer` is read from a `model_presets.json` config file; add GPT-2 (12 layers, ~30MB/layer bfloat16) as the test preset
|
||||
- The current tracker `_select_route` function is extended, not replaced — backward compat with stub nodes that don't send VRAM data (default to 8GB / bfloat16 if omitted)
|
||||
@@ -0,0 +1,34 @@
|
||||
# US-014 — Tracker-as-first-layer-node (inference entry point)
|
||||
|
||||
Merge the inference orchestration role of the gateway into tracker nodes that serve the first-layer shard. A tracker node for a model receives client requests directly, runs the head shard, routes activations through the network, collects the tail output, and streams tokens back to the client. The standalone gateway becomes a thin load-balancer that routes to whichever tracker node is least loaded.
|
||||
|
||||
## Context
|
||||
|
||||
Per ADR-0009:
|
||||
- Any node serving `layers[0..k]` for a model can act as its tracker node
|
||||
- Tracker nodes own: tokenizer, `embed_tokens`, `layers[0..k]`
|
||||
- Tracker nodes select the optimal onward route from the live coverage map (they already maintain it)
|
||||
- Multiple tracker nodes for the same model = horizontal scale at both the routing and the first-layer compute level
|
||||
- The standalone `meshnet-gateway` process from US-005 becomes a dumb round-robin load-balancer; it no longer orchestrates shard pipelines
|
||||
|
||||
This collapses a network hop and a process boundary: previously gateway → tracker → node; now client → tracker-node (which is both tracker and first-layer node).
|
||||
|
||||
## Acceptance Criteria
|
||||
|
||||
- A node started with `--tracker-mode` (or automatically when assigned `shard_start=0`) exposes both its existing `/forward` endpoint and a new `/v1/chat/completions` OpenAI-compatible endpoint
|
||||
- The tracker-node's `/v1/chat/completions` handler: tokenizes input, embeds, runs its own layers, selects onward route from coverage map, forwards binary activations to next node, receives tail output, decodes tokens, streams SSE back to client
|
||||
- Multiple tracker nodes for the same model can each independently handle requests (verified by sending requests to both and getting valid responses)
|
||||
- The existing `meshnet-gateway` is updated to proxy requests to tracker nodes (round-robin or least-connections) rather than orchestrating the pipeline itself
|
||||
- The gateway can discover tracker nodes from the tracker registry (`GET /v1/tracker-nodes/<model_preset>` returns the list of endpoints for tracker nodes for that model)
|
||||
- An integration test: register two tracker nodes and two mid-shard nodes for GPT-2; send 10 requests to the gateway; assert both tracker nodes received roughly equal load and all responses are valid
|
||||
- Existing US-005 OpenAI-compatible tests still pass (gateway still exposes `/v1/chat/completions`, just proxies to tracker nodes now)
|
||||
- `python -m pytest` passes
|
||||
- Commit only this story's changes
|
||||
|
||||
## Implementation Notes
|
||||
|
||||
- Tracker-mode detection: `shard_start == 0` OR `--tracker-mode` CLI flag
|
||||
- The tracker-node process runs two HTTP servers on different ports: port 7000 for node-to-node (`/forward`), port 8080 for client-facing (`/v1/chat/completions`) — or a single server with both route prefixes
|
||||
- Route selection inside the tracker node: same `GET /v1/route/<model_preset>` call to the central tracker registry, but the tracker node itself is already the first hop — the returned route starts from the *second* node
|
||||
- Streaming: tail node returns token IDs one by one (or in batches) back to the tracker node via chunked HTTP; tracker node converts to SSE and streams to client
|
||||
- The gateway's existing pipeline orchestration code (`_run_pipeline`) can be extracted into a shared `meshnet_common.pipeline` module reused by both tracker-node and the legacy gateway path
|
||||
Reference in New Issue
Block a user