diff --git a/scripts/verify-swift15-ple.py b/scripts/verify-swift15-ple.py new file mode 100755 index 0000000..a74f1c1 --- /dev/null +++ b/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 +""" +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(" 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)