88 lines
3.3 KiB
Python
Executable File
88 lines
3.3 KiB
Python
Executable File
#!/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
|
|
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")))
|
|
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
|
|
|
|
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:
|
|
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"]
|
|
a = np.frombuffer(mm, dtype=np.uint16, count=(end - off) // 2,
|
|
offset=base + off)
|
|
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
|
|
gfile = k
|
|
mm.close()
|
|
|
|
# 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]
|
|
|
|
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")
|
|
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)
|