Final card and build scripts: measured numbers, Engram caveat, pruning analysis

test_engram_shard.py

@@ -0,0 +1,70 @@
+#!/usr/bin/env python3
+"""Check the hand-written safetensors container round-trips before it meets a 94.6 GiB table.
+
+Builds a shard shaped like the real ones (one Engram table plus the small q/k/wkv tensors),
+runs write_engram_shard, then reopens the result with safetensors and verifies:
+  * it parses at all (the header length, padding and offsets are right)
+  * every tensor is present with the right dtype and shape
+  * pass-through tensors are byte-identical
+  * the packed table dequantizes to the same values the in-RAM path produced
+"""
+import os, sys, tempfile
+import torch
+from safetensors.torch import save_file
+from safetensors import safe_open
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+import dsv41_fp4_stream as Q
+
+ROWS, COLS = 250_000, 256
+NBLK = COLS // Q.SRC_BLOCK
+torch.manual_seed(1)
+
+raw = (torch.randn(ROWS, COLS) * torch.empty(ROWS, 1).uniform_(0.05, 4.0)).clamp(-448, 448)
+tensors = {
+    "layers.1.engram.embed.weight": raw.to(torch.float8_e4m3fn),
+    "layers.1.engram.embed.scale": torch.randint(110, 140, (ROWS, NBLK), dtype=torch.uint8).view(torch.float8_e8m0fnu),
+    "layers.1.engram.q_weight": torch.randn(4, 5120).bfloat16(),
+    "layers.1.engram.k_weight": torch.randn(4, 5120).bfloat16(),
+    "layers.1.engram.wkv.weight": torch.randn(256, 6144).to(torch.float8_e4m3fn),
+    "layers.1.engram.wkv.scale": torch.randint(120, 132, (8, 192), dtype=torch.uint8).view(torch.float8_e8m0fnu),
+}
+
+DEVS = ["cpu"] + (["cuda"] if torch.cuda.is_available() else [])
+with tempfile.TemporaryDirectory() as td:
+    src = os.path.join(td, "s.safetensors")
+    save_file(tensors, src, metadata={"format": "pt"})
+
+    for dev in DEVS:
+        dst = os.path.join(td, f"o-{dev}.safetensors")
+        st = Q.write_engram_shard(src, dst, device=dev, chunk_rows=100_000)
+        print(f"\n--- device={dev} --- engram={st['engram']} pass={st['pass']} "
+              f"cos={st['cos_sum']/st['cos_n']:.6f}")
+
+        with safe_open(dst, framework="pt") as g:
+            got = set(g.keys())
+            want = set(tensors)
+            assert got == want, f"key mismatch\n got {sorted(got)}\n want {sorted(want)}"
+            w = g.get_tensor("layers.1.engram.embed.weight")
+            sc = g.get_tensor("layers.1.engram.embed.scale")
+            assert w.dtype == torch.uint8 and tuple(w.shape) == (ROWS, COLS // 2), (w.dtype, w.shape)
+            assert sc.dtype == torch.float8_e8m0fnu and tuple(sc.shape) == (ROWS, NBLK)
+            for k in ["layers.1.engram.q_weight", "layers.1.engram.k_weight",
+                      "layers.1.engram.wkv.weight", "layers.1.engram.wkv.scale"]:
+                a, b = tensors[k], g.get_tensor(k)
+                assert a.dtype == b.dtype and a.shape == b.shape, (k, a.dtype, b.dtype)
+                assert bool((a.view(torch.uint8) == b.view(torch.uint8)).all()), f"{k} changed"
+            print(f"    keys/dtypes/shapes OK; pass-through tensors byte-identical")
+
+            # compare against the in-RAM path
+            with safe_open(src, framework="pt") as f2:
+                p2, s2, _, _ = Q.quantize_engram_to_fp4(
+                    f2, "layers.1.engram.embed.weight", "layers.1.engram.embed.scale",
+                    ROWS, COLS, chunk_rows=100_000, device=dev)
+            assert bool((p2 == w).all()), "streamed weights differ from the in-RAM path"
+            assert bool((s2.view(torch.uint8) == sc.view(torch.uint8)).all()), "scales differ"
+            print(f"    matches the in-RAM path bit-for-bit")
+
+        print(f"    file {os.path.getsize(dst):,} B  (source {os.path.getsize(src):,} B)")
+
+print("\nSTREAMED ENGRAM SHARD OK")