|
|
@@ -119,6 +119,9 @@ class Ledger:
|
|
|
|
|
|
def set(self, name: str, size: int, lo: int, hi: int, rows: int, samples: dict):
|
|
|
self.files[name] = {"size": size, "lo": lo, "hi": hi, "rows": rows, "sample": samples}
|
|
|
+ self.save()
|
|
|
+
|
|
|
+ def save(self) -> None:
|
|
|
self.path.parent.mkdir(parents=True, exist_ok=True)
|
|
|
self.path.write_text(json.dumps({"files": self.files}))
|
|
|
self.path.chmod(0o644)
|
|
|
@@ -132,6 +135,125 @@ def sample_rows(lo: int, hi: int, count: int) -> list:
|
|
|
return [lo + int(i * step) for i in range(count)]
|
|
|
|
|
|
|
|
|
+def sparse_recs(root: Path, layer: int):
|
|
|
+ """{kind: (src_path, rel_name, base_off, rows)} for the layer's engram
|
|
|
+ tensors — base_off is the tensor's first byte inside the shard."""
|
|
|
+ index = json.loads((root / "model.safetensors.index.json").read_text())["weight_map"]
|
|
|
+ recs = {}
|
|
|
+ for kind, width in (("weight", WEIGHT_BYTES), ("scale", SCALE_BYTES)):
|
|
|
+ tname = f"layers.{layer}.engram.embed.{kind}"
|
|
|
+ rel = index[tname]
|
|
|
+ path = root / rel
|
|
|
+ with path.open("rb") as f:
|
|
|
+ length = struct.unpack("<Q", f.read(8))[0]
|
|
|
+ header = json.loads(f.read(length))
|
|
|
+ info = header[tname]
|
|
|
+ assert info["dtype"] in ("F8_E4M3", "F8_E8M0"), f"unexpected dtype for {tname}"
|
|
|
+ assert info["shape"][1] == width
|
|
|
+ recs[kind] = (path, rel, 8 + length + info["data_offsets"][0], info["shape"][0])
|
|
|
+ return recs
|
|
|
+
|
|
|
+
|
|
|
+def row_hash_src(recs, row: int) -> str:
|
|
|
+ wpath, _, woff, _ = recs["weight"]
|
|
|
+ spath, _, soff, _ = recs["scale"]
|
|
|
+ return hashlib.sha256(read_row(wpath, woff, soff, row)).digest()[:16].hex()
|
|
|
+
|
|
|
+
|
|
|
+def row_hash_dst(wpath, woff, spath, soff, row: int) -> str:
|
|
|
+ return hashlib.sha256(read_row(wpath, woff, soff, row)).digest()[:16].hex()
|
|
|
+
|
|
|
+
|
|
|
+def copy_range(src_f, dst_f, off: int, length: int, started: float, tag: str) -> None:
|
|
|
+ """Stream one byte range src->dst at the same offset (sparse-friendly)."""
|
|
|
+ src_f.seek(off)
|
|
|
+ dst_f.seek(off)
|
|
|
+ done = 0
|
|
|
+ chunk = 1 << 26
|
|
|
+ while done < length:
|
|
|
+ n = min(chunk, length - done)
|
|
|
+ buf = src_f.read(n)
|
|
|
+ if len(buf) != n:
|
|
|
+ raise SystemExit(f"STAGING_FAIL short read at {tag} offset {off + done}")
|
|
|
+ dst_f.write(buf)
|
|
|
+ done += n
|
|
|
+ rate = done / max(time.monotonic() - started, 1e-6)
|
|
|
+ if done == length or done % (1 << 30) < chunk:
|
|
|
+ print(f"PACK sparse {tag}: {done/2**30:6.1f} GiB {rate/2**20:.0f} MiB/s",
|
|
|
+ flush=True)
|
|
|
+
|
|
|
+
|
|
|
+def stage_sparse_layer(out: Path, layer: int, recs, lo: int, hi: int, ledger: Ledger,
|
|
|
+ samples_n: int, force: bool, verify: str) -> None:
|
|
|
+ """One rank slice into node-local sparse shard copies at the ORIGINAL
|
|
|
+ byte offsets (the aiden16 `DSV41_ENGRAM_DIR` contract): the loader
|
|
|
+ mmaps the copied shard and preads rows, so holes are free and the
|
|
|
+ file size must match the source exactly. Verification is again
|
|
|
+ content-based: size + sampled row hashes vs source — a manually
|
|
|
+ staged tree (the old tools/engram_local.py output) is adopted into
|
|
|
+ the ledger instead of re-copied."""
|
|
|
+ key = f"sparse-l{layer}"
|
|
|
+ wpath, wrel, woff, _ = recs["weight"]
|
|
|
+ spath, srel, soff, _ = recs["scale"]
|
|
|
+ dst_w, dst_s = out / wrel, out / srel
|
|
|
+ size_w, size_s = wpath.stat().st_size, spath.stat().st_size
|
|
|
+ rec = ledger.files.get(key)
|
|
|
+
|
|
|
+ def files_ok() -> bool:
|
|
|
+ return (dst_w.exists() and dst_w.stat().st_size == size_w
|
|
|
+ and dst_s.exists() and dst_s.stat().st_size == size_s)
|
|
|
+
|
|
|
+ def check(samples) -> bool:
|
|
|
+ for row, want in samples:
|
|
|
+ if row_hash_src(recs, row) != want:
|
|
|
+ return False # source row changed: rebuild
|
|
|
+ if row_hash_dst(dst_w, woff, dst_s, soff, row) != want:
|
|
|
+ return False # staged row corrupted: rebuild
|
|
|
+ return True
|
|
|
+
|
|
|
+ t0 = time.monotonic()
|
|
|
+ verified_only = False
|
|
|
+ if not force and files_ok():
|
|
|
+ if rec is not None and rec.get("lo") == lo and rec.get("hi") == hi:
|
|
|
+ if verify == "presence" or check([(int(k), v) for k, v in rec["sample"].items()]):
|
|
|
+ verified_only = True
|
|
|
+ else:
|
|
|
+ # untracked tree (manual staging): adopt when samples match
|
|
|
+ if verify != "presence":
|
|
|
+ samples = [(r, row_hash_src(recs, r)) for r in sample_rows(lo, hi, samples_n)]
|
|
|
+ if check(samples):
|
|
|
+ ledger.files[key] = {"lo": lo, "hi": hi,
|
|
|
+ "sample": {str(r): h for r, h in samples}}
|
|
|
+ verified_only = True
|
|
|
+ if verified_only:
|
|
|
+ ledger.save()
|
|
|
+ print(f"STAGED sparse l{layer} rows={hi - lo} verified "
|
|
|
+ f"secs={time.monotonic() - t0:.0f}", flush=True)
|
|
|
+ return
|
|
|
+
|
|
|
+ # (re)build: ensure sparse-sized files, then stream our ranges in
|
|
|
+ for dst, size in ((dst_w, size_w), (dst_s, size_s)):
|
|
|
+ dst.parent.mkdir(parents=True, exist_ok=True)
|
|
|
+ if not dst.exists():
|
|
|
+ with open(dst, "xb"):
|
|
|
+ os.truncate(dst, size)
|
|
|
+ elif dst.stat().st_size != size:
|
|
|
+ os.truncate(dst, size)
|
|
|
+ with open(wpath, "rb", buffering=0) as sf, open(dst_w, "r+b", buffering=0) as df:
|
|
|
+ copy_range(sf, df, woff + lo * WEIGHT_BYTES, (hi - lo) * WEIGHT_BYTES,
|
|
|
+ t0, f"l{layer}-weight")
|
|
|
+ with open(spath, "rb", buffering=0) as sf, open(dst_s, "r+b", buffering=0) as df:
|
|
|
+ copy_range(sf, df, soff + lo * SCALE_BYTES, (hi - lo) * SCALE_BYTES,
|
|
|
+ t0, f"l{layer}-scale")
|
|
|
+ ledger.files[key] = {"lo": lo, "hi": hi,
|
|
|
+ "sample": {str(r): row_hash_src(recs, r)
|
|
|
+ for r in sample_rows(lo, hi, samples_n)}}
|
|
|
+ ledger.save()
|
|
|
+ print(f"STAGED sparse l{layer} rows={hi - lo} packed "
|
|
|
+ f"secs={time.monotonic() - t0:.0f}", flush=True)
|
|
|
+
|
|
|
+
|
|
|
+
|
|
|
def verify_file(path: Path, src: Path, w_off: int, s_off: int, lo: int, hi: int,
|
|
|
size: int, ledger: Ledger, name: str) -> "bool | None":
|
|
|
"""True staged-ok, False needs repack, None size-only trust (presence)."""
|
|
|
@@ -161,6 +283,9 @@ def main() -> int:
|
|
|
ap.add_argument("--rank", type=int, required=True)
|
|
|
ap.add_argument("--tp", type=int, required=True)
|
|
|
ap.add_argument("--layers", default="1,14")
|
|
|
+ ap.add_argument("--format", default="packed", choices=["packed", "sparse-copy"],
|
|
|
+ help="packed shards (row-store) or sparse shard copies at the "
|
|
|
+ "original offsets (safetensors-pread / aiden16 contract)")
|
|
|
ap.add_argument("--min-free", default=None)
|
|
|
ap.add_argument("--force", action="store_true")
|
|
|
ap.add_argument("--verify", default="sample-hash", choices=["sample-hash", "presence"])
|
|
|
@@ -199,6 +324,29 @@ def main() -> int:
|
|
|
print(f"STAGING_FAIL need {need/2**30:.1f} GiB, only {free/2**30:.1f} GiB at {out_dir}")
|
|
|
return 2
|
|
|
|
|
|
+ if args.format == "sparse-copy":
|
|
|
+ # the loader resolves tensors through the index, so it travels too
|
|
|
+ idx_src = root / "model.safetensors.index.json"
|
|
|
+ idx_dst = out_dir / "model.safetensors.index.json"
|
|
|
+ want = hashlib.sha256(idx_src.read_bytes()).hexdigest()
|
|
|
+ if not idx_dst.exists() or hashlib.sha256(idx_dst.read_bytes()).hexdigest() != want:
|
|
|
+ tmp = idx_dst.with_name(idx_dst.name + ".sm121tmp")
|
|
|
+ tmp.write_bytes(idx_src.read_bytes())
|
|
|
+ os.replace(tmp, idx_dst)
|
|
|
+ manifest = {}
|
|
|
+ for layer in layers:
|
|
|
+ recs = sparse_recs(root, layer)
|
|
|
+ rows = recs["weight"][3]
|
|
|
+ lo, hi = rows * args.rank // args.tp, rows * (args.rank + 1) // args.tp
|
|
|
+ stage_sparse_layer(out_dir, layer, recs, lo, hi, ledger,
|
|
|
+ args.samples, args.force, args.verify)
|
|
|
+ manifest[str(layer)] = [lo, hi]
|
|
|
+ tmp = out_dir / "engram-local.json.sm121tmp"
|
|
|
+ tmp.write_text(json.dumps({"layers": manifest}))
|
|
|
+ os.replace(tmp, out_dir / "engram-local.json")
|
|
|
+ print("STAGING_DONE", flush=True)
|
|
|
+ return 0
|
|
|
+
|
|
|
for layer in layers:
|
|
|
path, rows, w_off, s_off, lo, hi = spans[layer]
|
|
|
name = f"engram-l{layer}-r{args.rank}of{args.tp}.bin"
|
