File size: 5,868 Bytes
5cf6791 | 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 | """Load snapshots written by scripts/train.py (snap_*.pt / final.pt), full checkpoints, or an exported
directory (model.safetensors + config.json, as published on Hugging Face)."""
from __future__ import annotations
import json
import math
import os
import torch
from tiny_agent.model import ModelConfig, TinyAgentLM
def load_model(path: str, device="xpu", dtype=torch.float32, low_mem: bool = False) -> TinyAgentLM:
"""low_mem (serving): keep the Engram tables in bf16, memory-mapped read-only from model.safetensors, so
only the rows a prompt touches become resident. Engram.forward upcasts the gathered rows to the compute
dtype, so outputs match the full-upcast path exactly (bf16 -> fp32 is exact). Needs an exported directory
and leaves the tables frozen (requires_grad=False): not for training. The mapping stays valid as long as
model.safetensors is not rewritten in place (replacing it via rename is fine)."""
if low_mem and os.path.isdir(path):
return _load_low_mem(path, device, dtype)
if os.path.isdir(path):
from safetensors.torch import load_file
st = {"model": load_file(os.path.join(path, "model.safetensors")),
"config": json.load(open(os.path.join(path, "config.json")))}
else:
st = torch.load(path, map_location="cpu", weights_only=False)
cfg_d = st.get("config") or st["meta"]["config"]
for k in ("engram_layers", "engram_orders"):
cfg_d[k] = tuple(cfg_d[k])
model = TinyAgentLM(ModelConfig(**cfg_d))
model.load_state_dict({k: v.float() if v.is_floating_point() else v for k, v in st["model"].items()})
return model.to(device=device, dtype=dtype)
_ST_DTYPES = {"BF16": torch.bfloat16, "F16": torch.float16, "F32": torch.float32, "I64": torch.int64,
"I32": torch.int32}
def _st_view(mm, data_start: int, info: dict) -> torch.Tensor:
"""Zero-copy tensor over one safetensors entry of a read-only mmap (writing to it would fault)."""
import warnings
dt = _ST_DTYPES[info["dtype"]]
start, end = info["data_offsets"]
numel = math.prod(info["shape"])
assert end - start == numel * dt.itemsize
with warnings.catch_warnings(): # "non-writable buffer": intended, nothing writes to these
warnings.simplefilter("ignore", UserWarning)
t = torch.frombuffer(mm, dtype=dt, count=numel, offset=data_start + start) # holds a ref to mm
return t.view(info["shape"])
def _load_low_mem(path: str, device, dtype) -> TinyAgentLM:
import mmap
import struct
file = os.path.join(path, "model.safetensors")
cfg_d = json.load(open(os.path.join(path, "config.json")))
for k in ("engram_layers", "engram_orders"):
cfg_d[k] = tuple(cfg_d[k])
# build without running the weight initialisers (every parameter is replaced from the file below), so the
# torch.empty parameter storage is never touched and never becomes resident: no 1 GB fp32 table, and no
# second copy of the backbone while loading. Buffers (rope, Engram hashing) don't use nn.init.
init = torch.nn.init
real = {n: getattr(init, n) for n in dir(init) if n.endswith("_") and not n.startswith("_")}
for n in real:
setattr(init, n, lambda t, *a, **k: t)
try:
model = TinyAgentLM(ModelConfig(**cfg_d))
finally:
for n, fn in real.items():
setattr(init, n, fn)
table_keys = {k for k, _ in model.named_parameters() if k.endswith(".engram.table.weight")}
with open(file, "rb") as f:
hlen = struct.unpack("<Q", f.read(8))[0]
header = json.loads(f.read(hlen))
mm = mmap.mmap(f.fileno(), 0, access=mmap.ACCESS_READ) # stays valid after the file is closed
header.pop("__metadata__", None)
# everything but the tables: copy out (upcast like the default path), one tensor at a time
sd, v = {}, None
for k, info in header.items():
if k not in table_keys:
v = _st_view(mm, 8 + hlen, info)
sd[k] = v.float() if v.is_floating_point() else v.clone()
del v
missing, unexpected = model.load_state_dict(sd, strict=False, assign=True)
assert not unexpected and set(missing) == table_keys, (missing, unexpected)
del sd
for m in model.engram_modules():
m.table.weight = None # keep the (untouched) fp32 storage out of .to(); freed here
model.to(device=device, dtype=dtype)
if torch.device(device).type == "cpu":
# the tables live in a mapping of their own: hashed lookups are random, so no readahead
fd = os.open(file, os.O_RDONLY)
try:
tmm = mmap.mmap(fd, 0, access=mmap.ACCESS_READ)
finally:
os.close(fd)
if hasattr(tmm, "madvise") and hasattr(mmap, "MADV_RANDOM"):
tmm.madvise(mmap.MADV_RANDOM)
else:
tmm = mm
for i, b in enumerate(model.blocks):
if b.engram is not None:
w = _st_view(tmm, 8 + hlen, header[f"blocks.{i}.engram.table.weight"])
b.engram.table.weight = torch.nn.Parameter(w if tmm is not mm else w.to(device), requires_grad=False)
return model # mm (the loading map) is unmapped once nothing references it
def export(path: str, out_dir: str) -> None:
"""Write a checkpoint as out_dir/model.safetensors (bf16) + out_dir/config.json."""
from safetensors.torch import save_file
st = torch.load(path, map_location="cpu", weights_only=False)
cfg_d = dict(st.get("config") or st["meta"]["config"])
sd = {k: (v.to(torch.bfloat16) if v.is_floating_point() else v).contiguous() for k, v in st["model"].items()}
os.makedirs(out_dir, exist_ok=True)
save_file(sd, os.path.join(out_dir, "model.safetensors"))
json.dump({k: list(v) if isinstance(v, tuple) else v for k, v in cfg_d.items()},
open(os.path.join(out_dir, "config.json"), "w"), indent=1)
|