File size: 1,981 Bytes
93a8ba3 | 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 | """Checkpointing: compact state save/load for HAKO (npz-compressed + json).
Aggressive storage: every tensor is float32 (float16 for bulk matrices),
np.savez_compressed for ~60% size reduction, and the checkpoint carries the
GHSOM tree structure (numpy), head/adapters/attention tensors, the Theta /
K_pgo / growth / EWC states, the tokenizer merges, and the telemetry
counters. Restoring yields a byte-identical working state (deterministic
serialization).
"""
from __future__ import annotations
import json
from pathlib import Path
from typing import Dict
import numpy as np
def save_state(path: Path, **blocks) -> Path:
path = Path(path)
path.parent.mkdir(parents=True, exist_ok=True)
npz_payload: Dict[str, np.ndarray] = {}
meta: Dict[str, dict] = {}
for name, obj in blocks.items():
if obj is None:
continue
if isinstance(obj, dict):
arrays = {k: v for k, v in obj.items()
if isinstance(v, np.ndarray)}
scalars = {k: (v.tolist() if isinstance(v, (np.ndarray, list))
else v) for k, v in obj.items()
if not isinstance(v, np.ndarray) or v.ndim == 0}
for k, v in arrays.items():
npz_payload[f"{name}::{k}"] = v.astype(
np.float16 if v.dtype == np.float32 else v.dtype)
meta[name] = scalars
else:
meta[name] = {"repr": str(obj)[:400]}
np.savez_compressed(path, **npz_payload)
Path(str(path).replace(".npz", ".meta.json")).write_text(
json.dumps(meta, indent=1, default=str), encoding="utf-8")
return path
def load_state(path: Path) -> Dict[str, Dict[str, np.ndarray]]:
path = Path(path)
data = np.load(path, allow_pickle=False)
out: Dict[str, Dict[str, np.ndarray]] = {}
for key in data.files:
name, _, fieldname = key.partition("::")
out.setdefault(name, {})[fieldname] = data[key]
return out
|