File size: 2,802 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 | 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)
|