File size: 9,001 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 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 | # SPDX-License-Identifier: Apache-2.0
"""Host-only study (round 8): can a channel PERMUTATION make bfp8 activations accurate enough?
bfp8_b shares one exponent per 16 consecutive channels of a token (one row of a 16x16 face), so a few
large-magnitude channels force a coarse step on their 15 neighbours. Permuting channels so that channels of
similar magnitude share a group is exact algebra for the model (the weights are permuted to match on the host),
but it changes how much the bfp8 rounding hurts.
Fake-quantizes (CPU fp32 reference model, forward pre-hooks):
ln1 = qkv input (the model already runs it in bfp8), ln2 = fc1 input, g = fc2 input (GELU output)
with identity or magnitude-sorted channel orders, and reports the 157-d output PCC vs the unquantized model over
the 32 acc.py crops (calibration of the orders on 8 DIFFERENT crops: the same hands with other bbox jitter).
usage: bfp8_perm_study.py [n_eval]"""
from __future__ import annotations
import json
import sys
from pathlib import Path
import torch
_REPO = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(_REPO))
sys.path.insert(0, str(_REPO / "scripts"))
from hamer import Hamer, load_hamer_checkpoint # noqa: E402
from hamer.preprocess import preprocess # noqa: E402
from acc import crops # noqa: E402
torch.set_num_threads(16)
import os
TRUNC = os.environ.get('BFP8_TRUNC') == '1' # round toward zero (packer-truncation hypothesis)
def bfp8(x: torch.Tensor, perm: torch.Tensor | None) -> torch.Tensor:
"""bfp8_b fake quantization along the last dim (groups of 16 after `perm`), round to nearest."""
if perm is not None:
x = x[..., perm]
s = x.shape
g = x.reshape(*s[:-1], s[-1] // 16, 16)
mx = g.abs().amax(-1, keepdim=True).clamp_min(1e-30)
e = torch.floor(torch.log2(mx))
step = torch.pow(2.0, e - 6)
r = torch.trunc(g / step) if TRUNC else torch.round(g / step)
q = (r.clamp(-127, 127) * step).reshape(s)
if perm is not None:
q = q[..., torch.argsort(perm)]
return q
def calib_crops():
media = _REPO.parent / "media"
boxes = json.loads((media / "bboxes.json").read_text())
out = []
for name in sorted(boxes):
for e in boxes[name]:
x1, y1, x2, y2 = e["bbox"]
w, h = x2 - x1, y2 - y1
bb = (x1 - 0.05 * w, y1 + 0.05 * h, x2 - 0.05 * w, y2 + 0.05 * h)
x, _ = preprocess(media / name, bbox_xyxy=bb, is_right=bool(e.get("is_right", 1)))
out.append(x)
return out
def main():
n_eval = int(sys.argv[1]) if len(sys.argv) > 1 else 32
hf = Path.home() / ".cache/huggingface/hub/models--changh95--hamer-weights/snapshots/e8217b46d54b4e8afaaf0dd6efb6b62341c97d32"
ref = Hamer().eval()
load_hamer_checkpoint(ref, hf / "hamer.ckpt", hf / "mano_mean_params.npz")
blocks = ref.backbone.blocks
L = len(blocks)
# ---- calibration: per-channel mean |x| of ln1 / ln2 / g per layer
stats = {k: [torch.zeros(1280 if k != "g" else 5120) for _ in range(L)] for k in ("ln1", "ln2", "g")}
hooks = []
for i, b in enumerate(blocks):
for k, mod in (("ln1", b.attn.qkv), ("ln2", b.mlp.fc1), ("g", b.mlp.fc2)):
def pre(m, inp, i=i, k=k):
stats[k][i] += inp[0].detach().abs().reshape(-1, inp[0].shape[-1]).mean(0)
hooks.append(mod.register_forward_pre_hook(pre))
with torch.no_grad():
for x in calib_crops():
ref(x)
for h in hooks:
h.remove()
# per-layer orders (hidden channels: free per layer); the stream orders must be ONE global order
order = {k: [torch.argsort(stats[k][i], descending=True) for i in range(L)] for k in stats}
# data-free order of the GELU outputs (round 8): E[relu(z)] for z ~ N(mu, sigma^2), mu = b1 + W1 beta2,
# sigma = |W1 (.) gamma2| (LN output ~ unit variance per channel): needs the weights only
def wf_order(b):
W = b.mlp.fc1.weight.detach().double() # [5120, 1280]
g2, b2 = b.norm2.weight.detach().double(), b.norm2.bias.detach().double()
mu = b.mlp.fc1.bias.detach().double() + W @ b2
sg = (W * g2).norm(dim=1).clamp_min(1e-12)
a = mu / sg
nd = torch.distributions.Normal(0.0, 1.0)
er = sg * torch.exp(-0.5 * a * a) / (2 * torch.pi) ** 0.5 + mu * nd.cdf(a)
return torch.argsort(er.float(), descending=True)
order["gw"] = [wf_order(b) for b in blocks]
# power-of-two per-channel scales of the LN outputs (fold 2^-k into gamma / beta, 2^k into the next weight's K
# rows: exact in bf16 / bfp8): data-free |h_c| ~ sqrt(gamma_c^2 + beta_c^2), or calibrated mean |h_c|
def p2(v):
return torch.pow(2.0, torch.round(torch.log2(v.clamp_min(1e-12))))
scale = {"ln1": [p2((b.norm1.weight.detach() ** 2 + b.norm1.bias.detach() ** 2).sqrt()) for b in blocks],
"ln2": [p2((b.norm2.weight.detach() ** 2 + b.norm2.bias.detach() ** 2).sqrt()) for b in blocks]}
scale_c = {k: [p2(stats[k][i] / stats[k][i].median()) for i in range(L)] for k in ("ln1", "ln2")}
# rank agreement with the calibrated order (Spearman)
def spear(o1, o2):
r1 = torch.empty_like(o1); r1[o1] = torch.arange(len(o1)); r2 = torch.empty_like(o2); r2[o2] = torch.arange(len(o2))
return float(torch.corrcoef(torch.stack([r1.double(), r2.double()]))[0, 1])
sp = [spear(order["g"][i], order["gw"][i]) for i in range(L)]
print(f"[calib] data-free vs calibrated g order, Spearman per layer: min {min(sp):.3f} median {sorted(sp)[L // 2]:.3f}")
# ONE stream order serves ln1 and ln2 (both are in the residual stream's channel order)
tot = sum(s / s.mean() for s in stats["ln1"]) + sum(s / s.mean() for s in stats["ln2"])
g1 = torch.argsort(tot, descending=True)
glob = {"ln1": g1, "ln2": g1}
# ratio of the largest to the median channel magnitude (outlier severity)
for k in stats:
r = [float(stats[k][i].max() / stats[k][i].median()) for i in range(L)]
print(f"[calib] {k}: max/median channel |x| per layer: min {min(r):.1f} median {sorted(r)[L // 2]:.1f} max {max(r):.1f}")
xs = crops()[:n_eval]
with torch.no_grad():
base = [ref(x) for _, x in xs]
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 run(cfg: dict, tag: str):
hs = []
for i, b in enumerate(blocks):
for k, mod in (("ln1", b.attn.qkv), ("ln2", b.mlp.fc1), ("g", b.mlp.fc2)):
mode = cfg.get(k)
if mode is None:
continue
if mode in ("scale", "scalec"):
sc = (scale if mode == "scale" else scale_c)[k][i]
hs.append(mod.register_forward_pre_hook(lambda m, inp, sc=sc: (bfp8(inp[0] / sc, None) * sc,)))
continue
perm = None if mode == "id" else (glob[k] if mode == "glob" else (order["gw"][i] if mode == "wf" else order[k][i]))
hs.append(mod.register_forward_pre_hook(lambda m, inp, perm=perm: (bfp8(inp[0], perm),)))
with torch.no_grad():
outs = [ref(x) for _, x in xs]
for h in hs:
h.remove()
p = [pcc(o, r) for o, r in zip(outs, base)]
cam = [float(o[0, -3] - r[0, -3]) for o, r in zip(outs, base)]
betas = max(float((o[0, -13:-3] - r[0, -13:-3]).abs().max()) for o, r in zip(outs, base))
print(f"[{tag:34s}] PCC mean {sum(p) / len(p):.7f} min {min(p):.7f} cam[0] signed {sum(cam) / len(cam):+.5f}"
f" abs {sum(abs(c) for c in cam) / len(cam):.5f} betas_max {betas:.4f}", flush=True)
if "trunc" in sys.argv:
run({"ln1": "id"}, "ln1 bfp8")
run({"ln1": "id", "ln2": "id"}, "ln1+ln2 bfp8")
return
if "scale" in sys.argv:
run({"ln1": "id"}, "ln1 bfp8 (served today)")
run({"ln1": "scale"}, "ln1 bfp8, p2 scale (gamma/beta)")
run({"ln1": "scalec"}, "ln1 bfp8, p2 scale (calibrated)")
run({"ln1": "id", "ln2": "id"}, "ln1+ln2 bfp8")
run({"ln1": "scale", "ln2": "scale"}, "ln1+ln2 bfp8, p2 scale (gamma/beta)")
run({"ln1": "scalec", "ln2": "scalec"}, "ln1+ln2 bfp8, p2 scale (calibrated)")
return
if "gonly" in sys.argv:
run({"ln1": "id", "g": "layer"}, "ln1 + g bfp8, per-layer sorted g")
run({"ln1": "id", "g": "wf"}, "ln1 + g bfp8, data-free sorted g")
return
run({"ln1": "id"}, "ln1 bfp8 (served today)")
run({"ln1": "glob"}, "ln1 bfp8, global sorted order")
run({"ln1": "id", "ln2": "id"}, "ln1+ln2 bfp8")
run({"ln1": "glob", "ln2": "glob"}, "ln1+ln2 bfp8, global order (ln2 stats)")
run({"ln1": "id", "g": "id"}, "ln1 + g bfp8")
run({"ln1": "id", "g": "layer"}, "ln1 + g bfp8, per-layer sorted g")
run({"ln1": "glob", "ln2": "glob", "g": "layer"}, "all bfp8, sorted")
if __name__ == "__main__":
main()
|