| |
| """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}") |
|
|
