# SPDX-License-Identifier: Apache-2.0 """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): # q [8,64], k/v [8,192,64] 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}")