File size: 4,932 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 | # SPDX-License-Identifier: Apache-2.0
"""Camera-scale error attribution of the TT_FUSED_HEAD=hg head (round 6): the device ViT feature map is the common
input; every fused head op's LOCAL error (device output minus the fp64 op on the device input) is injected alone into
an fp64 CPU head; the change of cam[0] attributes the end-to-end camera-scale bias to op types. usage: attrib_hg.py [n]"""
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")
os.environ["TT_FUSED_HEAD"] = "hg"
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, td, D = ref.head, ref.head.transformer, 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", "sa", "q", "sdpa", "co", "fc1", "fc2", "dec"]
r0 = lambda t, n=None: t.reshape(-1, t.shape[-1])[0, : (n or t.shape[-1])].double()
lin = lambda m, x: F.linear(x, m.weight.double(), None if m.bias is None else m.bias.double())
ln = lambda m, x: F.layer_norm(x, (D,), m.weight.double(), m.bias.double(), m.eps)
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 sa_f(blk, x):
return lin(blk.sa.to_out, lin(blk.sa.to_qkv, ln(blk.norm_sa, x)).chunk(3, -1)[2])
def cpu(ctx, E, inj):
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)
for i, blk in enumerate(td.layers):
k, v = lin(blk.ca.to_kv, ctx).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:
x = x + sa_f(blk, x) + e(f"sa@{i}")
q = (lin(blk.ca.to_q, ln(blk.norm_ca, x)) + e(f"q@{i}")).reshape(8, 64)
else:
q = q0.double().reshape(8, 64)
c = attn(q, k, v) + e(f"sdpa@{i}")
x = x + lin(blk.ca.to_out, c.reshape(512)) + e(f"co@{i}")
f = F.gelu(lin(blk.ff.net[0], ln(blk.norm_ff, x))) + e(f"fc1@{i}")
x = x + lin(blk.ff.net[2], f) + e(f"fc2@{i}")
return float(lin(head.deccam, x)[0] + e("dec@")), x
acc = defaultdict(list)
with torch.no_grad():
for n, img in xs:
tok = tt._upload_tokens(img)
fdev = tv.forward(tok, tt._vit_params, fused=True)
ctx = ttnn.to_torch(fdev).double().reshape(192, 1280)
mh._DBG = []
dec = mh.forward_hg(fdev, tt._head_params)
R = dict(mh._DBG)
mh._DBG = None
ttnn.deallocate(fdev); ttnn.deallocate(dec)
E = {}
for i, blk in enumerate(td.layers):
k, v = lin(blk.ca.to_kv, ctx).chunk(2, -1)
kd, vd = R["kvs"][2 * i][0].double(), R["kvs"][2 * i + 1][0].double()
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)
if i > 0:
xi, xo = r0(R[f"sa{i}"][0]), r0(R[f"sa{i}"][1])
E[f"sa@{i}"] = xo - (xi + sa_f(blk, xi))
xq, qd = r0(R[f"q{i}"][0]), r0(R[f"q{i}"][1], 512)
E[f"q@{i}"] = qd - lin(blk.ca.to_q, ln(blk.norm_ca, xq))
qh, ca = R[f"sdpa{i}"][0][0, :, 0, :].double(), R[f"sdpa{i}"][1][0, :, 0, :].double()
E[f"sdpa@{i}"] = ca - attn(qh, kd, vd)
cc, xb, xo = r0(R[f"co{i}"][0], 512), r0(R[f"co{i}"][1]), r0(R[f"co{i}"][2])
E[f"co@{i}"] = xo - (xb + lin(blk.ca.to_out, cc))
xa, fd, xo = r0(R[f"ff{i}"][0]), r0(R[f"ff{i}"][1]), r0(R[f"ff{i}"][2])
E[f"fc1@{i}"] = fd - F.gelu(lin(blk.ff.net[0], ln(blk.norm_ff, xa)))
E[f"fc2@{i}"] = xo - (xa + lin(blk.ff.net[2], fd))
xl, dd = r0(R["dec"][0]), r0(R["dec"][1], 109)
E["dec@"] = float(dd[106] - lin(head.deccam, xl)[0])
base, _ = cpu(ctx, E, set())
acc["all"].append(cpu(ctx, E, set(TYPES))[0] - base)
for t in TYPES:
acc[t].append(cpu(ctx, E, {t})[0] - base)
acc["device"].append(float(dd[106]) - base)
print(n, f"all {acc['all'][-1]:+.4f} dev {acc['device'][-1]:+.4f}", flush=True)
tt.close()
for k, v in acc.items():
print(f"{k:8s} 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}")
|