File size: 2,540 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 | # SPDX-License-Identifier: Apache-2.0
"""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:
# torch Linear weight is [N, K]; device layout is [K, N] -> quantize in that orientation
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()
|