spb/localvm-research Public License

1# =============================================================================2#  Project   : localvm-research3#  File      : experiments/candidate_01/implementation/streaming_verifier.py4#  Purpose   : Layer-streamed q8 verification for models larger than free RAM —5#              per-layer materialize → compute → re-lazify on unified memory6#  Author    : Simon-Pierre Boucher7#  Contact   : contact@spboucher.ai8#  Created   : 2026-08-129#  Modified  : 2026-08-1210#  Platform  : macOS / Apple Silicon (arm64) — MLX / Metal11#  License   : All rights reserved (research code)12# =============================================================================13"""StreamingVerifier — runs a verification forward pass through a quantized14model whose weights do NOT fit in free memory alongside the resident base.1516Mechanism: the model is built with lazy (mmap-backed) weights. During a17forward pass we walk the layers manually; each layer's weights materialize18on first use, and immediately after the layer's output is evaluated we19re-assign that layer's parameters to FRESH lazy arrays (a new mx.load view),20dropping the concrete buffers. Peak residency ≈ resident base + a few21layers, while the SSD sees one sequential pass over the checkpoint per22sweep — exactly the expH-friendly access pattern.23"""2425from __future__ import annotations2627import glob28import time29from pathlib import Path3031import mlx.core as mx32from mlx_lm import load as mlx_load33from mlx_lm.models.base import create_attention_mask343536class StreamingVerifier:37    def __init__(self, model_path: str):38        self.path = Path(model_path)39        # lazy=True: parameters are mmap-backed lazy arrays, nothing evaluated40        self.model, self.tokenizer = mlx_load(str(model_path), lazy=True)41        self.shards = sorted(glob.glob(str(self.path / "*.safetensors")))42        self.weight_bytes = sum(Path(s).stat().st_size for s in self.shards)43        self.last_sweep_io_s = 0.04445    def _fresh_lazy_weights(self) -> dict:46        w = {}47        for s in self.shards:48            w.update(mx.load(s))  # lazy by default: no eval performed49        return w5051    def _relazify(self, weights: dict, prefix: str) -> None:52        subset = [(k, v) for k, v in weights.items() if k.startswith(prefix)]53        if subset:54            self.model.load_weights(subset, strict=False)5556    def forward_chunk(self, chunk_ids: list[int], cache) -> mx.array:57        """Teacher-force `chunk_ids` through the model with per-layer weight58        streaming. `cache` is a make_prompt_cache(self.model) list; it is59        advanced by len(chunk_ids). Returns logits (T, vocab)."""60        t0 = time.perf_counter()61        fresh = self._fresh_lazy_weights()62        inner = self.model.model63        x = mx.array(chunk_ids)[None]64        h = inner.embed_tokens(x)65        mx.eval(h)66        self._relazify(fresh, "model.embed_tokens")67        mask = create_attention_mask(h, cache)68        for i, layer in enumerate(inner.layers):69            h = layer(h, mask, cache=cache[i] if cache else None)70            mx.eval(h)71            self._relazify(fresh, f"model.layers.{i}.")72            if (i + 1) % 8 == 0:73                mx.clear_cache()  # release Metal allocator pools74        h = inner.norm(h)75        if hasattr(self.model, "lm_head"):76            logits = self.model.lm_head(h)77        else:  # tied embeddings78            logits = inner.embed_tokens.as_linear(h)79        logits = logits[0]80        mx.eval(logits)81        self._relazify(fresh, "model.norm")82        self._relazify(fresh, "lm_head")83        mx.clear_cache()84        self.last_sweep_io_s = time.perf_counter() - t085        return logits86