File size: 3,050 Bytes
be3fd62 | 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 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 | #!/usr/bin/env python3
"""Export PyTorch tern_tc model to TC02 binary format for C inference."""
import argparse
import struct
import torch
import sys
def quantize_ternary(w):
g = w.abs().mean().clamp(min=1e-5)
wq = (w / g).round().clamp(-1, 1).to(torch.int8)
return wq, g.item()
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--ckpt", required=True, help="Path to PyTorch .pt model")
parser.add_argument("--out", required=True, help="Path to output .bin file")
args = parser.parse_args()
print(f"Loading {args.ckpt}...")
sd = torch.load(args.ckpt, map_location="cpu")
if "model" in sd:
sd = sd["model"]
# Determine config from weights
vocab_size, d_model = sd["wte.weight"].shape
n_layer = max(int(k.split(".")[1]) for k in sd.keys() if k.startswith("blocks.")) + 1
# Check layer 0 shapes
q_out, q_in = sd["blocks.0.attn.q.weight"].shape
k_out, k_in = sd["blocks.0.attn.k.weight"].shape
gate_out, gate_in = sd["blocks.0.mlp.gate.weight"].shape
assert q_in == d_model
d_ff = gate_out
# Calculate n_head and n_kv_head
# We know from tern_tc: n_head=5, n_kv_head=1, head_dim = d_model / n_head
# But wait, head_dim is usually d_model / n_head or explicit.
# Let's assume head_dim = 64 (standard for 320 d_model / 5)
head_dim = 64
n_head = q_out // head_dim
n_kv_head = k_out // head_dim
print(f"Detected Config: n_layer={n_layer}, n_head={n_head}, n_kv_head={n_kv_head}")
print(f" d_model={d_model}, d_ff={d_ff}, head_dim={head_dim}, vocab={vocab_size}")
with open(args.out, "wb") as f:
# Magic: TC02
f.write(b"TC02")
# Header: n_layer, n_head, n_kv_head, d_model, d_ff, head_dim, vocab_size
f.write(struct.pack("<7i", n_layer, n_head, n_kv_head, d_model, d_ff, head_dim, vocab_size))
def write_f32(t):
f.write(t.float().numpy().tobytes())
def write_ternary(name):
w = sd[name]
wq, g = quantize_ternary(w)
f.write(wq.numpy().tobytes())
f.write(struct.pack("<f", g))
print("Writing embeddings...")
write_f32(sd["wte.weight"])
for i in range(n_layer):
print(f"Writing layer {i}...")
prefix = f"blocks.{i}"
write_f32(sd[f"{prefix}.n1.weight"])
write_ternary(f"{prefix}.attn.q.weight")
write_ternary(f"{prefix}.attn.k.weight")
write_ternary(f"{prefix}.attn.v.weight")
write_ternary(f"{prefix}.attn.o.weight")
write_f32(sd[f"{prefix}.n2.weight"])
write_ternary(f"{prefix}.mlp.gate.weight")
write_ternary(f"{prefix}.mlp.up.weight")
write_ternary(f"{prefix}.mlp.down.weight")
print("Writing final norm...")
write_f32(sd["norm.weight"])
print(f"Exported to {args.out} successfully.")
if __name__ == "__main__":
main()
|