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:
@@ -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])
|
||||
Reference in New Issue
Block a user