HAKO upload: hako/train/restore.py

hako/train/restore.py

@@ -0,0 +1,95 @@
+"""Restore a HAKOSystem from a checkpoint (npz + meta) for phase resume.
+
+The checkpoint stores every trainable block as numpy arrays (see
+HAKOSystem.state_blocks); restore re-instantiates the components, registers
+the decomposed sources (adapters re-seeded then overwritten by checkpoint
+values), rebuilds the GHSOM grid geometry from the stored W tensors, and
+reloads the scalar state (step counter, growth, game level).
+"""
+from __future__ import annotations
+
+import logging
+from pathlib import Path
+
+import numpy as np
+import torch
+
+from hako.core.ghsom import GHSOMNode
+from hako.core.heads import neuron_key
+from hako.train.trainer import HAKOSystem
+
+log = logging.getLogger("hako.restore")
+
+
+def restore_system(cfg, tel, mem, ckpt_path: Path) -> tuple:
+    sys_ = HAKOSystem(cfg, tel, mem)
+
+    # sources (register => adapters seeded; checkpoint overwrites below)
+    qwen_npz = Path(cfg.artifacts_dir) / "qwen05b_decomposed.npz"
+    gran_npz = Path(cfg.artifacts_dir) / "granite1b_decomposed.npz"
+    have = {}
+    if qwen_npz.exists():
+        sys_.register_source("qwen", qwen_npz, cfg.qwen_d)
+        have["qwen"] = True
+    if gran_npz.exists():
+        sys_.register_source("granite", gran_npz, cfg.granite_proto_dim)
+        have["granite"] = True
+
+    from hako.checkpoint import load_state
+    state = load_state(ckpt_path)
+    block = state.get("moe", {})
+    if block:
+        sys_.moe.load_numpy({k: v.astype(np.float32) for k, v in block.items()})
+    block = state.get("heads", {})
+    if block:
+        sys_.heads.load_numpy({k: v.astype(np.float32) for k, v in block.items()})
+    block = state.get("chain", {})
+    if block:
+        with torch.no_grad():
+            sys_.chain.theta.copy_(torch.as_tensor(block["theta"]))
+            sys_.chain.L_hat.copy_(torch.as_tensor(block["L_hat"]))
+            sys_.chain.W_beta.copy_(torch.as_tensor(block["W_beta"]))
+    block = state.get("tgnn", {})
+    if block:
+        with torch.no_grad():
+            sys_.tgnn.W_att.copy_(torch.as_tensor(block["W_att"]))
+            sys_.tgnn.W_msg.copy_(torch.as_tensor(block["W_msg"]))
+    block = state.get("router", {})
+    if block:
+        sd = {k: torch.as_tensor(v, dtype=torch.float32)
+              for k, v in block.items()}
+        sys_.router.load_state_dict(sd, strict=False)
+    block = state.get("diffusion", {})
+    if block:
+        sd = {k: torch.as_tensor(v, dtype=torch.float32)
+              for k, v in block.items()}
+        sys_.diffusion.load_state_dict(sd, strict=False)
+    gblock = state.get("ghsom", {})
+    if gblock:
+        for key, arr in gblock.items():
+            nid = int(key.split("[")[1].rstrip("]"))
+            W = torch.as_tensor(arr, dtype=torch.float32)
+            if nid not in sys_.ghsom.nodes:
+                node = GHSOMNode(nid, None, None, 0, W)
+                sys_.ghsom.nodes[nid] = node
+            else:
+                sys_.ghsom.nodes[nid].W = W
+                sys_.ghsom.nodes[nid].rows, sys_.ghsom.nodes[nid].cols = \
+                    W.shape[0], W.shape[1]
+        sys_.ghsom.root_id = min(sys_.ghsom.nodes)
+        sys_.ghsom.next_id = max(sys_.ghsom.nodes) + 1
+    block = state.get("theta", {})
+    if len(block):
+        sys_.theta.theta = np.asarray(block["theta"], dtype=np.float32)
+    block = state.get("scalars", {})
+    if block:
+        sys_.step = int(block.get("step", 0))
+        sys_.ghsom.growth_events = int(block.get("growth_events", 0))
+    # attach heads for all grid neurons and rebuild optimizer
+    for nid, node in sys_.ghsom.nodes.items():
+        for idx in range(node.rows * node.cols):
+            sys_.heads.attach(neuron_key(nid, idx))
+    sys_.rebuild_optimizer()
+    log.info("restored system from %s (step=%d, sources=%s)",
+             Path(ckpt_path).name, sys_.step, sorted(have))
+    return sys_, have