File size: 3,351 Bytes
ee52468 | 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 | """How well can the experts of layer j be predicted before segment j runs? Take segment j's input hidden state (the
residual stream one layer earlier), apply layer j's own router (weights read from the IR) and compare its top-k
with the experts the real router picks at the end of segment j, while generating real answers.
usage: python predict_experts.py MODEL_DIR"""
import sys
import numpy as np
import openvino as ov
from onw.chat import ChatEngine
from onw.runtime import SegmentedModel, read_segment
def routers(d, meta):
"""{layer: (W [E,H] f32, gamma [H] f32)} from the 1-token segments' IR."""
core, out = ov.Core(), {}
for name, m in sorted(meta["segments"].items()):
if m["S"] != 1 or m.get("router_layer") is None:
continue
model = read_segment(core, d, name, m, False, {}, meta["T"], meta["T"])
tk = next(n for n in model.get_ops() if n.get_type_name() == "TopK")
mm = tk.input(0).get_source_output().get_node().input(0).get_source_output().get_node() \
.input(0).get_source_output().get_node() # softmax <- reshape <- matmul
assert mm.get_type_name() == "MatMul"
w = mm.input(1).get_source_output().get_node()
while w.get_type_name() != "Constant":
w = w.input(0).get_source_output().get_node()
g = mm.input(0).get_source_output().get_node() # xn * gamma
gam = next(i.get_source_output().get_node() for i in g.inputs()
if i.get_source_output().get_node().get_type_name() == "Constant")
out[m["router_layer"]] = (np.asarray(w.get_data(), np.float32), np.asarray(gam.get_data(), np.float32))
return out
def main():
d = sys.argv[1]
e = ChatEngine(d, "NPU", pld=False)
R = routers(d, e.model.meta)
print(f"routers read for {len(R)} layers ({sum(w.nbytes / 2 + g.nbytes for w, g in R.values()) / 2**20:.0f} MB as f16)")
K = e.model.K
stats = {k: [0, 0] for k in (K, K + 4, K + 8, 2 * K)}
orig = SegmentedModel._run_segment
def spy(self, req, m, names, carry, common, route, S, n, cur, keep=False):
L = m.get("router_layer")
pred = None
if S == 1 and L is not None and L in R and "carry.h" in carry:
h = carry["carry.h"].astype(np.float32).reshape(-1)
xn = h / np.sqrt((h * h).mean() + 1e-6) * R[L][1]
pred = np.argsort(-(R[L][0] @ xn))
out = orig(self, req, m, names, carry, common, route, S, n, cur, keep)
if pred is not None:
real = set(int(x) for x in out[1][0][0])
for k in stats:
stats[k][0] += len(real & set(int(x) for x in pred[:k]))
stats[k][1] += len(real)
return out
SegmentedModel._run_segment = spy
for q in ["NPUとGPUの違いを、身近なたとえを使って説明してください。",
"次のPython関数に型ヒントを付けてください。\n\ndef add(a, b):\n return a + b",
"Explain the difference between a process and a thread."]:
e.checkpoint = None
list(e.stream_chat([{"role": "user", "content": q}], 96))
for k, (hit, allp) in stats.items():
print(f"predict top-{k:2d}: {hit / allp * 100:5.1f}% of the real top-{K} experts caught")
if __name__ == "__main__":
main()
|