| |
| """MANO-head error attribution (32 acc.py crops). The device ViT feature map is the common input; the device |
| head runs with hamer.mano_head._DBG recording every op's input/output. Each op's LOCAL error (device output minus |
| the fp32 op applied to the device input) is injected alone into an fp32 CPU head; the resulting change of the |
| camera scale cam[0] (and its sign count) attributes the end-to-end camera-scale bias to op types.""" |
| import os, sys, statistics |
| from collections import defaultdict |
| from pathlib import Path |
| sys.path.insert(0, str(Path(__file__).resolve().parents[1])) |
| sys.path.insert(0, str(Path(__file__).resolve().parent)) |
| os.environ.setdefault("TT_FUSED_HASH", "0") |
| import torch, torch.nn.functional as F, ttnn |
| from hamer import Hamer, build_from_reference, load_hamer_checkpoint |
| from hamer import vit as tv, mano_head as mh |
| from hamer.fused import head_layer0_constants |
| from acc import crops |
|
|
| hf = Path.home() / ".cache/huggingface/hub/models--changh95--hamer-weights/snapshots/e8217b46d54b4e8afaaf0dd6efb6b62341c97d32" |
| torch.manual_seed(0) |
| ref = Hamer().eval() |
| load_hamer_checkpoint(ref, hf / "hamer.ckpt", hf / "mano_mean_params.npz") |
| head = ref.head |
| td = head.transformer |
| D = 1024 |
| xs = crops() |
| if len(sys.argv) > 1: |
| xs = xs[: int(sys.argv[1])] |
| tt = build_from_reference(ref, device_id=0, strict=True) |
| TYPES = ["kv", "x1", "q0", "ln_sa", "sa", "add_sa", "ln_ca", "q", "sdpa", "ca_out", "add_ca", "ln_ff", "fc1", "gelu", |
| "fc2", "add_ff", "dec"] |
|
|
|
|
| def tok(t, n=D): |
| return t.reshape(-1, t.shape[-1])[0, :n].double() |
|
|
|
|
| def attn(q, k, v): |
| a = torch.softmax(torch.einsum("hd,hnd->hn", q, k) * 64 ** -0.5, -1) |
| return torch.einsum("hn,hnd->hd", a, v) |
|
|
|
|
| def lin(m, x): |
| return F.linear(x, m.weight.double(), None if m.bias is None else m.bias.double()) |
|
|
|
|
| def ln(m, x): |
| return F.layer_norm(x, (D,), m.weight.double(), m.bias.double(), m.eps) |
|
|
|
|
| def cpu_head(ctx, E, inj): |
| """fp32(64) head on ctx [192,1280] with local errors E[name] added where name's type is in inj; returns cam[0] |
| contribution of dec (the cam scale before init_cam).""" |
| e = lambda n: E[n] if n.split("@")[0] in inj else 0.0 |
| x1, q0 = head_layer0_constants(td) |
| x = x1.double().reshape(D) + e("x1@") |
| res = [] |
| for i, blk in enumerate(td.layers): |
| kv = lin(blk.ca.to_kv, ctx) |
| k, v = kv.chunk(2, -1) |
| k = k.reshape(192, 8, 64).permute(1, 0, 2) + e(f"kv@{i}k") |
| v = v.reshape(192, 8, 64).permute(1, 0, 2) + e(f"kv@{i}v") |
| if i > 0: |
| h = ln(blk.norm_sa, x) + e(f"ln_sa@{i}") |
| sa = lin(blk.sa.to_out, lin(blk.sa.to_qkv, h).chunk(3, -1)[2]) + e(f"sa@{i}") |
| x = x + sa + e(f"add_sa@{i}") |
| h = ln(blk.norm_ca, x) + e(f"ln_ca@{i}") |
| q = (lin(blk.ca.to_q, h) + e(f"q@{i}")).reshape(8, 64) |
| else: |
| q = q0.double().reshape(8, 64) + e("q0@").reshape(8, 64) if "q0" in inj else q0.double().reshape(8, 64) |
| c = attn(q, k, v) + e(f"sdpa@{i}") |
| co = lin(blk.ca.to_out, c.reshape(512)) + e(f"ca_out@{i}") |
| x = x + co + e(f"add_ca@{i}") |
| h = ln(blk.norm_ff, x) + e(f"ln_ff@{i}") |
| f = lin(blk.ff.net[0], h) + e(f"fc1@{i}") |
| g = F.gelu(f) + e(f"gelu@{i}") |
| f2 = lin(blk.ff.net[2], g) + e(f"fc2@{i}") |
| x = x + f2 + e(f"add_ff@{i}") |
| cam = lin(head.deccam, x) + e("dec@") |
| return float(cam[0]) |
|
|
|
|
| acc = defaultdict(list) |
| with torch.no_grad(): |
| for n, img in xs: |
| tokd = tt._upload_tokens(img) |
| fdev = tv.forward(tokd, tt._vit_params, fused=True) |
| ctx = ttnn.to_torch(fdev).double().reshape(192, 1280) |
| mh._DBG = [] |
| dec = mh.forward_ds(fdev, tt._head_params) |
| R = {} |
| for name, ts in mh._DBG: |
| R.setdefault(name, ts) |
| mh._DBG = None |
| ttnn.deallocate(fdev) |
| E = {} |
| x1, q0 = head_layer0_constants(td) |
| E["x1@"] = tok(R["x_in0"][0]) - x1.double().reshape(D) |
| E["q0@"] = R["sdpa0"][0][0, :, 0, :].double() - q0.double().reshape(8, 64) |
| for i, blk in enumerate(td.layers): |
| kd, vd = R["kvs"][2 * i][0].double(), R["kvs"][2 * i + 1][0].double() |
| k, v = lin(blk.ca.to_kv, ctx).chunk(2, -1) |
| E[f"kv@{i}k"] = kd - k.reshape(192, 8, 64).permute(1, 0, 2) |
| E[f"kv@{i}v"] = vd - v.reshape(192, 8, 64).permute(1, 0, 2) |
| xin = tok(R[f"x_in{i}"][0]) |
| if i > 0: |
| h = tok(R[f"ln_sa{i}"][0]); E[f"ln_sa@{i}"] = h - ln(blk.norm_sa, xin) |
| sa = tok(R[f"sa{i}"][0]); E[f"sa@{i}"] = sa - lin(blk.sa.to_out, lin(blk.sa.to_qkv, h).chunk(3, -1)[2]) |
| xs_ = tok(R[f"sa{i}"][1]); E[f"add_sa@{i}"] = xs_ - (xin + sa) |
| h = tok(R[f"ln_ca{i}"][0]); E[f"ln_ca@{i}"] = h - ln(blk.norm_ca, xs_) |
| q = tok(R[f"ln_ca{i}"][1], 512); E[f"q@{i}"] = q - lin(blk.ca.to_q, h) |
| xb = xs_ |
| else: |
| xb = xin |
| qh = R[f"sdpa{i}"][0][0, :, 0, :].double() |
| c = R[f"sdpa{i}"][1][0, :, 0, :].double(); E[f"sdpa@{i}"] = c - attn(qh, kd, vd) |
| cm = tok(R[f"ca_out{i}"][0], 512); co = tok(R[f"ca_out{i}"][1]) |
| E[f"ca_out@{i}"] = co - lin(blk.ca.to_out, cm) |
| xa = tok(R[f"add_ca{i}"][0]); E[f"add_ca@{i}"] = xa - (xb + co) |
| h = tok(R[f"ff{i}"][0]); E[f"ln_ff@{i}"] = h - ln(blk.norm_ff, xa) |
| f = tok(R[f"ff{i}"][1], 4 * D if False else None) if False else R[f"ff{i}"][1].reshape(-1, R[f"ff{i}"][1].shape[-1])[0].double() |
| f = f[: blk.ff.net[0].out_features] |
| E[f"fc1@{i}"] = f - lin(blk.ff.net[0], h) |
| g = R[f"fc2_{i}"][0].reshape(-1, R[f"fc2_{i}"][0].shape[-1])[0].double()[: blk.ff.net[0].out_features] |
| E[f"gelu@{i}"] = g - F.gelu(f) |
| f2 = tok(R[f"fc2_{i}"][1]); E[f"fc2@{i}"] = f2 - lin(blk.ff.net[2], g) |
| xo = tok(R[f"add_ff{i}"][0]); E[f"add_ff@{i}"] = xo - (xa + f2) |
| xl = tok(R["dec"][0]); dd = tok(R["dec"][1], 109) |
| E["dec@"] = (dd - lin(torch.nn.Linear(1, 1), torch.zeros(1)) * 0 if False else dd[106:109] - lin(head.deccam, xl)) |
| base = cpu_head(ctx, E, set()) |
| acc["all"].append(cpu_head(ctx, E, set(TYPES)) - base) |
| for t in TYPES: |
| acc[t].append(cpu_head(ctx, E, {t}) - base) |
| acc["device_dec_vs_cpu"].append(float(dd[108 - 2]) - base) |
| print(n, f"all {acc['all'][-1]:+.4f} dev {acc['device_dec_vs_cpu'][-1]:+.4f}", flush=True) |
| tt.close() |
| for k, v in acc.items(): |
| print(f"{k:18s} mean {statistics.mean(v):+.5f} pos {sum(t > 0 for t in v)}/{len(v)} mean|.| {statistics.mean(abs(t) for t in v):.5f}") |
|
|
