File size: 4,515 Bytes
d8b3c96 | 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 | """Pack / unpack native 1.58-bit ternary weights (5 trits per byte)."""
from __future__ import annotations
import json
from pathlib import Path
import numpy as np
import torch
from model import Boopit, BoopitConfig, ternary_and_scale
MAX_PACKED_BYTES = 7 * 1024 * 1024
def pack_trits(t: np.ndarray) -> tuple[bytes, int]:
"""Map {-1,0,1} -> {0,1,2} and pack 5 trits into each byte (3^5=243)."""
t = t.astype(np.int8).ravel()
coded = (t + 1).astype(np.uint8)
pad = (5 - (coded.size % 5)) % 5
if pad:
coded = np.concatenate([coded, np.full(pad, 1, dtype=np.uint8)]) # 0-trit
packed = (
coded[0::5].astype(np.uint16)
+ coded[1::5].astype(np.uint16) * 3
+ coded[2::5].astype(np.uint16) * 9
+ coded[3::5].astype(np.uint16) * 27
+ coded[4::5].astype(np.uint16) * 81
).astype(np.uint8)
return packed.tobytes(), pad
def unpack_trits(data: bytes, n: int, pad: int) -> np.ndarray:
packed = np.frombuffer(data, dtype=np.uint8).astype(np.uint16)
out = np.empty(packed.size * 5, dtype=np.int8)
v = packed
for i in range(5):
out[i::5] = (v % 3).astype(np.int8)
v //= 3
if pad:
out = out[: out.size - pad]
return (out[:n] - 1).astype(np.int8)
def extract_packed(model: Boopit) -> dict:
tensors = {}
seen: set[int] = set()
for name, param in model.named_parameters():
if id(param) in seen:
continue
seen.add(id(param))
w = param.detach()
if w.ndim == 2:
t, scale = ternary_and_scale(w)
packed, pad = pack_trits(t.cpu().numpy())
tensors[name] = {
"kind": "ternary",
"shape": list(w.shape),
"scale": float(scale.cpu()),
"pad": pad,
"data": packed,
}
else:
tensors[name] = {
"kind": "fp16",
"shape": list(w.shape),
"data": w.detach().to(torch.float16).cpu().numpy().tobytes(),
}
return tensors
def packed_nbytes(tensors: dict) -> int:
return sum(len(v["data"]) for v in tensors.values())
def save_packed(model: Boopit, path: Path, extra: dict | None = None) -> int:
path = Path(path)
tensors = extract_packed(model)
nbytes = packed_nbytes(tensors)
payload = {
"format": "boopit-1.58",
"config": model.config.to_dict(),
"tied": ["tok_emb.weight", "lm_head.weight"],
"extra": extra or {},
"tensors": {},
}
blob_parts = []
offset = 0
for name, spec in tensors.items():
data = spec["data"]
entry = {k: v for k, v in spec.items() if k != "data"}
entry["offset"] = offset
entry["nbytes"] = len(data)
payload["tensors"][name] = entry
blob_parts.append(data)
offset += len(data)
header = json.dumps(payload, separators=(",", ":")).encode("utf-8")
header_len = len(header).to_bytes(8, "little")
blob = b"".join(blob_parts)
raw = header_len + header + blob
if len(raw) >= MAX_PACKED_BYTES:
raise RuntimeError(f"packed model is {len(raw)} bytes, limit is {MAX_PACKED_BYTES}")
path.write_bytes(raw)
return len(raw)
def load_packed(path: Path, device: torch.device | str = "cpu") -> Boopit:
raw = Path(path).read_bytes()
header_len = int.from_bytes(raw[:8], "little")
payload = json.loads(raw[8 : 8 + header_len].decode("utf-8"))
blob = raw[8 + header_len :]
cfg = BoopitConfig.from_dict(payload["config"])
model = Boopit(cfg)
named = dict(model.named_parameters())
loaded: set[int] = set()
with torch.no_grad():
for name, spec in payload["tensors"].items():
param = named[name]
chunk = blob[spec["offset"] : spec["offset"] + spec["nbytes"]]
if spec["kind"] == "ternary":
t = unpack_trits(chunk, int(np.prod(spec["shape"])), spec["pad"])
w = torch.from_numpy(t.astype(np.float32).reshape(spec["shape"])) * float(spec["scale"])
param.copy_(w.to(dtype=param.dtype))
else:
arr = np.frombuffer(chunk, dtype=np.float16).reshape(spec["shape"])
param.copy_(torch.from_numpy(arr.copy()).to(dtype=param.dtype))
loaded.add(id(param))
# tied lm_head
if id(model.lm_head.weight) not in loaded:
model.lm_head.weight = model.tok_emb.weight
return model.to(device)
|