neuron-tai/docs/issues/44-tracker-shard-source-partial-download.md

US-044 — Tracker as model-file source; nodes download only their shard

Status: in progress 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).
   • The allow_patterns subset must not depend on the tracker having a local snapshot: when the tracker has no --models-dir match for a repo (or hasn't cached it yet — the common case for a fresh public tracker), model_sources comes back empty and download_shard falls straight to _download_huggingface_subset(..., allow_patterns=None), i.e. the full repo. Reported 2026-07-06: a CPU node assigned layers 0–2 of unsloth/Qwen3.6-35B-A3B (42 safetensor shards) sat downloading the entire model unauthenticated because of this. Fix: fetch model.safetensors.index.json + config.json directly from HF (a few KB) and compute the same layer-scoped file subset client-side, so the HF-fallback path is filtered even with an empty model_sources.
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

• Tracker can be started with --models-dir PATH / MESHNET_MODELS_DIR and advertises a local model-file source in assignment responses when it has a matching HF snapshot.
• Tracker serves a tar stream containing only the safetensors files selected for the assigned layer range plus config/tokenizer/index metadata.
• Node downloader keeps exact-shard peers first, then races tracker model sources against a HuggingFace snapshot_download(..., allow_patterns=...) subset download, using the first successful source.
• When no tracker model source is available at all, the HuggingFace fallback still computes allow_patterns from the repo's own model.safetensors.index.json (fetched directly, not via the tracker) — it never silently downloads the full model just because the tracker has nothing cached.
• Real PyTorch model startup can use tracker full_url sources to fetch the full local snapshot over LAN before from_pretrained, so local-network testing no longer has to pull from HuggingFace first.
• Two-machine test: machine A (tracker + node, holds full snapshot) serves layers 0–k; 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

Implementation notes

• 2026-07-06: Added the tracker/node download path. For immediate Qwen3.6-35B LAN testing, real PyTorch nodes fetch the full snapshot from the tracker via full_url; HuggingFace remains fallback-only, and when it is used the node computes allow_patterns from the repo's remote SafeTensors index so it stays layer-filtered even without tracker-cached files. Remaining hard half is true partial model materialization: the backend can prefer a downloaded local model directory, but Transformers still needs a meta-device load path that materializes only assigned layers.