|
|
@@ -0,0 +1,256 @@
|
|
|
+#!/usr/bin/env python3
|
|
|
+"""sm121 sparse-rows staging engine (spec: staging-feature.md, P2).
|
|
|
+
|
|
|
+Pack one launch rank's slice of the Engram tables into a node-local
|
|
|
+contiguous shard dir, so the engine reads its rows from local NVMe instead
|
|
|
+of two strided NFS reads per cache miss.
|
|
|
+
|
|
|
+Packed layout (the runtime's read contract):
|
|
|
+ engram-l<layer>-r<rank>of<tp>.bin
|
|
|
+ 4096 B header: <Q magic, layer, lo, hi, rows, record_bytes>
|
|
|
+ rows*record_bytes of interleaved records (weight + scale per row)
|
|
|
+
|
|
|
+The row split (lo = rows*rank//tp, hi = rows*(rank+1)//tp) is derived from
|
|
|
+the checkpoint header, never hardcoded. Idempotency: exact expected size
|
|
|
+plus a per-file ledger of sampled combined (weight+scale) row hashes
|
|
|
+checked against the source — mtime/size alone never count.
|
|
|
+"""
|
|
|
+from __future__ import annotations
|
|
|
+
|
|
|
+import argparse
|
|
|
+import hashlib
|
|
|
+import json
|
|
|
+import os
|
|
|
+from pathlib import Path
|
|
|
+import struct
|
|
|
+import sys
|
|
|
+import time
|
|
|
+
|
|
|
+MAGIC = 0x31344E4531565344 # "DSV41EN1"
|
|
|
+HEADER_BYTES = 4096
|
|
|
+try:
|
|
|
+ import numpy as np
|
|
|
+
|
|
|
+ WEIGHT_BYTES, SCALE_BYTES = 256, 8
|
|
|
+ RECORD_BYTES = WEIGHT_BYTES + SCALE_BYTES
|
|
|
+except ImportError: # tiny fallback, numpy not required
|
|
|
+ np = None
|
|
|
+ WEIGHT_BYTES, SCALE_BYTES, RECORD_BYTES = 256, 8, 264
|
|
|
+
|
|
|
+
|
|
|
+def parse_size(s: str) -> int:
|
|
|
+ s = s.strip().upper()
|
|
|
+ mult = 1
|
|
|
+ if s.endswith("G"):
|
|
|
+ mult = 2**30
|
|
|
+ elif s.endswith("T"):
|
|
|
+ mult = 2**40
|
|
|
+ elif s.endswith("M"):
|
|
|
+ mult = 2**20
|
|
|
+ return int(float(s.rstrip("GMT")) * mult)
|
|
|
+
|
|
|
+
|
|
|
+def tensor_span(root: Path, layer: int):
|
|
|
+ """(shard_path, rows, weight_offset, scale_offset) for one Engram layer."""
|
|
|
+ index = json.loads((root / "model.safetensors.index.json").read_text())["weight_map"]
|
|
|
+ prefix = f"layers.{layer}.engram.embed."
|
|
|
+ w_name = index[prefix + "weight"]
|
|
|
+ assert index[prefix + "scale"] == w_name, "weight and scale must share a shard"
|
|
|
+ path = root / w_name
|
|
|
+ with path.open("rb") as f:
|
|
|
+ length = struct.unpack("<Q", f.read(8))[0]
|
|
|
+ header = json.loads(f.read(length))
|
|
|
+ w, s = header[prefix + "weight"], header[prefix + "scale"]
|
|
|
+ assert w["dtype"] == "F8_E4M3" and s["dtype"] == "F8_E8M0", "unexpected dtype"
|
|
|
+ assert w["shape"][1] == WEIGHT_BYTES and s["shape"][1] == SCALE_BYTES
|
|
|
+ base = 8 + length
|
|
|
+ return path, w["shape"][0], base + w["data_offsets"][0], base + s["data_offsets"][0]
|
|
|
+
|
|
|
+
|
|
|
+def read_row(path: Path, w_off: int, s_off: int, row: int) -> bytes:
|
|
|
+ with path.open("rb", buffering=0) as f:
|
|
|
+ f.seek(w_off + row * WEIGHT_BYTES)
|
|
|
+ w = f.read(WEIGHT_BYTES)
|
|
|
+ f.seek(s_off + row * SCALE_BYTES)
|
|
|
+ s = f.read(SCALE_BYTES)
|
|
|
+ return w + s
|
|
|
+
|
|
|
+
|
|
|
+def interleave(weights: bytes, scales: bytes, n: int) -> bytes:
|
|
|
+ if np is not None:
|
|
|
+ buf = np.empty((n, RECORD_BYTES), dtype=np.uint8)
|
|
|
+ buf[:, :WEIGHT_BYTES] = np.frombuffer(weights, dtype=np.uint8).reshape(n, WEIGHT_BYTES)
|
|
|
+ buf[:, WEIGHT_BYTES:] = np.frombuffer(scales, dtype=np.uint8).reshape(n, SCALE_BYTES)
|
|
|
+ return buf.tobytes()
|
|
|
+ out = bytearray(n * RECORD_BYTES)
|
|
|
+ w = memoryview(weights)
|
|
|
+ s = memoryview(scales)
|
|
|
+ for i in range(n):
|
|
|
+ k = i * RECORD_BYTES
|
|
|
+ out[k : k + WEIGHT_BYTES] = w[i * WEIGHT_BYTES : (i + 1) * WEIGHT_BYTES]
|
|
|
+ out[k + WEIGHT_BYTES : k + SCALE_BYTES + WEIGHT_BYTES] = s[i * SCALE_BYTES : (i + 1) * SCALE_BYTES]
|
|
|
+ return bytes(out)
|
|
|
+
|
|
|
+
|
|
|
+def expected_size(rows: int, lo: int, hi: int) -> int:
|
|
|
+ return HEADER_BYTES + (hi - lo) * RECORD_BYTES
|
|
|
+
|
|
|
+
|
|
|
+class Ledger:
|
|
|
+ """Per-block ledger in the output dir: per-file {size, lo, hi, rows,
|
|
|
+ sample: {row: sha256(128B-slice of combined row)}} — content, not mtime."""
|
|
|
+
|
|
|
+ def __init__(self, path: Path):
|
|
|
+ self.path = path
|
|
|
+ self.files = {}
|
|
|
+ if path.exists():
|
|
|
+ self.data = json.loads(path.read_text())
|
|
|
+ self.files = self.data.get("files", {})
|
|
|
+
|
|
|
+ def known(self, name: str, size: int, lo: int, hi: int, rows: int) -> bool:
|
|
|
+ rec = self.files.get(name)
|
|
|
+ return (
|
|
|
+ rec is not None
|
|
|
+ and rec["size"] == size
|
|
|
+ and rec["lo"] == lo
|
|
|
+ and rec["hi"] == hi
|
|
|
+ and rec["rows"] == rows
|
|
|
+ )
|
|
|
+
|
|
|
+ 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.path.parent.mkdir(parents=True, exist_ok=True)
|
|
|
+ self.path.write_text(json.dumps({"files": self.files}))
|
|
|
+ self.path.chmod(0o644)
|
|
|
+
|
|
|
+
|
|
|
+def sample_rows(lo: int, hi: int, count: int) -> list:
|
|
|
+ n = hi - lo
|
|
|
+ if n <= count:
|
|
|
+ return list(range(lo, hi))
|
|
|
+ step = n / count
|
|
|
+ return [lo + int(i * step) for i in range(count)]
|
|
|
+
|
|
|
+
|
|
|
+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)."""
|
|
|
+ if not path.exists() or path.stat().st_size != size:
|
|
|
+ return False
|
|
|
+ rec = ledger.files.get(name)
|
|
|
+ if rec is None or not rec.get("sample"):
|
|
|
+ return None # size ok, no content samples — presence pass
|
|
|
+ with path.open("rb", buffering=0) as f:
|
|
|
+ for row in sorted(int(k) for k in rec["sample"]):
|
|
|
+ want = rec["sample"][str(row)]
|
|
|
+ h = hashlib.sha256(read_row(src, w_off, s_off, row)).digest()[:16].hex()
|
|
|
+ if want != h:
|
|
|
+ return False # source row changed
|
|
|
+ f.seek(HEADER_BYTES + (row - lo) * RECORD_BYTES)
|
|
|
+ h2 = hashlib.sha256(f.read(RECORD_BYTES)).digest()[:16].hex()
|
|
|
+ if h != h2:
|
|
|
+ return False # staged record corrupted
|
|
|
+ return True
|
|
|
+
|
|
|
+
|
|
|
+def main() -> int:
|
|
|
+ ap = argparse.ArgumentParser()
|
|
|
+ ap.add_argument("--id", required=True)
|
|
|
+ ap.add_argument("--model", required=True, help="checkpoint snapshot dir")
|
|
|
+ ap.add_argument("--out", required=True, help="node-local output dir")
|
|
|
+ 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("--min-free", default=None)
|
|
|
+ ap.add_argument("--force", action="store_true")
|
|
|
+ ap.add_argument("--verify", default="sample-hash", choices=["sample-hash", "presence"])
|
|
|
+ ap.add_argument("--chunk-rows", type=int, default=1 << 16)
|
|
|
+ ap.add_argument("--samples", type=int, default=128)
|
|
|
+ args = ap.parse_args()
|
|
|
+
|
|
|
+ root, out_dir = Path(args.model), Path(args.out)
|
|
|
+ out_dir.mkdir(parents=True, exist_ok=True)
|
|
|
+ ledger = Ledger(out_dir / f".sm121-staging-{args.id}.json")
|
|
|
+ layers = [int(x) for x in args.layers.split(",") if x.strip()]
|
|
|
+ if args.tp < 1 or not (0 <= args.rank < args.tp):
|
|
|
+ print(f"STAGING_FAIL invalid rank/tp {args.rank}/{args.tp}")
|
|
|
+ return 2
|
|
|
+
|
|
|
+ # min-free gate before any write (sum across layers)
|
|
|
+ need = 0
|
|
|
+ spans = {}
|
|
|
+ for layer in layers:
|
|
|
+ try:
|
|
|
+ path, rows, w_off, s_off = tensor_span(root, layer)
|
|
|
+ except Exception as e: # noqa: BLE001
|
|
|
+ print(f"STAGING_FAIL cannot read engram layer {layer} of {root}: {e}")
|
|
|
+ return 2
|
|
|
+ lo, hi = rows * args.rank // args.tp, rows * (args.rank + 1) // args.tp
|
|
|
+ spans[layer] = (path, rows, w_off, s_off, lo, hi)
|
|
|
+ need += expected_size(rows, lo, hi)
|
|
|
+ if args.min_free:
|
|
|
+ need_free = parse_size(args.min_free)
|
|
|
+ free = os.statvfs(out_dir).f_bavail * os.statvfs(out_dir).f_frsize
|
|
|
+ if free < need_free:
|
|
|
+ print(f"STAGING_FAIL {out_dir} free {free/2**30:.0f} GiB < min_free {args.min_free}")
|
|
|
+ return 2
|
|
|
+ free = os.statvfs(out_dir).f_bavail * os.statvfs(out_dir).f_frsize
|
|
|
+ if free < need:
|
|
|
+ print(f"STAGING_FAIL need {need/2**30:.1f} GiB, only {free/2**30:.1f} GiB at {out_dir}")
|
|
|
+ return 2
|
|
|
+
|
|
|
+ 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"
|
|
|
+ target = out_dir / name
|
|
|
+ size = expected_size(rows, lo, hi)
|
|
|
+ if not args.force and target.exists():
|
|
|
+ ok = verify_file(target, path, w_off, s_off, lo, hi, size, ledger, name)
|
|
|
+ if ok is True:
|
|
|
+ print(f"STAGED {name} rows={hi-lo} verified")
|
|
|
+ continue
|
|
|
+ if ok is None:
|
|
|
+ print(f"STAGED {name} rows={hi-lo} presence-ok")
|
|
|
+ continue
|
|
|
+ started = time.monotonic()
|
|
|
+ header = bytearray(HEADER_BYTES)
|
|
|
+ struct.pack_into("<6Q", header, 0, MAGIC, layer, lo, hi, rows, RECORD_BYTES)
|
|
|
+ partial = target.with_name(target.name + ".partial")
|
|
|
+ samples: dict = {}
|
|
|
+ rows_written = 0
|
|
|
+ with open(path, "rb", buffering=0) as src, open(partial, "wb", buffering=0) as dst:
|
|
|
+ dst.write(header)
|
|
|
+ for start in range(lo, hi, args.chunk_rows):
|
|
|
+ stop = min(start + args.chunk_rows, hi)
|
|
|
+ n = stop - start
|
|
|
+ src.seek(w_off + start * WEIGHT_BYTES)
|
|
|
+ weights = src.read(n * WEIGHT_BYTES)
|
|
|
+ src.seek(s_off + start * SCALE_BYTES)
|
|
|
+ scales = src.read(n * SCALE_BYTES)
|
|
|
+ if len(weights) != n * WEIGHT_BYTES or len(scales) != n * SCALE_BYTES:
|
|
|
+ print(f"STAGING_FAIL short read from {path} at row {start}")
|
|
|
+ partial.unlink(missing_ok=True)
|
|
|
+ return 2
|
|
|
+ dst.write(interleave(weights, scales, n))
|
|
|
+ rows_written += n
|
|
|
+ if rows_written % (16 * args.chunk_rows) == 0:
|
|
|
+ rate = rows_written * RECORD_BYTES / max(time.monotonic() - started, 1e-6)
|
|
|
+ print(f"PACK {name}: {100.0*rows_written/(hi-lo):5.1f}% {rate/2**20:.0f} MiB/s",
|
|
|
+ flush=True)
|
|
|
+ dst.flush()
|
|
|
+ os.fsync(dst.fileno())
|
|
|
+ partial.rename(target)
|
|
|
+ if args.verify == "sample-hash":
|
|
|
+ for row in sample_rows(lo, hi, args.samples):
|
|
|
+ h = hashlib.sha256(read_row(path, w_off, s_off, row)).digest()[:16].hex()
|
|
|
+ samples[row] = h
|
|
|
+ ledger.set(name, size, lo, hi, rows, samples)
|
|
|
+ secs = time.monotonic() - started
|
|
|
+ print(f"STAGED {name} rows={hi-lo} packed {size/2**30:.1f} GiB in {secs:.0f}s",
|
|
|
+ flush=True)
|
|
|
+ print("STAGING_DONE")
|
|
|
+ return 0
|
|
|
+
|
|
|
+
|
|
|
+if __name__ == "__main__":
|
|
|
+ sys.exit(main())
|
