diff --git a/scripts/verify-swift15-ple.py b/scripts/verify-swift15-ple.py index 2d83010..069193f 100755 --- a/scripts/verify-swift15-ple.py +++ b/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)