File size: 6,820 Bytes
2049fc8 | 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 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 | """TinyLLama-v0 int4 for the 20-bit WROM, with per-group bases, row scales and tables."""
import os
import struct
import sys
import numpy as np
from safetensors import safe_open
HERE = os.path.dirname(os.path.abspath(__file__))
SRC = os.environ.get("TL_SRC", os.path.join(os.path.dirname(os.path.abspath(__file__)), "model.safetensors"))
WORDS_PER_UNIT = 15 * 64
DIM, NL, NH, HID, VOCAB, MAXPOS = 64, 8, 16, 256, 32000, 512
HS, KVD = DIM // NH, DIM
def rope_perm(w):
idx = [(m // 2) + (m % 2) * (HS // 2) for m in range(HS)]
return w.reshape(NH, HS, -1)[:, idx, :].reshape(NH * HS, -1)
def packed(unit, off):
return (unit << 10) | off
class Placer:
def __init__(self):
self.unit = 0
self.off = 0
self.cells = bytearray()
def _grow(self, wp):
need = 8 * (wp + 1)
if len(self.cells) < need:
self.cells.extend(bytes(need - len(self.cells)))
def add_nibbles(self, nibs):
first = None
for i in range(0, len(nibs), 8):
if self.off + 1 > WORDS_PER_UNIT:
self.unit += 1
self.off = 0
if first is None:
first = packed(self.unit, self.off)
wp = packed(self.unit, self.off)
self._grow(wp)
for k in range(8):
self.cells[8 * wp + k] = int(nibs[i + k]) & 15
self.off += 1
return first
def add(self, q):
rows, n = q.shape
wpr = n // 8
gw = 4 * wpr
first = None
bases = []
for g0 in range(0, rows, 4):
if self.off + gw > WORDS_PER_UNIT:
self.unit += 1
self.off = 0
if first is None:
first = packed(self.unit, self.off)
wp0 = packed(self.unit, self.off)
bases.append(wp0)
for r in range(4):
row = q[g0 + r] if g0 + r < rows else np.zeros(n, np.int8)
for w in range(wpr):
wp = wp0 + r * wpr + w
self._grow(wp)
for k in range(8):
self.cells[8 * wp + k] = int(row[8 * w + k] + 8) & 15
self.off += gw
return first
def quantize(w, grid=13):
w = w.astype(np.float32)
best = None
for c in np.linspace(0.5, 1.0, grid):
s = np.abs(w).max(-1, keepdims=True) * c / 7.0
s[s == 0] = 1e-12
q = np.clip(np.round(w / s), -8, 7)
e = ((q * s - w) ** 2).sum(-1, keepdims=True)
if best is None:
best = (q, s, e)
else:
m = e < best[2]
best = (np.where(m, q, best[0]), np.where(m, s, best[1]), np.minimum(e, best[2]))
return best[0].astype(np.int8), best[1][..., 0].astype(np.float64)
def carr(name, ctype, a):
a = np.asarray(a).ravel()
return f"static const {ctype} {name}[{len(a)}] = {{{','.join(str(int(v)) for v in a)}}};\n"
def load():
t = {}
with safe_open(SRC, "pt") as f:
for k in f.keys():
t[k] = f.get_tensor(k).float().numpy()
return t
def build():
t = load()
P = Placer()
Q, S, W = {}, {}, {}
SC = {}
def pack_scales(v):
rows = len(v)
nibs = []
for g in range(0, rows, 4):
for pair in range(2):
for r in range(2):
idx = g + pair * 2 + r
x = int(v[idx]) if idx < rows else 0
nibs += [(x >> (4 * j)) & 15 for j in range(4)]
return nibs
def put(name, w):
q, s = quantize(w)
Q[name], S[name] = q, np.round(s * 2**15).astype(np.int64)
assert S[name].max() < 65536, (name, S[name].max())
if name.startswith("down"):
W[name] = [P.add(q[:, 64 * c:64 * c + 64]) for c in range(4)]
else:
W[name] = P.add(q)
SC[name] = P.add_nibbles(pack_scales(S[name]))
fin = t["model.norm.weight"]
put("lm_head", t["lm_head.weight"] * fin[None, :])
put("embed", t["model.embed_tokens.weight"])
for l in range(NL):
ga = t[f"model.layers.{l}.input_layernorm.weight"]
gf = t[f"model.layers.{l}.post_attention_layernorm.weight"]
put(f"q{l}", rope_perm(t[f"model.layers.{l}.self_attn.q_proj.weight"] * ga[None, :]))
put(f"k{l}", rope_perm(t[f"model.layers.{l}.self_attn.k_proj.weight"] * ga[None, :]))
put(f"v{l}", t[f"model.layers.{l}.self_attn.v_proj.weight"] * ga[None, :])
put(f"o{l}", t[f"model.layers.{l}.self_attn.o_proj.weight"])
put(f"gate{l}", t[f"model.layers.{l}.mlp.gate_proj.weight"] * gf[None, :])
put(f"up{l}", t[f"model.layers.{l}.mlp.up_proj.weight"] * gf[None, :])
put(f"down{l}", t[f"model.layers.{l}.mlp.down_proj.weight"])
return P, Q, S, W, SC
def emit(P, Q, S, W, SC):
h = ["#include <stdint.h>\n"]
h.append(f"#define DIM {DIM}\n#define HID {HID}\n#define NL {NL}\n#define NH {NH}\n"
f"#define HS {HS}\n#define KVD {KVD}\n#define VOCAB {VOCAB}\n#define MAXPOS {MAXPOS}\n")
h.append(f"#define ATT_C {round(2 ** 31 / (127 ** 2 * np.sqrt(HS)))}\n")
h.append(f"#define G_HEAD {W['lm_head']}\n#define G_EMB {W['embed']}\n")
for nm in ("q", "k", "v", "o", "gate", "up"):
h.append(carr(f"G_{nm}", "int32_t", [W[f"{nm}{l}"] for l in range(NL)]))
h.append(carr("G_down", "int32_t", [W[f"down{l}"][c] for l in range(NL) for c in range(4)]))
h.append(f"#define SC_HEAD {SC['lm_head']}\n#define SC_EMB {SC['embed']}\n")
for nm in ("q", "k", "v", "o", "gate", "up", "down"):
h.append(carr(f"SC_{nm}", "int32_t", [SC[f"{nm}{l}"] for l in range(NL)]))
h.append(carr("SILU", "int32_t",
[round((x := (i - 256) / 32) / (1 + np.exp(-x)) * 4096) for i in range(512)]))
cs = []
for pos in range(MAXPOS):
for i in range(0, HS, 2):
fr = 1.0 / 10000 ** (i / HS)
cs += [round(np.cos(pos * fr) * 16384), round(np.sin(pos * fr) * 16384)]
h.append(carr("ROPE", "int16_t", cs))
h.append(carr("EXP2", "uint16_t", [round(2 ** (-i / 256) * 32768) for i in range(256)]))
import json
kv = json.load(open(HERE + "/kvmax.json"))
h.append(carr("KMAX", "int32_t", kv["kmax"]))
h.append(carr("VMAX", "int32_t", kv["vmax"]))
h.append(carr("P8", "uint8_t", [round(127 * np.exp(-i / 32)) for i in range(256)]))
open(HERE + "/weights_tl.h", "w").write("".join(h))
open(HERE + "/wrom_tl.c", "w").write(
"#include <stdint.h>\n__attribute__((section(\".wrom\"), used)) const uint8_t "
f"WROM_DATA[{len(P.cells)}] = {{{','.join(map(str, P.cells))}}};\n")
print(f"header y wrom escritos: {len(P.cells)} celdas, {P.unit + 1} unidades")
if __name__ == "__main__":
P, Q, S, W, SC = build()
emit(P, Q, S, W, SC)
|