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

dsv41_fp4_stream.py

@@ -50,7 +50,7 @@ Source tensor naming (flat, no "model." prefix, "ffn" not "mlp"):
     layers.{L}.engram.embed.scale                e8m0   [rows, 8]
     mtp.{M}.ffn.experts.{E}...                   same as backbone
 """
-import argparse, json, os, re, shutil, sys, time
+import argparse, json, os, re, shutil, struct, sys, time
 import torch
 from safetensors import safe_open
 from safetensors.torch import save_file
@@ -97,7 +97,8 @@ def quantize_e2m1_codes(x: torch.Tensor) -> torch.Tensor:
     code, matching what the FP4 hardware conversion does.
     """
     mag = x.abs().clamp(max=FP4_MAX)
-    grid = FP4_TABLE[:8]                                   # 8 magnitudes, ascending
+    # boundaries must live on x's device -- torch.bucketize will not cross devices
+    grid = FP4_TABLE[:8].to(x.device)                      # 8 magnitudes, ascending
     # midpoints between consecutive grid points: .25 .75 1.25 1.75 2.5 3.5 5.0
     mid = (grid[1:] + grid[:-1]) / 2
     code = torch.bucketize(mag, mid)                       # 0..7
@@ -200,6 +201,124 @@ def quantize_engram_to_fp4(f, name, sname, rows, cols, chunk_rows=1 << 20, devic
 
 # ------------------------------------------------------------------ shard loop
 
+
+# ------------------------------------------- engram shard: stream straight to disk
+
+ST_DTYPE = {torch.uint8: ("U8", 1), torch.int8: ("I8", 1),
+            torch.float8_e4m3fn: ("F8_E4M3", 1), torch.float8_e8m0fnu: ("F8_E8M0", 1),
+            torch.bfloat16: ("BF16", 2), torch.float16: ("F16", 2),
+            torch.float32: ("F32", 4), torch.int32: ("I32", 4), torch.int64: ("I64", 8)}
+
+
+def write_engram_shard(src_path, dst_path, device="cpu", chunk_rows=1 << 20):
+    """Quantize and write an Engram shard without ever holding the table in RAM.
+
+    The obvious implementation -- build the packed tensor, hand the dict to save_file --
+    needs ~47 GiB resident for the weights alone, and on a GB10 the GPU allocator draws on
+    that same pool. It got OOM-killed at 83%. So we emit the safetensors container by hand:
+    the format is an 8-byte little-endian header length, that many bytes of JSON naming
+    each tensor's dtype/shape/byte-range, then the raw buffers back to back. Every offset
+    is known before a single value is computed, so the packed table can be streamed to the
+    file a chunk at a time.
+
+    Scales are the one thing still buffered (~3 GB for 384M rows x 8) because they are
+    produced by the same pass that produces the packed rows but live elsewhere in the file.
+
+    Peak memory is one chunk plus the scale buffer.
+    """
+    stats = {"engram": 0, "pass": 0, "cos_sum": 0.0, "cos_n": 0}
+    with safe_open(src_path, framework="pt") as f:
+        keys = list(f.keys())
+        emb = [k for k in keys if ENGRAM_RE.match(k)]
+        consumed = set()
+        plan = []                                   # (name, dtype_str, shape, kind, src)
+        for name in emb:
+            base = name[:-len(".weight")]
+            sname = base + ".scale"
+            rows, cols = f.get_slice(name).get_shape()
+            nblk = cols // SRC_BLOCK
+            plan.append((base + ".weight", "U8", [rows, cols // 2], "stream", (name, sname)))
+            plan.append((base + ".scale", "F8_E8M0", [rows, nblk], "scale", base))
+            consumed.add(name); consumed.add(sname)
+        for name in sorted(keys):
+            if name in consumed:
+                continue
+            plan.append((name, None, list(f.get_slice(name).get_shape()), "copy", name))
+
+        # resolve dtypes for the copied tensors (cheap: they are all small here)
+        resolved = []
+        for name, dts, shape, kind, srcref in plan:
+            if kind == "copy":
+                dts = ST_DTYPE[f.get_tensor(name).dtype][0]
+            resolved.append((name, dts, shape, kind, srcref))
+
+        header, off = {}, 0
+        for name, dts, shape, kind, srcref in resolved:
+            n = 1
+            for d in shape:
+                n *= d
+            nbytes = n * {"U8": 1, "I8": 1, "F8_E4M3": 1, "F8_E8M0": 1,
+                          "BF16": 2, "F16": 2, "F32": 4, "I32": 4, "I64": 8}[dts]
+            header[name] = {"dtype": dts, "shape": shape, "data_offsets": [off, off + nbytes]}
+            off += nbytes
+        header["__metadata__"] = {"format": "pt"}
+        blob = json.dumps(header).encode("utf-8")
+        pad = (-len(blob)) % 8                       # keep the data 8-byte aligned
+        blob += b" " * pad
+
+        scale_buf = {}
+        with open(dst_path, "wb", buffering=1 << 22) as out:
+            out.write(struct.pack("<Q", len(blob)))
+            out.write(blob)
+            for name, dts, shape, kind, srcref in resolved:
+                if kind == "stream":
+                    wname, sname = srcref
+                    base = name[:-len(".weight")]
+                    rows, cols = f.get_slice(wname).get_shape()
+                    nblk = cols // SRC_BLOCK
+                    sb = torch.empty((rows, nblk), dtype=torch.float8_e8m0fnu)
+                    wsl, ssl = f.get_slice(wname), f.get_slice(sname)
+                    t0 = time.time()
+                    for lo in range(0, rows, chunk_rows):
+                        hi = min(lo + chunk_rows, rows)
+                        n = hi - lo
+                        v = (wsl[lo:hi].to(device).float().view(n, nblk, SRC_BLOCK)
+                             * ssl[lo:hi].to(device).float().unsqueeze(-1))
+                        amax = v.abs().amax(dim=-1, keepdim=True)
+                        exp = torch.ceil(torch.log2((amax / FP4_MAX).clamp(min=1e-38))).clamp(-127, 127)
+                        ns = torch.pow(torch.tensor(2.0, device=v.device), exp)
+                        codes = quantize_e2m1_codes(v / ns)
+                        out.write(pack_e2m1(codes.view(n, cols)).to("cpu")
+                                  .flatten().view(torch.uint8).numpy().tobytes())
+                        sb[lo:hi] = ns.squeeze(-1).to(torch.float8_e8m0fnu).to("cpu")
+
+                        k = min(4096, n)
+                        a = v[:k].flatten(1)
+                        b = (FP4_TABLE.to(v.device)[codes[:k].long()].view(k, nblk, SRC_BLOCK)
+                             * ns[:k]).flatten(1)
+                        num = (a * b).sum(1); den = a.norm(dim=1) * b.norm(dim=1)
+                        ok = den > 0
+                        if ok.any():
+                            stats["cos_sum"] += float((num[ok] / den[ok]).sum())
+                            stats["cos_n"] += int(ok.sum())
+                        del v, codes, ns, amax, exp, a, b, num, den, ok
+                        print(f"      {base}: {hi}/{rows} ({100*hi/rows:.1f}%) "
+                              f"cos={stats['cos_sum']/max(stats['cos_n'],1):.6f} "
+                              f"[{time.time()-t0:.0f}s]", flush=True)
+                    scale_buf[base + ".scale"] = sb
+                    stats["engram"] += 1
+                elif kind == "scale":
+                    out.write(scale_buf.pop(srcref + ".scale")
+                              .flatten().view(torch.uint8).numpy().tobytes())
+                else:
+                    t = f.get_tensor(srcref).contiguous()
+                    # reinterpret as raw bytes for every dtype: numpy has no bfloat16 or
+                    # float8, so .numpy() is not an option on most of what passes through
+                    out.write(t.flatten().view(torch.uint8).numpy().tobytes())
+                    stats["pass"] += 1
+    return stats
+
+
 DEVICE = "cpu"
 
 
@@ -280,6 +399,14 @@ def main():
         sp = os.path.join(args.src, s)
         if not os.path.exists(sp):
             print(f"  SKIP (absent) {s}", flush=True); continue
+        has_engram = any(ENGRAM_RE.match(k) for k, v in index["weight_map"].items() if v == s)
+        if has_engram and do_g:
+            st = write_engram_shard(sp, os.path.join(args.out, s), device=DEVICE)
+            stats["engram"] += st["engram"]; stats["pass"] += st["pass"]
+            stats["engram_cos_sum"] += st["cos_sum"]; stats["engram_cos_n"] += st["cos_n"]
+            print(f"  {s}  [streamed engram shard] engram={st['engram']} pass={st['pass']} "
+                  f"cos={st['cos_sum']/max(st['cos_n'],1):.6f} [{time.time()-t0:.0f}s]", flush=True)
+            continue
         wmap, tb = process_shard(sp, os.path.join(args.out, s), do_e, do_g, stats)
         full_map.update(wmap); total += tb
         cos = (f" engram_cos={stats['engram_cos_sum']/stats['engram_cos_n']:.6f}"