feat: harden node placement and partial model loading

This commit is contained in:
Dobromir Popov
2026-07-13 21:58:08 +02:00
parent a6bcc69288
commit 5d87e81bc9
21 changed files with 497 additions and 55 deletions

View File

@@ -86,10 +86,10 @@ What exists already (build on it, don't duplicate):
- [ ] 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
- [x] Machine B's resident memory scales with its shard size, not model size
- [x] Checksums verified end-to-end; corrupted transfer falls back cleanly
- [x] Single-node/full-model flows unchanged
- [ ] `python -m pytest` passes from repo root
- [x] `python -m pytest` passes from repo root
## Implementation notes
@@ -98,6 +98,13 @@ What exists already (build on it, don't duplicate):
`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
is 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.
- 2026-07-13: Partial LOAD implemented. `_load_partial_model_from_snapshot` builds
on `meta` via `init_empty_weights`, materializes only layer-scoped checkpoint
tensors, and finalizes device placement without copying unmaterialized meta
weights (`_finalize_active_shard_modules_on_device`). Tests cover memory
scaling (`test_partial_snapshot_resident_weight_numel_scales_with_shard`)
and real-torch meta-vs-materialized counts. Remaining: live two-machine LAN
verification.