File size: 2,014 Bytes
939ace7 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 | #!/usr/bin/env python3
"""Synthetic Engram shards for server tests: same safetensors layout as model-00047/48 (an F8_E4M3 [rows, 256]
weight + F8_E8M0 [rows, 8] scale for one layer, then a trailing non-Engram tensor), random bytes, plus a test
manifest with their size and sha256. make_synth.py OUTDIR [--rows 2000003 --rows14 1500007]"""
import argparse
import hashlib
import json
import os
import struct
ap = argparse.ArgumentParser()
ap.add_argument("out")
ap.add_argument("--rows", type=int, default=2000003)
ap.add_argument("--rows14", type=int, default=1500007)
a = ap.parse_args()
os.makedirs(a.out, exist_ok=True)
manifest = {}
for layer, rows, name in ((1, a.rows, "model-00047-of-00048.safetensors"), (14, a.rows14, "model-00048-of-00048.safetensors")):
w, s, tail = rows * 256, rows * 8, 40960
hdr = {
f"layers.{layer}.engram.embed.weight": {"dtype": "F8_E4M3", "shape": [rows, 256], "data_offsets": [0, w]},
f"layers.{layer}.engram.embed.scale": {"dtype": "F8_E8M0", "shape": [rows, 8], "data_offsets": [w, w + s]},
f"layers.{layer}.engram.q_weight": {"dtype": "BF16", "shape": [4, 5120], "data_offsets": [w + s, w + s + tail]},
}
hb = json.dumps(hdr, separators=(",", ":")).encode()
if layer == 14:
hb += b" " * 8 # a different header length than layer 1: offsets that are not 4 KiB aligned either way
path = os.path.join(a.out, name)
h = hashlib.sha256()
with open(path, "wb") as f:
for part in (struct.pack("<Q", len(hb)), hb):
f.write(part)
h.update(part)
left = w + s + tail
while left:
blk = os.urandom(min(left, 8 << 20))
f.write(blk)
h.update(blk)
left -= len(blk)
manifest[name] = {"repo": "synthetic", "rev": "test", "size": os.path.getsize(path), "sha256": h.hexdigest()}
print(name, rows, manifest[name]["size"], manifest[name]["sha256"])
json.dump(manifest, open(os.path.join(a.out, "manifest.json"), "w"), indent=1)
|