swift15: PLE audit adds over-cap census + next-bf16 scale suggestion

scripts/verify-swift15-ple.py

@@ -21,7 +21,10 @@ import sys
 SCALE = 1.9931793212890625e-04  # bf16 0x3951, the official plateau scale
 FP8_E4M3_MAX = 448.0
 
-import numpy as np
+gmax = 0.0
+gfile = ""
+over_cap = 0
+total = 0
 
 ckpt = sys.argv[1]
 index = json.load(open(os.path.join(ckpt, "model.safetensors.index.json")))
@@ -36,8 +39,10 @@ print(f"n-gram shards: {files} files, {tensors} tensors")
 
 import math
 
-gmax = 0.0
-gfile = ""
+cap = FP8_E4M3_MAX * SCALE
+# largest bf16 <= cap: truncate the float32 bits (cap > 0); bf16
+# magnitudes are monotone in their bits, so |w| > cap <=> bits > cap_bits
+cap_bits = int((np.float32(cap).view(np.uint32) >> 16) & 0xFFFF)
 for fname, keys in sorted(shards.items()):
     path = os.path.join(ckpt, fname)
     with open(path, "rb") as fh:
@@ -50,7 +55,10 @@ for fname, keys in sorted(shards.items()):
             off, end = meta["data_offsets"]
             a = np.frombuffer(mm, dtype=np.uint16, count=(end - off) // 2,
                               offset=base + off)
-            m = int((a & 0x7FFF).max())
+            v = a & 0x7FFF
+            m = int(v.max())
+            over_cap += int((v > np.uint16(cap_bits)).sum())
+            total += v.size
             del a
             if m > gmax:
                 gmax = m
@@ -60,7 +68,7 @@ for fname, keys in sorted(shards.items()):
 # rebuild the BF16 value: the 16 bits are the HIGH half of the float32
 absmax = np.frombuffer(b"\x00\x00" + np.uint16(gmax).tobytes(),
                         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}")
@@ -68,4 +76,11 @@ 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")
+print(f"values > cap  = {over_cap} of {total} ({over_cap / total:.3e})")
+nxt_bits = int((np.float32(absmax / FP8_E4M3_MAX).view(np.uint32)
+                >> 16) & 0xFFFF) + 1
+nxt_scale = np.frombuffer(b"\x00\x00" + np.uint16(nxt_bits).tobytes(),
+                          dtype=np.float32)[0]
+print(f"next-bf16 scale covering absmax: {nxt_scale:.8g} "
+      f"(0x{nxt_bits:04x}, headroom {FP8_E4M3_MAX * nxt_scale / absmax:.4f}x)")
 sys.exit(0 if ok else 1)