| """Load a paper checkpoint for evaluation; for arm C, round the ternary weights ONCE so inference runs at float speed. |
| (train.TLinear re-rounds weights on every forward pass, which is right for training and 2x slower for evaluation.)""" |
| import numpy as np, torch, torch.nn.functional as F |
| import train |
|
|
|
|
| class FrozenTernary(torch.nn.Module): |
| def __init__(self, lin): |
| super().__init__() |
| w = lin.weight.detach() |
| sw = 1.0 / w.abs().mean().clamp_min(1e-5) |
| self.register_buffer("wq", ((w * sw).round().clamp(-1, 1) / sw)) |
| self.bias = lin.bias |
|
|
| def forward(self, x): |
| xf = x.float() |
| xn = xf * torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + 1e-6) |
| sx = 127.0 / xn.abs().amax(-1, keepdim=True).clamp_min(1e-5) |
| xq = (xn * sx).round().clamp(-128, 127) / sx |
| return F.linear(xq.to(x.dtype), self.wq.to(x.dtype), self.bias) |
|
|
|
|
| def load(arm, ckpt, counts, learned_bias=None): |
| if learned_bias is None: |
| learned_bias = arm == "C" |
| m = train.Model(arm, counts, learned_bias=learned_bias).cuda().eval() |
| st = torch.load(ckpt, map_location="cuda", weights_only=False) |
| m.load_state_dict(st["model"]) |
| if arm == "C": |
| for mod in list(m.body.layers.modules()): |
| for name, child in list(mod.named_children()): |
| if isinstance(child, train.TLinear): |
| setattr(mod, name, FrozenTernary(child)) |
| return m, st.get("tag", arm) |
|
|
