swift15: PLE fixed-scale audit (streaming BF16 absmax vs 448*scale cap)

scripts/verify-swift15-ple.py

@@ -0,0 +1,71 @@
+#!/usr/bin/env python3
+"""Streaming absmax audit of the Swift1.5 PLE n-gram table (BF16).
+
+The PLE fp8 encoding uses ONE fixed scale (1.9931793212890625e-04, the
+official checkpoint's plateau value). fp8_e4m3 saturates at 448, so the
+fixed scale is legal iff absmax <= 448 * SCALE. This walks every
+ngram_embedding.shard_*.weight tensor in the checkpoint (mmap, no copy),
+computes the absmax from the BF16 bit pattern (clear the sign bit, max
+the remaining uint16 — BF16 magnitudes are monotone in their bits), and
+reports the headroom. Exit 0 only when headroom >= 1.5x.
+
+Run from any cwd EXCEPT /tmp on dankserver (/tmp/inspect.py shadows the
+stdlib). Usage: verify-swift15-ple.py <ckpt-dir>
+"""
+import json
+import mmap
+import os
+import struct
+import sys
+
+SCALE = 1.9931793212890625e-04  # bf16 0x3951, the official plateau scale
+FP8_E4M3_MAX = 448.0
+
+import numpy as np
+
+ckpt = sys.argv[1]
+index = json.load(open(os.path.join(ckpt, "model.safetensors.index.json")))
+shards: dict[str, list[str]] = {}
+for k, f in index["weight_map"].items():
+    if "ngram_embedding.shard_" in k and k.endswith(".weight"):
+        shards.setdefault(f, []).append(k)
+
+files = len(shards)
+tensors = sum(len(v) for v in shards.values())
+print(f"n-gram shards: {files} files, {tensors} tensors")
+
+import math
+
+gmax = 0.0
+gfile = ""
+for fname, keys in sorted(shards.items()):
+    path = os.path.join(ckpt, fname)
+    with open(path, "rb") as fh:
+        hlen = struct.unpack("<Q", fh.read(8))[0]
+        header = json.loads(fh.read(hlen))
+        base = 8 + hlen
+        mm = mmap.mmap(fh.fileno(), 0, prot=mmap.PROT_READ)
+        for k in keys:
+            meta = header[k]
+            off, end = meta["data_offsets"]
+            assert meta["dtype"] == "BF16", f"{k}: dtype {meta['dtype']}"
+            a = np.frombuffer(mm, dtype=np.uint16, count=(end - off) // 2,
+                              offset=base + off)
+            m = int((a & 0x7FFF).max())
+            if m > gmax:
+                gmax = m
+                gfile = k
+        mm.close()
+
+# rebuild the BF16 value from the sign-cleared uint16 bit pattern
+absmax = np.frombuffer(np.uint16(gmax).tobytes() + b"\x00\x00",
+                       dtype=np.float32)[0]
+cap = FP8_E4M3_MAX * SCALE
+print(f"absmax        = {absmax!r}  (tensor: {gfile})")
+print(f"2^log2        = 2^{math.log2(absmax):.3f}")
+print(f"fp8 cap       = 448 * {SCALE:.8g} = {cap:.8g}")
+print(f"headroom      = {cap / absmax:.3f}x")
+ok = cap / absmax >= 1.5
+print("VERDICT:", "PASS (>= 1.5x headroom over the fixed scale)" if ok
+      else "FAIL — fixed scale does not cover the Swift PLE")
+sys.exit(0 if ok else 1)