| """numpy fp32 forward of SlothE-T 12M mirroring upstream slothe.cpp (ggml) math: F16 islands, |
| ternary linears = SubLN -> (A8: per-256-block absmax int8, ggml Q8_K) -> {-1,0,1}*scale. |
| Used to check the Core ML graph structure and activation ranges; ggml itself is the reference.""" |
| import json |
| import numpy as np |
|
|
| W = np.load("work/w12m.npz"); C = json.load(open("work/cfg12m.json")) |
| EPS = 1e-6 |
| f16 = lambda a: a.astype(np.float16).astype(np.float32) |
|
|
| def rms(x, w): |
| return x / np.sqrt((x * x).mean(-1, keepdims=True) + EPS) * f16(w) |
|
|
| def a8(x): |
| L, n = x.shape |
| xb = x.reshape(L, n // 256, 256) |
| am = np.abs(xb).max(-1, keepdims=True) |
| am = np.where(am == 0, 1, am) |
| return (np.round(xb * (127.0 / am)) * (am / 127.0)).reshape(L, n) |
|
|
| STATS = {} |
| TRACE = [] |
| def lin(i, s, x, fp, act8=True): |
| if fp: |
| return f16(x) @ f16(W[f"{i}.{s}.w"]).T |
| x = rms(x, W[f"{i}.{s}.pre"]) |
| STATS["lin_in"] = max(STATS.get("lin_in", 0), float(np.abs(x).max())) |
| if act8: |
| x = a8(x) |
| return (x @ W[f"{i}.{s}.q"].astype(np.float32).T) * f16(W[f"{i}.{s}.s"]) |
|
|
| def rope_tab(L, hd=32): |
| half = hd // 2 |
| freq = 1.0 / (10000 ** (np.arange(half) / half)) |
| ang = np.arange(L)[:, None] * freq[None] |
| return np.concatenate([np.cos(ang)] * 2, -1), np.concatenate([np.sin(ang)] * 2, -1) |
|
|
| def forward(ids, act8=True): |
| L = len(ids); D = C["depth"]; H, KV, hd = C["heads"], C["kv"], C["dim"] // C["heads"] |
| cos, sin = rope_tab(L, hd); TRACE.clear() |
| x = rms(f16(W["embed"])[ids], W["embed_norm"]) |
| for i in range(D): |
| fp = i == 0 or i == D - 1 |
| h = rms(x, W[f"{i}.n1"]) |
| q = lin(i, "q", h, fp, act8).reshape(L, H, hd).transpose(1, 0, 2) |
| k = lin(i, "k", h, fp, act8).reshape(L, KV, hd).transpose(1, 0, 2) |
| v = lin(i, "v", h, fp, act8).reshape(L, KV, hd).transpose(1, 0, 2) |
| q, k = rms(q, W[f"{i}.qn"]), rms(k, W[f"{i}.kn"]) |
| rot = lambda t: np.concatenate([-t[..., hd // 2:], t[..., :hd // 2]], -1) |
| q, k = q * cos + rot(q) * sin, k * cos + rot(k) * sin |
| k, v = np.repeat(k, H // KV, 0), np.repeat(v, H // KV, 0) |
| s = q @ k.transpose(0, 2, 1) / np.sqrt(hd) |
| s = np.exp(s - s.max(-1, keepdims=True)); s /= s.sum(-1, keepdims=True) |
| o = (s @ v).transpose(1, 0, 2).reshape(L, H * hd) |
| x = x + lin(i, "o", o, fp, act8) |
| h2 = rms(x, W[f"{i}.n2"]) |
| g, u = lin(i, "w1", h2, fp, act8), lin(i, "w3", h2, fp, act8) |
| ff = g / (1 + np.exp(-g)) * u |
| STATS["ffn_act"] = max(STATS.get("ffn_act", 0), float(np.abs(ff).max())) |
| x = x + lin(i, "w2", ff, fp, act8) |
| STATS["resid"] = max(STATS.get("resid", 0), float(np.abs(x).max())) |
| TRACE.append(x.copy()) |
| x = rms(x, W["norm"]) |
| return f16(x) @ f16(W["head"]).T |
|
|
