# Tensor Roll — Recursive CUDA-Q Model Quantizer # Copyright (C) 2026 SnapKitty Collective # SPDX-License-Identifier: AGPL-3.0-or-later """Out-of-core execution — RULE ZERO: never require the teacher in RAM. The transformation must operate incrementally over a model larger than system RAM (and larger than RAM+VRAM combined): 30B MODEL STORAGE -> MANIFEST READER -> SHARD INDEX -> MMAP READER -> ONE TENSOR / BLOCK -> TENSOR ROLL -> 8B TARGET SHARD -> release source memory -> NEXT BLOCK New components (spec section 1): TensorSource manifest + shard index; never materializes the model ShardReader lazy, memory-mapped tensor reads MappedTensor memmap-backed tensor with exact shape/dtype TensorWindow recursive windowing of tensors larger than the budget RollWorkspace bounded working-set accounting with a hard ceiling StudentWriter incremental 8B-shard emission + verified shard index At no point does ingestion call load_model_30b() and hold all weights. The process respects --memory-budget within allocator overhead; the RollWorkspace refuses allocations that would exceed the ceiling. """ from __future__ import annotations import hashlib import json import os import time from dataclasses import dataclass, field from typing import Dict, Iterator, List, Optional, Tuple import numpy as np from .util import ensure_dir, utc_now_iso # --------------------------------------------------------------------------- # helpers # --------------------------------------------------------------------------- def parse_bytes(s) -> int: """'5G' -> bytes. Accepts int (bytes) or strings like '512M', '2G'.""" if isinstance(s, (int, float)): return int(s) s = str(s).strip().upper() for suffix, mult in (("TB", 1024 ** 4), ("GB", 1024 ** 3), ("G", 1024 ** 3), ("MB", 1024 ** 2), ("M", 1024 ** 2), ("KB", 1024), ("K", 1024), ("B", 1)): if s.endswith(suffix): return int(float(s[: -len(suffix)]) * mult) return int(float(s)) _DTYPE_ITEMSIZE = {"float32": 4, "float64": 8, "float16": 2, "int8": 1, "int32": 4, "int64": 8, "uint8": 1, "uint16": 2} def dtype_itemsize(dtype: str) -> int: return _DTYPE_ITEMSIZE[np.dtype(dtype).name] class MemoryBudgetExceeded(MemoryError): pass # --------------------------------------------------------------------------- # Manifests — the model is described, never loaded # --------------------------------------------------------------------------- @dataclass class TensorSpec: """One tensor's address in storage. No data attached.""" name: str shape: Tuple[int, ...] dtype: str nbytes: int shard: str # shard file name (relative to manifest dir) offset: int # byte offset inside the shard @property def params(self) -> int: n = 1 for d in self.shape: n *= d return n def to_dict(self) -> dict: return {"name": self.name, "shape": list(self.shape), "dtype": self.dtype, "nbytes": self.nbytes, "shard": self.shard, "offset": self.offset} @staticmethod def from_dict(d: dict) -> "TensorSpec": return TensorSpec(d["name"], tuple(d["shape"]), d["dtype"], d["nbytes"], d["shard"], d["offset"]) class VirtualManifest: """A model manifest with NO weight data — for structural scale tests. Lets the SCAN/PLAN machinery address a ~30B model without allocating a single weight. Any attempt to read data raises; only metadata flows. """ def __init__(self, specs: List[TensorSpec], *, model_name: str = "virtual", provenance: dict = None): self.specs = specs self.model_name = model_name self.provenance = provenance or {} @property def total_params(self) -> int: return sum(s.params for s in self.specs) @property def total_bytes(self) -> int: return sum(s.nbytes for s in self.specs) def to_dict(self) -> dict: return {"model": self.model_name, "virtual": True, "total_params": self.total_params, "total_bytes": self.total_bytes, "created": utc_now_iso(), "provenance": self.provenance, "tensors": [s.to_dict() for s in self.specs]} def write(self, path: str) -> dict: with open(path, "w") as f: json.dump(self.to_dict(), f, indent=2) return {"path": path, "tensors": len(self.specs), "total_params": self.total_params} @staticmethod def read(path: str) -> "VirtualManifest": with open(path) as f: d = json.load(f) return VirtualManifest( [TensorSpec.from_dict(s) for s in d["tensors"]], model_name=d.get("model", "virtual"), provenance=d.get("provenance", {})) def reader(self): raise RuntimeError( "virtual manifest has no weight data to read — " "metadata only (30B STRUCTURAL SCALE TEST)") def glimmer30b_manifest() -> VirtualManifest: """Realistic ~30B dense-transformer manifest (decoder-heavy, like the published Glimmer 30B: ~2B perception encoder + ~28B text decoder). Shapes are realistic transformer dimensions; NO weights are allocated. Used exclusively by the 30B STRUCTURAL SCALE TEST — this does not validate model quality, only that storage/planning address 30B under bounded RAM. """ specs: List[TensorSpec] = [] off = 0 def add(name, shape, dtype="float16", shard="decoder.safetensors"): nonlocal off nbytes = int(np.prod(shape)) * dtype_itemsize(dtype) specs.append(TensorSpec(name, tuple(shape), dtype, nbytes, shard, off)) off += nbytes # --- perception encoder (~2B): a few large projection blocks -------- d_v, n_v = 2048, 24 for l in range(n_v): add(f"vision.L{l}.Wqkv", (d_v, 3 * d_v)) add(f"vision.L{l}.Wo", (d_v, d_v)) add(f"vision.L{l}.W1", (d_v, 4 * d_v)) add(f"vision.L{l}.W2", (4 * d_v, d_v)) add(f"vision.L{l}.ln1", (d_v,), shard="decoder.safetensors") add(f"vision.L{l}.ln2", (d_v,), shard="decoder.safetensors") add("vision.proj", (d_v, 7168)) # --- text decoder (~28B): 44 layers, d=7168 -------------------------- d, dff, layers, vocab = 7168, 28672, 44, 128256 for s_idx in range(6): shard = f"decoder-{s_idx:05d}.safetensors" l0, l1 = s_idx * 8, min(layers, (s_idx + 1) * 8) for l in range(l0, l1): add(f"L{l}.Wq", (d, d), shard=shard) add(f"L{l}.Wk", (d, d), shard=shard) add(f"L{l}.Wv", (d, d), shard=shard) add(f"L{l}.Wo", (d, d), shard=shard) add(f"L{l}.W1", (d, dff), shard=shard) add(f"L{l}.W2", (dff, d), shard=shard) add(f"L{l}.ln1", (d,), shard=shard) add(f"L{l}.ln2", (d,), shard=shard) add("emb", (vocab, d), shard="decoder-00000.safetensors") add("Whead", (d, vocab), shard="decoder-00005.safetensors") add("lnf", (d,), shard="decoder-00005.safetensors") total = sum(s.params for s in specs) return VirtualManifest( specs, model_name="glimmer-30b-virtual", provenance={"note": "30B STRUCTURAL SCALE TEST manifest — shapes " "only, no weights allocated", "total_params": total}) # --------------------------------------------------------------------------- # Real shard storage: write once, read via mmap # --------------------------------------------------------------------------- class ShardWriter: """Writes tensors as raw little-endian blobs into shard files and emits a manifest.json describing every tensor's address.""" def __init__(self, root: str, shard_size_bytes: int = 256 * 1024 * 1024, dtype: str = "float32"): self.root = ensure_dir(root) self.shard_size = int(shard_size_bytes) self.specs: List[TensorSpec] = [] self._cur: Optional[str] = None self._fh = None self._off = 0 def _rotate(self): if self._fh: self._fh.close() idx = len([s for s in self.specs]) # unique shard index self._cur = f"shard-{idx:05d}.bin" self._fh = open(os.path.join(self.root, self._cur), "wb") self._off = 0 def add(self, name: str, arr: np.ndarray): arr = np.ascontiguousarray(arr) blob = arr.tobytes() if self._fh is None or self._off + len(blob) > self.shard_size: self._rotate() self._fh.write(blob) self.specs.append(TensorSpec( name, tuple(arr.shape), str(arr.dtype), len(blob), self._cur, self._off)) self._off += len(blob) def finalize(self, model_name: str = "teacher") -> str: if self._fh: self._fh.close() self._fh = None man = {"model": model_name, "virtual": False, "total_params": sum(s.params for s in self.specs), "total_bytes": sum(s.nbytes for s in self.specs), "created": utc_now_iso(), "tensors": [s.to_dict() for s in self.specs]} path = os.path.join(self.root, "manifest.json") with open(path, "w") as f: json.dump(man, f, indent=2) return path class TensorSource: """A model addressable tensor-by-tensor. Backed by real shard files (mmap reads) or by a virtual manifest (metadata only).""" def __init__(self, root: str): self.root = root man_path = os.path.join(root, "manifest.json") with open(man_path) as f: man = json.load(f) self.virtual = bool(man.get("virtual")) self.model_name = man.get("model", "?") self.specs = [TensorSpec.from_dict(s) for s in man["tensors"]] @property def total_params(self) -> int: return sum(s.params for s in self.specs) def names(self) -> List[str]: return [s.name for s in self.specs] def spec(self, name: str) -> TensorSpec: for s in self.specs: if s.name == name: return s raise KeyError(name) def reader(self) -> "ShardReader": if self.virtual: raise RuntimeError("virtual manifest has no weight data to read") return ShardReader(self) class ShardReader: """Lazy memory-mapped tensor reads. One tensor (or window) at a time.""" def __init__(self, source: TensorSource): self.source = source self._mmaps: Dict[str, np.memmap] = {} def _mmap(self, spec: TensorSpec) -> np.memmap: key = spec.shard if key not in self._mmaps: path = os.path.join(self.source.root, spec.shard) # map the whole shard file once; tensors are views into it size = os.path.getsize(path) self._mmaps[key] = np.memmap(path, dtype=np.uint8, mode="r", shape=(size,)) return self._mmaps[key] def read(self, name: str) -> "MappedTensor": spec = self.source.spec(name) raw = self._mmap(spec) view = raw[spec.offset:spec.offset + spec.nbytes].view( np.dtype(spec.dtype)).reshape(spec.shape) return MappedTensor(name, view, spec) def close(self): self._mmaps = {} class MappedTensor: """A memmap-backed tensor. `.window(slices)` yields TensorWindows.""" def __init__(self, name: str, view: np.ndarray, spec: TensorSpec): self.name = name self.view = view # memmap view — no copy self.spec = spec @property def shape(self): return self.view.shape def window(self, slices: Tuple[slice, ...]) -> "TensorWindow": return TensorWindow(self, slices) def windows(self, budget_bytes: int) -> Iterator["TensorWindow"]: """Recursive windowing until each window fits the budget.""" for sl in plan_windows(self.view.shape, self.view.dtype.itemsize, budget_bytes): yield TensorWindow(self, sl) @dataclass class TensorWindow: """One bounded piece of a (possibly huge) tensor, with exact reconstruction metadata. No silent shape corruption: the window records its slices, shape, and global offsets.""" parent: MappedTensor slices: Tuple[slice, ...] @property def name(self) -> str: idx = ",".join(f"{s.start}:{s.stop}" for s in self.slices) return f"{self.parent.name}@[{idx}]" @property def shape(self) -> Tuple[int, ...]: return self.data.shape @property def nbytes(self) -> int: return self.data.nbytes @property def data(self) -> np.ndarray: # materialize ONLY this window (bounded by the budget) return np.asarray(self.parent.view[self.slices]) def meta(self) -> dict: return {"tensor": self.parent.name, "slices": [[s.start, s.stop] for s in self.slices], "shape": list(self.shape), "global_shape": list(self.parent.shape)} def plan_windows(shape: Tuple[int, ...], itemsize: int, budget_bytes: int) -> List[Tuple[slice, ...]]: """Recursively split along the largest axis until every chunk's byte size <= budget_bytes. Returns exact slice tuples covering the tensor with no overlap and no gaps.""" total = int(np.prod(shape)) * itemsize if total <= budget_bytes or len(shape) == 0: return [tuple(slice(0, d) for d in shape)] # split the largest axis in half, recurse on both halves ax = int(np.argmax(shape)) mid = shape[ax] // 2 if mid == 0: raise MemoryBudgetExceeded( f"cannot window shape {shape}: axis {ax} has size 1 but a " f"single element ({itemsize}B) exceeds the budget") out = [] for lo, hi in ((0, mid), (mid, shape[ax])): sub = list(shape) sub[ax] = hi - lo for sl in plan_windows(tuple(sub), itemsize, budget_bytes): full = list(sl) s = full[ax] full[ax] = slice(s.start + lo, s.stop + lo) out.append(tuple(full)) return out def reassemble(windows: List[TensorWindow], shape: Tuple[int, ...], dtype=np.float64) -> np.ndarray: """Rebuild the full tensor from windows; verifies exact coverage.""" out = np.zeros(shape, dtype=dtype) seen = np.zeros(shape, dtype=bool) for w in windows: m = w.meta() sl = tuple(slice(a, b) for a, b in m["slices"]) assert w.data.shape == tuple(b - a for a, b in m["slices"]), \ f"window shape mismatch for {w.name}" out[sl] = w.data assert not seen[sl].any(), f"overlapping windows at {w.name}" seen[sl] = True assert seen.all(), "windows do not cover the tensor exactly" return out # --------------------------------------------------------------------------- # Bounded working set # --------------------------------------------------------------------------- class RollWorkspace: """Tracks the live working set against a hard memory ceiling. Every materialized window/buffer must be allocated through here. Allocations that would exceed the ceiling raise MemoryBudgetExceeded instead of silently paging the machine to death. """ def __init__(self, ceiling_bytes: int): self.ceiling = int(ceiling_bytes) self.live: Dict[str, int] = {} self.peak = 0 self.allocations = 0 @property def used(self) -> int: return sum(self.live.values()) def alloc(self, tag: str, nbytes: int): nbytes = int(nbytes) if self.used + nbytes > self.ceiling: raise MemoryBudgetExceeded( f"workspace ceiling {self.ceiling}B exceeded: " f"used={self.used}B + {nbytes}B ({tag})") self.live[tag] = self.live.get(tag, 0) + nbytes self.peak = max(self.peak, self.used) self.allocations += 1 def free(self, tag: str): self.live.pop(tag, None) def report(self) -> dict: return {"ceiling_bytes": self.ceiling, "peak_bytes": self.peak, "current_bytes": self.used, "allocations": self.allocations, "within_budget": self.peak <= self.ceiling} # --------------------------------------------------------------------------- # Incremental student emission # --------------------------------------------------------------------------- class StudentWriter: """Writes student shards incrementally + a verified shard index. Each emitted tensor is checksummed; finalize() writes shard_index.json with per-tensor sha256 and a total parameter count. Crash safety: flush_index() rewrites shard_index.json after every tensor (atomic replace). On resume, the writer reloads the existing index so previously emitted shards are kept and file names continue without collision. """ INDEX_NAME = "shard_index.json" def __init__(self, root: str, resume: bool = False): self.root = ensure_dir(root) self.entries: List[dict] = [] if resume: idx_path = os.path.join(root, self.INDEX_NAME) if os.path.exists(idx_path): with open(idx_path) as f: index = json.load(f) self.entries = index.get("tensors", []) def emit(self, name: str, arr: np.ndarray, meta: dict = None) -> dict: arr = np.ascontiguousarray(arr) fname = f"student-{len(self.entries):05d}.npy" path = os.path.join(self.root, fname) # np.save appends .npy when missing, so the temp name keeps it tmp = path + ".tmp.npy" np.save(tmp, arr) os.replace(tmp, path) # atomic: no half-written shards digest = hashlib.sha256(arr.tobytes()).hexdigest() entry = {"name": name, "file": fname, "shape": list(arr.shape), "dtype": str(arr.dtype), "params": int(arr.size), "sha256": digest, "meta": meta or {}} self.entries.append(entry) return entry def flush_index(self) -> dict: """Rewrite the index now (per-tensor checkpointing).""" return self.finalize() def finalize(self) -> dict: index = {"created": utc_now_iso(), "n_tensors": len(self.entries), "total_params": sum(e["params"] for e in self.entries), "tensors": self.entries} payload = json.dumps(index, sort_keys=True).encode() index["index_sha256"] = hashlib.sha256(payload).hexdigest() tmp = os.path.join(self.root, self.INDEX_NAME + ".tmp") with open(tmp, "w") as f: json.dump(index, f, indent=2) os.replace(tmp, os.path.join(self.root, self.INDEX_NAME)) return index @staticmethod def verify(root: str) -> dict: with open(os.path.join(root, "shard_index.json")) as f: index = json.load(f) for e in index["tensors"]: arr = np.load(os.path.join(root, e["file"])) assert list(arr.shape) == e["shape"], f"shape drift: {e['name']}" assert (hashlib.sha256( np.ascontiguousarray(arr).tobytes()).hexdigest() == e["sha256"]), f"checksum mismatch: {e['name']}" return {"tensors": index["n_tensors"], "total_params": index["total_params"], "ok": True} # --------------------------------------------------------------------------- # Checkpoints — crash-safe, corruption-rejecting # --------------------------------------------------------------------------- class Checkpoint: """JSON state + sha256. Corrupted checkpoints are refused, never loaded.""" def __init__(self, path: str): self.path = path def save(self, state: dict) -> dict: payload = json.dumps(state, sort_keys=True).encode() digest = hashlib.sha256(payload).hexdigest() tmp = self.path + ".tmp" with open(tmp, "w") as f: json.dump({"sha256": digest, "state": state}, f) os.replace(tmp, self.path) # atomic return {"path": self.path, "sha256": digest} def load(self) -> dict: with open(self.path) as f: doc = json.load(f) payload = json.dumps(doc["state"], sort_keys=True).encode() if hashlib.sha256(payload).hexdigest() != doc["sha256"]: raise ValueError( f"checkpoint checksum mismatch: {self.path} — refusing") return doc["state"] def exists(self) -> bool: return os.path.exists(self.path) # --------------------------------------------------------------------------- # Out-of-core driver: SCAN -> PLAN -> EXECUTE with resume # --------------------------------------------------------------------------- class OutOfCoreDriver: """Streams a teacher through TensorRoll without ever holding it in RAM. scan: lightweight per-tensor stats from the manifest (no weights). plan: global budget allocation (delegates to planner.plan). execute: per tensor (windowed as needed): mmap -> roll -> transform -> emit student shard -> release. Checkpointed per tensor so --resume continues after a crash. """ def __init__(self, source: TensorSource, workdir: str, memory_budget: str = "5G", target_params: float = 8.0e9, cap_params: float = 8.5e9): self.source = source self.workdir = ensure_dir(workdir) self.ws = RollWorkspace(parse_bytes(memory_budget)) self.target = float(target_params) self.cap = float(cap_params) self.ckpt = Checkpoint(os.path.join(workdir, "ooc_checkpoint.json")) self.student_dir = ensure_dir(os.path.join(workdir, "student_shards")) def scan(self) -> List[dict]: """Lightweight stats per tensor — manifest only, no weight data.""" from . import core as C # deferred: keeps import graph light rows = [] for s in self.source.specs: rows.append({ "tensor": s.name, "shape": list(s.shape), "dtype": s.dtype, "params": s.params, "bytes": s.nbytes, "shard": s.shard, # shape-derived proxies, measured from metadata only "is_matrix": len(s.shape) == 2 and min(s.shape) > 1, "aspect": (max(s.shape) / max(1, min(s.shape)) if s.shape else 1.0), }) return rows def plan(self, rows: List[dict]) -> dict: from . import planner as P return P.plan(rows, target=self.target, cap=self.cap) def execute(self, plan: dict, transform, *, resume: bool = False, fault_after: int = -1) -> dict: """Run the plan. `transform(name, window_data, w_target, ws, meta)` is the caller-supplied per-window transform returning either one student array or a dict {suffix: array} (e.g. factor pairs). Checkpointed per tensor; the student index is flushed per tensor too, so `fault_after` crash tests resume with zero emitted-data loss. `fault_after` injects a crash (tests only, never production). """ state = self.ckpt.load() if (resume and self.ckpt.exists()) else { "done": [], "emitted_params": 0} done = set(state["done"]) writer = StudentWriter(self.student_dir, resume=resume) # re-verify previously emitted shards on resume if resume and done: StudentWriter.verify(self.student_dir) reader = self.source.reader() processed, emitted = 0, state["emitted_params"] try: for spec in self.source.specs: if spec.name in done: continue mt = reader.read(spec.name) for w in mt.windows(self.ws.ceiling): tag = f"win:{w.name}" self.ws.alloc(tag, w.nbytes) try: data = w.data target = plan["tensors"][spec.name]["target_params"] # split the tensor's target budget across windows # proportionally to window params w_target = max( 1, int(target * w.data.size / spec.params)) out = transform(spec.name, data, w_target, self.ws, w.meta()) finally: self.ws.free(tag) parts = out if isinstance(out, dict) else {"": out} for suffix, arr in parts.items(): arr = np.asarray(arr) ename = w.name + (f".{suffix}" if suffix else "") writer.emit(ename, arr, {"parent": spec.name, **w.meta()}) emitted += int(arr.size) # release the memmap view for this tensor before next del mt processed += 1 done.add(spec.name) state = {"done": sorted(done), "emitted_params": emitted} self.ckpt.save(state) writer.flush_index() if fault_after >= 0 and processed > fault_after: raise RuntimeError("injected crash (fault_after test)") finally: reader.close() index = writer.finalize() return {"processed_tensors": processed, "emitted_params": emitted, "index_total_params": index["total_params"], "workspace": self.ws.report(), "budget_ok": index["total_params"] <= self.cap} # --------------------------------------------------------------------------- # External teacher formats: safetensors (true mmap) and npz (streaming repack) # --------------------------------------------------------------------------- _SAFETENSORS_DTYPES = { "F32": "float32", "F16": "float16", "I32": "int32", "I64": "int64", "U8": "uint8", "I8": "int8", "BOOL": "bool", } def _read_safetensors_header(path: str) -> Tuple[dict, int]: """Returns (header_dict, data_start_offset). Metadata only.""" import struct with open(path, "rb") as f: n = struct.unpack(" abs path for fname in sorted(os.listdir(path)): if not fname.endswith(".safetensors"): continue fpath = os.path.join(path, fname) header, data_start = _read_safetensors_header(fpath) self._files[fname] = fpath for name, info in header.items(): if name == "__metadata__": continue dt = info["dtype"] if dt == "BF16": try: import ml_dtypes # noqa np_dtype = "bfloat16" except ImportError: raise RuntimeError( f"tensor {name}: BF16 needs the ml_dtypes " f"package for CPU decode; refusing to guess") elif dt not in _SAFETENSORS_DTYPES: raise RuntimeError( f"tensor {name}: unsupported dtype {dt}") else: np_dtype = _SAFETENSORS_DTYPES[dt] shape = tuple(info["shape"]) start, end = info["data_offsets"] nbytes = end - start self.specs.append(TensorSpec( f"{fname}:{name}", shape, np_dtype, nbytes, fname, data_start + start)) if not self.specs: raise ValueError(f"no safetensors tensors found in {path}") @property def total_params(self) -> int: return sum(s.params for s in self.specs) def names(self) -> List[str]: return [s.name for s in self.specs] def spec(self, name: str) -> TensorSpec: for s in self.specs: if s.name == name: return s raise KeyError(name) def reader(self) -> "SafetensorsReader": return SafetensorsReader(self) class SafetensorsReader: """Lazy memory-mapped reads over safetensors files.""" def __init__(self, source: SafetensorsSource): self.source = source self._mmaps: Dict[str, np.memmap] = {} def _mmap(self, spec: TensorSpec) -> np.memmap: if spec.shard not in self._mmaps: path = self.source._files[spec.shard] size = os.path.getsize(path) self._mmaps[spec.shard] = np.memmap( path, dtype=np.uint8, mode="r", shape=(size,)) return self._mmaps[spec.shard] def read(self, name: str) -> MappedTensor: spec = self.source.spec(name) raw = self._mmap(spec) view = raw[spec.offset:spec.offset + spec.nbytes].view( np.dtype(spec.dtype)).reshape(spec.shape) return MappedTensor(name, view, spec) def close(self): self._mmaps = {} def _scan_npz(path: str) -> List[Tuple[str, tuple, str, int]]: """Metadata-only scan of one .npz (zip headers, no array data).""" import zipfile from numpy.lib import format as npfmt specs = [] with zipfile.ZipFile(path) as zf: for entry in zf.namelist(): if not entry.endswith(".npy"): continue with zf.open(entry) as f: npfmt.read_magic(f) shape, _fortran, dtype = npfmt.read_array_header_1_0(f) name = entry[:-4] specs.append((name, tuple(shape), np.dtype(dtype).name, int(np.prod(shape, dtype=np.int64)) * np.dtype(dtype).itemsize)) return specs def prepare_teacher_source(teacher_dir: str, cache_dir: str): """Open a teacher weight directory as a streaming source. .safetensors -> true mmap source (no copy). .npz -> one-time streaming repack into the shard format (one array in RAM at a time), then a TensorSource over the repacked shards. Raises informatively when the directory holds neither. """ teacher_dir = os.path.abspath(teacher_dir) st_files = sorted(f for f in os.listdir(teacher_dir) if f.endswith(".safetensors")) npz_files = sorted(f for f in os.listdir(teacher_dir) if f.endswith(".npz")) if st_files: return SafetensorsSource(teacher_dir) if npz_files: shard_root = ensure_dir(os.path.join(cache_dir, "repacked_shards")) man_path = os.path.join(shard_root, "manifest.json") if not os.path.exists(man_path): import zipfile writer = ShardWriter(shard_root) for fname in npz_files: specs = _scan_npz(os.path.join(teacher_dir, fname)) zf = zipfile.ZipFile(os.path.join(teacher_dir, fname)) for name, _shape, _dtype, _nbytes in specs: # one array in RAM at a time — the streaming invariant with zf.open(name + ".npy") as f: buf = f.read() import io arr = np.load(io.BytesIO(buf)) writer.add(f"{fname}:{name}", np.ascontiguousarray(arr)) del arr, buf zf.close() writer.finalize("teacher-repacked") return TensorSource(shard_root) raise ValueError( f"no .safetensors or .npz weights in {teacher_dir}") # --------------------------------------------------------------------------- # Streaming compression: per-window truncated SVD under the planned budget # --------------------------------------------------------------------------- def stream_compress(source, plan: dict, memory_budget: str, out_dir: str, *, resume: bool = False, fault_after: int = -1, dtype=np.float32) -> dict: """End-to-end out-of-core compression. For each tensor: mmap -> recursive windows (each <= budget) -> truncated SVD at the tensor's planned rank share -> emit low-rank factor pairs (A = U_r S_r, B = V_r^T) as student shards. Student parameter count = sum of factor params, honestly accounted against the plan. Per-tensor checkpoint + per-tensor index flush; resume continues after a crash with zero emitted-data loss. """ workdir = ensure_dir(out_dir) driver = OutOfCoreDriver( source, workdir, memory_budget=memory_budget, target_params=plan.get("total_target_params", 8.0e9), cap_params=plan.get("cap", 8.5e9)) stats: Dict[str, dict] = {} rep_path = os.path.join(workdir, "stream_compress_report.json") # on resume, keep the previous run's per-tensor stats for tensors we # skip this time (otherwise a resume-noop would report empty stats) prior_stats: Dict[str, dict] = {} if resume and os.path.exists(rep_path): try: with open(rep_path) as f: prior_stats = json.load(f).get("tensor_stats", {}) except (json.JSONDecodeError, OSError): prior_stats = {} def transform(name, data, w_target, ws, meta): parent = meta["tensor"] W = np.asarray(data) st = stats.setdefault(parent, {"windows": 0, "e_num": 0.0, "e_den": 0.0, "rank": 0}) if W.ndim == 2 and min(W.shape) > 1: m, n = W.shape R = max(1, int(w_target // (m + n))) st["rank"] = max(st["rank"], R) Wd = W.astype(np.float64) U, s, Vt = np.linalg.svd(Wd, full_matrices=False) r = max(1, min(m, n, R, s.size)) A = (U[:, :r] * s[:r]).astype(dtype) B = Vt[:r, :].astype(dtype) st["e_num"] += float((s[:r] ** 2).sum()) st["e_den"] += float((s ** 2).sum()) st["windows"] += 1 return {"A": A, "B": B} st["windows"] += 1 return {"raw": np.asarray(W).astype(dtype)} t0 = time.perf_counter() res = driver.execute(plan, transform, resume=resume, fault_after=fault_after) dt = time.perf_counter() - t0 # per-tensor retained energy (measured); re-run tensors overwrite for _name, st in stats.items(): st["retained_energy"] = (st["e_num"] / st["e_den"] if st["e_den"] > 0 else 1.0) del st["e_num"], st["e_den"] merged_stats = {**prior_stats, **stats} res.update({"elapsed_s": dt, "tensor_stats": merged_stats, "source_params": source.total_params, "compression_ratio": (res["index_total_params"] / max(1, source.total_params))}) rep_path = os.path.join(workdir, "stream_compress_report.json") with open(rep_path, "w") as f: json.dump(res, f, indent=2, default=float) res["report"] = rep_path return res