File size: 6,706 Bytes
58ba6df | 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 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 | # 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}")
|