File size: 1,706 Bytes
67876a1 | 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 | """NPU vs CPU, segment by segment, for one block of S tokens (same inputs fed to both at every segment).
usage: python probe_segments.py MODEL_DIR [S=1]"""
import sys
import numpy as np
import openvino as ov
from onw.runtime import SegmentedModel
def main():
d = sys.argv[1]
S = int(sys.argv[2]) if len(sys.argv) > 2 else 1
core = ov.Core()
cpu = SegmentedModel(core, d, "CPU", {"INFERENCE_PRECISION_HINT": "f32"}, sizes=[1, S] if S > 1 else [1])
npu = SegmentedModel(core, d, "NPU", {"CACHE_DIR": d + "/npu_cache"}, sizes=[1, S] if S > 1 else [1])
ids = list(range(1000, 1000 + S))
# run CPU normally, capturing every segment's feed; replay each feed on the NPU segment
feeds = []
for req, m, names in cpu.segs[S]:
orig = req.infer
def cap(feed, *a, _o=orig, **k):
feeds.append(dict(feed))
return _o(feed, *a, **k)
req.infer = cap
cpu.step(ids)
cpu.reset() # replay from the same (zero) state on both devices
npu.reset()
for i, ((creq, m, names), (nreq, _, _)) in enumerate(zip(cpu.segs[S], npu.segs[S])):
f = feeds[i]
r_cpu = {o.get_any_name(): v.astype(np.float32) for o, v in creq.infer(f).items()}
r_npu = {o.get_any_name(): v.astype(np.float32) for o, v in nreq.infer(f).items()}
errs = {k: float(np.linalg.norm(r_npu[k] - r_cpu[k]) / (np.linalg.norm(r_cpu[k]) + 1e-12)) for k in r_cpu
if r_cpu[k].dtype.kind == "f"}
worst = sorted(errs.items(), key=lambda kv: -kv[1])[:12]
print(f"seg {i}: " + " ".join(f"{k} {v:.3f}" for k, v in worst), flush=True)
if __name__ == "__main__":
main()
|