File size: 7,958 Bytes
0f38e6b 9f4477b 0f38e6b 9f4477b 0f38e6b 9f4477b 0f38e6b 9f4477b 0f38e6b 9f4477b 0f38e6b 9f4477b 0f38e6b 9f4477b 0f38e6b 9f4477b 0f38e6b | 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 | #!/usr/bin/env python3
"""Resident-QUANTIZED any-k base-3 (ubit) — the cheat code generalized past k5.
Cody, 2026-09-17: "remember the k5 cheat code. we don't need to unpack." grouped_place's --out
path DEQUANTED every placed linear to a dense bf16 checkpoint (54 GB, and the 248320-row head's
unpack OOMs). That threw the packing away. Here the placed PackedKTritLinear buffers (packed
bytes + group scales) are saved AS-IS and reloaded resident, decoding per-forward one linear at a
time — never a global unpack, ~10 GB on disk and in VRAM. Mirrors k5_resident.load_k5_resident.
"""
import os, json, torch, torch.nn as nn
import rotation as ROT
from base3_pack import PackedKTritLinear
from ternary_gptq import _RotatedInput
def save_packed(model, out_dir, rot_map, grp_map, log=print):
"""Save the PLACED model without dequant: each PackedKTritLinear's bytes+scale + a manifest of
(k,g,rot,in,out,bias) per linear, plus the dense small tensors (embed/norms/conv/biases)."""
from safetensors.torch import save_file
os.makedirs(out_dir, exist_ok=True)
packed_stems = set()
tensors, manifest = {}, {}
wrapped = {id(m.lin) for _, m in model.named_modules() if isinstance(m, _RotatedInput)}
for name, mod in model.named_modules():
if isinstance(mod, _RotatedInput):
inner, blk = mod.lin, int(getattr(mod, "blk", 0))
elif isinstance(mod, PackedKTritLinear) and id(mod) not in wrapped:
inner, blk = mod, 0 # bare (unrotated) linear
else:
continue # skip the .lin child of a wrapper (dup)
if isinstance(inner, PackedKTritLinear):
tensors[f"{name}.bytes"] = inner.bytes.cpu().contiguous()
tensors[f"{name}.scale"] = inner.scale.cpu().contiguous()
if inner.bias is not None:
tensors[f"{name}.bias"] = inner.bias.detach().to(torch.bfloat16).cpu().contiguous()
manifest[name] = {"k": int(inner.k), "g": int(inner.group_size), "blk": blk,
"in": int(inner.in_features), "out": int(inner.out_features),
"bias": inner.bias is not None}
packed_stems.add(name)
# dense remainder: every state_dict tensor not owned by a packed linear
for k, v in model.state_dict().items():
if any(k == f"{s}.bytes" or k == f"{s}.scale" or k == f"{s}.bias" or
k.startswith(f"{s}.") for s in packed_stems):
continue
tensors[f"__dense__{k}"] = (v.to(torch.bfloat16) if v.is_floating_point() else v).cpu().contiguous()
save_file(tensors, os.path.join(out_dir, "packed.safetensors"))
json.dump(manifest, open(os.path.join(out_dir, "ubit_manifest.json"), "w"))
log(f"[save-packed] {len(manifest)} packed linears + {sum(1 for k in tensors if k.startswith('__dense__'))} "
f"dense tensors -> {out_dir}/packed.safetensors (NO dequant)")
return len(manifest)
def load_ktrit_resident(ck, device, dtype=torch.bfloat16, log=print):
"""Rebuild the placed model resident: each manifest linear as a PackedKTritLinear (rotation-
wrapped if blk>1), decode-per-forward. No global unpack."""
from transformers import AutoConfig, AutoModelForCausalLM
import transformers
from safetensors.torch import load_file
manifest = json.load(open(os.path.join(ck, "ubit_manifest.json")))
blob = load_file(os.path.join(ck, "packed.safetensors"))
config = AutoConfig.from_pretrained(ck)
dk = "dtype" if int(transformers.__version__.split(".")[0]) >= 5 else "torch_dtype"
is_vl = getattr(config, "vision_config", None) is not None # retrofitted vision+MTP -> VL class
with torch.device("meta"):
if is_vl:
from transformers import AutoModelForImageTextToText
model = AutoModelForImageTextToText.from_config(config, **{dk: dtype})
else:
model = AutoModelForCausalLM.from_config(config, **{dk: dtype})
model.eval()
# 1) swap in the packed linears
def _set(root, dotted, module):
*ps, leaf = dotted.split("."); p = root
for x in ps: p = getattr(p, x)
setattr(p, leaf, module)
skipped = []
for name, m in manifest.items():
bias = blob.get(f"{name}.bias")
pk = PackedKTritLinear(blob[f"{name}.bytes"], blob[f"{name}.scale"], m["in"], m["k"],
bias.float() if bias is not None else None, m["out"], m["g"]).to(device)
try:
_set(model, name, _RotatedInput(pk, m["blk"]) if m["blk"] and m["blk"] > 1 else pk)
except AttributeError:
skipped.append(name) # no module in this arch (e.g. MTP head absent in this transformers)
if skipped:
log(f"[load-packed] {len(skipped)} packed keys have no module in this arch (kept in file, not attached): e.g. {skipped[:3]}")
# 2) load the dense remainder (embed/norms/conv/dt_bias/…) onto meta
dense = {k[len("__dense__"):]: v for k, v in blob.items() if k.startswith("__dense__")}
model.load_state_dict(dense, strict=False, assign=True)
# 3) materialize meta rotary/pos buffers (recomputed from config, never in the checkpoint) by
# re-instantiating their owning module -- LM via text_config, vision via its dim. Any leftover
# meta tensor makes model.to(device) throw "Cannot copy out of meta tensor".
rope = getattr(config, "text_config", None) or config
vcfg = getattr(config, "vision_config", None)
def _set(root, dotted, module):
*ps, leaf = dotted.split("."); p = root
for x in ps: p = getattr(p, x)
setattr(p, leaf, module)
for mn, mod in list(model.named_modules()):
metab = [bn for bn, b in mod._buffers.items() if b is not None and b.is_meta]
if not metab:
continue
cfg_for = vcfg if "visual" in mn else rope
ib = mod._buffers.get("inv_freq")
trials = [lambda: type(mod)(config=cfg_for, device=device), lambda: type(mod)(config=cfg_for), lambda: type(mod)(cfg_for)]
if ib is not None:
dim = 2 * ib.shape[0]
trials += [lambda: type(mod)(dim), lambda: type(mod)(dim=dim), lambda: type(mod)(dim, device=device)]
made = None
for t in trials:
try: made = t(); break
except Exception: continue
if made is not None:
_set(model, mn, made.to(device))
else:
for bn in metab:
mod._buffers[bn] = torch.zeros(mod._buffers[bn].shape, dtype=torch.float32, device=device)
log(f"[load-packed] WARN zeroed meta {metab} in {mn} (rope may be off)")
model = model.to(device)
log(f"[load-packed] {len(manifest)} resident PackedKTritLinear + {len(dense)} dense tensors")
return model
if __name__ == "__main__": # toy round-trip smoke: packed save/load must be bit-exact
import tempfile
from base3_pack import quantize_ktrit, pack_ktrit
torch.manual_seed(0)
for k, g, blk in [(1, 64, 0), (2, 64, 64), (3, 128, 128)]:
lin = nn.Linear(256, 512)
codes, scale = quantize_ktrit(lin.weight.data, k, g)
b, _, _ = pack_ktrit(codes, k)
pk = PackedKTritLinear(b, scale, 256, k, lin.bias.data.float(), 512, g)
wrapped = _RotatedInput(pk, blk) if blk > 1 else pk
x = torch.randn(4, 256)
y0 = wrapped(x)
# emulate save_packed's per-module capture + load reconstruction
st = {"bytes": pk.bytes.clone(), "scale": pk.scale.clone(), "bias": pk.bias.clone()}
pk2 = PackedKTritLinear(st["bytes"], st["scale"], 256, k, st["bias"], 512, g)
w2 = _RotatedInput(pk2, blk) if blk > 1 else pk2
y1 = w2(x)
assert torch.equal(y0, y1), f"packed round-trip MISMATCH k{k} blk{blk}"
print(f"k{k} g{g} blk{blk}: packed save/load forward EXACT")
print("========== ktrit_resident round-trip OK ==========")
|