Files
neuron-tai/docs/issues/44-tracker-shard-source-partial-download.md
2026-07-06 17:11:53 +03:00

66 lines
3.4 KiB
Markdown
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
# US-044 — Tracker as model-file source; nodes download only their shard
Status: planned
Priority: High (blocks multi-machine big-model serving; pairs with US-042)
Stage: Designed (grill remaining decisions before build)
## Context
Common deployment: the tracker and the first node share a machine that already
holds the model files (e.g. `/run/media/popov/DATA/llm/safetensor`). When a
second node joins with no model selected, the tracker assigns it the uncovered
layer range — and today that node then downloads the **entire snapshot from
HuggingFace**, even for a 20-layer shard of a 160 GB model.
What exists already (build on it, don't duplicate):
- Nodes serve their shard dir as a tar at `GET /v1/shards/download` with
checksum verification; `download_shard` tries assignment-provided `peers`
before HF (`downloader.py`). But it only matches **identical layer ranges**,
and the HF fallback runs `snapshot_download` of the whole repo.
- The torch path (`--model-id`) bypasses `download_shard` entirely:
`TorchModelShard``from_pretrained` downloads **and loads into RAM** the
full model, then executes only the assigned layers. Sharding currently saves
compute, not memory or bandwidth.
## Design
1. **Tracker `--models-dir PATH`** (env `MESHNET_MODELS_DIR`). When set, the
tracker indexes HF-layout snapshots under it and advertises itself as a
model-file source in `/v1/nodes/assign` responses.
2. **Layer-aware file selection.** For safetensors models, read
`model.safetensors.index.json` and map the assigned layer range → the
subset of weight files containing those layers, plus the always-needed
files (config, tokenizer, index, embeddings/head files for head/tail
shards). Serve exactly that subset (tar stream, per-file checksums).
GGUF (US-042): single file or naive byte-range — phase 2.
3. **Node download order**: exact-shard peer (existing) → tracker/peer file
subset (new) → HF `snapshot_download` with `allow_patterns` for the same
subset (new — stop downloading the whole repo even from HF) → full snapshot
(last resort).
4. **Partial LOAD (the hard half).** Downloading a subset is wasted unless the
node stops instantiating the full model: build the model skeleton on the
`meta` device, materialize only assigned layers (+embeddings/norm/head as
role requires) from the local files, leave the rest on meta. Without this,
an 80 GB machine can never hold a shard of a 160 GB model regardless of
how the bytes arrive. This is the acceptance bar for the issue.
## Open questions (grill before building)
- Trust: joining nodes fetch weights from the tracker/peers — checksum against
what root of trust? (HF etag/sha vs tracker-signed manifest.)
- Disk layout: partial snapshots must not corrupt the HF cache dir; probably
a meshnet-owned layout keyed by repo+revision.
- Serving cost: a 100 GB tar stream per joining node on the tracker box —
rate-limit/queue? LAN-only heuristic?
## Acceptance criteria
- Two-machine test: machine A (tracker + node, holds full snapshot) serves
layers 0k; machine B joins with no model and receives **only** the files
for its assigned range from A — nothing fetched from HF
- Machine B's resident memory scales with its shard size, not model size
- Checksums verified end-to-end; corrupted transfer falls back cleanly
- Single-node/full-model flows unchanged
- `python -m pytest` passes from repo root