| """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) |
|
|
