| |
| """Host-only: end-to-end sensitivity of the real-checkpoint output (32 acc.py crops, fp32 CPU) to the |
| ViT MLP GELU implementation. PCC vs the exact-erf reference for several approximations.""" |
| import sys, math |
| 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 |
| from acc import crops |
| from hamer import Hamer, load_hamer_checkpoint |
|
|
| 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") |
| cs = crops() |
| stats = {"n": 0, "lt-3": 0, "gt3": 0, "lt-4": 0, "lt-5.5": 0, "absmax": 0.0} |
|
|
| class Act(torch.nn.Module): |
| def __init__(self, f): super().__init__(); self.f = f |
| def forward(self, x): |
| if self.f is None: |
| stats["n"] += x.numel(); stats["lt-3"] += int((x < -3.1).sum()); stats["gt3"] += int((x > 3.1).sum()) |
| stats["lt-4"] += int((x < -4).sum()); stats["lt-5.5"] += int((x < -5.5).sum()); stats["absmax"] = max(stats["absmax"], float(x.abs().max())) |
| return torch.nn.functional.gelu(x) |
| return self.f(x) |
|
|
| def set_act(f): |
| for b in ref.backbone.blocks: |
| b.mlp.act = Act(f) |
|
|
| def run(): |
| with torch.no_grad(): |
| return [ref(x).reshape(-1).double() for _, x in cs] |
|
|
| def pcc(u, v): |
| u = u - u.mean(); v = v - v.mean(); return float((u @ v) / (u.norm() * v.norm())) |
|
|
| set_act(None); base = run(); print("fc1-out stats", {k: (v / stats["n"] if k not in ("n", "absmax") else v) for k, v in stats.items()}) |
| erf_gelu = lambda x: torch.nn.functional.gelu(x) |
| def clamp_tails(lo, hi): |
| return lambda x: torch.where(x > hi, x, torch.where(x < lo, torch.zeros_like(x), torch.nn.functional.gelu(x))) |
| _C8 = [8.040229593e-11, -6.988785106e-09, 2.712364733e-07, -6.289864814e-06, 9.886907355e-05, -1.133757411e-03, |
| 9.877118282e-03, -6.641025841e-02, 3.989226222e-01] |
| def poly8(x): |
| """hamer/kernels/gelu SFPU recipe in fp32: x * (0.5 + x * P(x^2)), x > 4 -> x, x < -4 -> 0.""" |
| x = x.float(); u = x * x; acc = torch.zeros_like(x) |
| for c in _C8: |
| acc = acc * u + c |
| g = x * (0.5 + x * acc) |
| return torch.where(x > 4, x, torch.where(x < -4, torch.zeros_like(x), g)) |
| variants = { |
| "poly8": poly8, |
| "poly8-bf16": lambda x: poly8(x).bfloat16().float(), |
| "tanh": lambda x: torch.nn.functional.gelu(x, approximate="tanh"), |
| "bf16-round": lambda x: torch.nn.functional.gelu(x).bfloat16().float(), |
| "tails 3.1": clamp_tails(-3.1, 3.1), |
| "tails 4": clamp_tails(-4.0, 4.0), |
| "tails -5.5/4": clamp_tails(-5.5, 4.0), |
| } |
| extra = sys.argv[1:] |
| for name, f in variants.items(): |
| if extra and name not in extra: |
| continue |
| set_act(f); out = run() |
| p = [pcc(a, b) for a, b in zip(out, base)] |
| print(f"{name:14s} mean {sum(p)/len(p):.8f} min {min(p):.8f}", flush=True) |
|
|
