| import sys |
|
|
| sys.path.insert(0, "/tmp/opencode/tl") |
| import numpy as np |
| import torch |
|
|
| import export_full as E |
| import int4hf as H |
| from transformers import AutoModelForCausalLM |
|
|
| GREEDY_FLOAT = [4612, 2619, 260, 647, 31844, 635, 393, 260, 1496, 2789, 3557, 21075, 31843, 1100, |
| 4838, 289, 826, 351, 600, 17360, 1086, 31843, 1591, 1124, 31844, 700, 2815, 260, |
| 1421, 3476, 288, 600, 2103, 31843, 1100, 393, 900, 6055, 291, 2598] |
|
|
|
|
| def deq(w): |
| q, s = E.quantize(w) |
| sr = np.round(s * 2 ** 15) / 2 ** 15 |
| return torch.tensor(q.astype(np.float32) * sr[:, None].astype(np.float32)) |
|
|
|
|
| base = AutoModelForCausalLM.from_pretrained("/tmp/opencode/tl", dtype=torch.float32).eval() |
| t = E.load() |
|
|
|
|
| def build(groups): |
| m = H.int4_model() |
| with torch.no_grad(): |
| if "head" in groups: |
| m.lm_head.weight.copy_(torch.tensor(t["lm_head.weight"] * t["model.norm.weight"][None, :])) |
| if "embed" in groups: |
| m.model.embed_tokens.weight.copy_(torch.tensor(t["model.embed_tokens.weight"])) |
| for l in range(8): |
| L = m.model.layers[l] |
| ga = t[f"model.layers.{l}.input_layernorm.weight"] |
| gf = t[f"model.layers.{l}.post_attention_layernorm.weight"] |
| pre = f"model.layers.{l}." |
| if "qkv" in groups: |
| L.self_attn.q_proj.weight.copy_(torch.tensor(t[pre + "self_attn.q_proj.weight"] * ga[None, :])) |
| L.self_attn.k_proj.weight.copy_(torch.tensor(t[pre + "self_attn.k_proj.weight"] * ga[None, :])) |
| L.self_attn.v_proj.weight.copy_(torch.tensor(t[pre + "self_attn.v_proj.weight"] * ga[None, :])) |
| if "o" in groups: |
| L.self_attn.o_proj.weight.copy_(torch.tensor(t[pre + "self_attn.o_proj.weight"])) |
| if "gateup" in groups: |
| L.mlp.gate_proj.weight.copy_(torch.tensor(t[pre + "mlp.gate_proj.weight"] * gf[None, :])) |
| L.mlp.up_proj.weight.copy_(torch.tensor(t[pre + "mlp.up_proj.weight"] * gf[None, :])) |
| if "down" in groups: |
| L.mlp.down_proj.weight.copy_(torch.tensor(t[pre + "mlp.down_proj.weight"])) |
| return m |
|
|
|
|
| def greedy(m, n=40): |
| ids = torch.tensor([[1]]) |
| out = [] |
| with torch.no_grad(): |
| for _ in range(n): |
| lg = m(ids).logits[0, -1] |
| k = int(lg.argmax()) |
| out.append(k) |
| ids = torch.cat([ids, torch.tensor([[k]])], 1) |
| return out |
|
|
|
|
| def hit(out): |
| n = 0 |
| for a, b in zip(out, GREEDY_FLOAT): |
| if a != b: |
| break |
| n += 1 |
| return n |
|
|
|
|
| print("base int4: ", hit(greedy(build(set()))), "/40", flush=True) |
| for g in (["head"], ["embed"], ["qkv"], ["o"], ["gateup"], ["down"]): |
| print(f"solo {g[0]:8s} en float:", hit(greedy(build(set(g)))), "/40", flush=True) |
|
|
