| |
| """Host-only (no device) sensitivity check of the ViT weight-precision change. |
| |
| Quantizes ViT weights to bfp8 on the HOST (ttnn.from_torch(..., bfloat8_b, TILE) -> to_torch, no |
| device opened) and runs the fp32 torch CPU model on the same 32 crops as acc.py, PCC vs the |
| unquantized fp32 reference. It isolates the effect of the declared weight-precision change: |
| baseline config : fc1 weights bfp8 |
| optimized config: qkv + proj + fc1 + fc2 weights bfp8 |
| It does NOT measure device numerics (bf16 activations, LN/SDPA/gelu kernels, fp32 vs bf16 DST); |
| those need acc.py on a healthy chip. Run with TT_VISIBLE_DEVICES=none. |
| """ |
| import copy, json, sys |
| from pathlib import Path |
|
|
| _REPO = Path(__file__).resolve().parents[1] |
| sys.path.insert(0, str(_REPO)); sys.path.insert(0, str(_REPO / "scripts")) |
| import torch |
| import ttnn |
| from acc import crops |
| from hamer import Hamer, load_hamer_checkpoint |
|
|
|
|
| def q8(w): |
| return ttnn.to_torch(ttnn.from_torch(w.detach().float(), dtype=ttnn.bfloat8_b, layout=ttnn.TILE_LAYOUT)).float() |
|
|
|
|
| def pcc(u, v): |
| u = u.reshape(-1).double(); v = v.reshape(-1).double() |
| u = u - u.mean(); v = v - v.mean() |
| return float((u @ v) / (u.norm() * v.norm()).clamp_min(1e-12)) |
|
|
|
|
| def quantized(ref, names): |
| m = copy.deepcopy(ref) |
| for blk in m.backbone.blocks if hasattr(m, "backbone") else m.vit.blocks: |
| mods = {"qkv": blk.attn.qkv, "proj": blk.attn.proj, "fc1": blk.mlp.fc1, "fc2": blk.mlp.fc2} |
| for n in names: |
| |
| mods[n].weight.data = q8(mods[n].weight.data.t().contiguous()).t().contiguous() |
| return m |
|
|
|
|
| def main(): |
| 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") |
| xs = crops() |
| cfgs = {"baseline(fc1)": ["fc1"], "optimized(qkv,proj,fc1,fc2)": ["qkv", "proj", "fc1", "fc2"]} |
| with torch.no_grad(): |
| refs = [ref(x) for _, x in xs] |
| for tag, names in cfgs.items(): |
| m = quantized(ref, names) |
| p = [pcc(m(x), r) for (_, x), r in zip(xs, refs)] |
| print("HOSTQ", json.dumps({"cfg": tag, "n": len(p), "mean": sum(p) / len(p), "min": min(p), |
| "below_0.9999": sum(v < 0.9999 for v in p)}), flush=True) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
