Repository Public Research preview
1"""Deterministic training of a Tier-0 toy model on a single-answer-token micro-world.23Loss is the teacher-forced cross-entropy of the answer (scratchpad) tokens only. Training is used to4*manufacture* donor competence and recipient weakness; it is documented as a separate stage and5is never part of a transplant.6"""7from __future__ import annotations89import json10import math11import time12from pathlib import Path1314import numpy as np15import torch16import torch.nn.functional as F1718from wp_ai_0001.config import ToyModelConfig, TrainConfig19from wp_ai_0001.datasets.common import Split, Vocab, encode_split20from wp_ai_0001.models.decoding import greedy_decode21from wp_ai_0001.models.toy_transformer import ToyTransformer, resolve_device222324def set_seed(seed: int) -> None:25 torch.manual_seed(seed)26 np.random.seed(seed)272829@torch.no_grad()30def accuracy(model: ToyTransformer, enc: dict[str, np.ndarray], device: torch.device, bs: int = 1024) -> tuple[float, float, np.ndarray, np.ndarray]:31 """Greedy-decoding accuracy: (final-answer acc, full-trace acc, final_correct [N], trace_correct [N])."""32 model.eval()33 K = enc["answers"].shape[1]34 fin, tr = [], []35 for i in range(0, len(enc["ids"]), bs):36 ids = torch.from_numpy(enc["ids"][i : i + bs]).to(device)37 pl = torch.from_numpy(enc["prompt_len"][i : i + bs]).to(device)38 na = torch.from_numpy(enc["n_answer"][i : i + bs]).to(device)39 gold = torch.from_numpy(enc["answers"][i : i + bs]).to(device)40 pred = greedy_decode(model, ids, pl, na, K)41 ar = torch.arange(len(pl), device=device)42 fin.append((pred[ar, na - 1] == gold[ar, na - 1]).cpu().numpy())43 valid = torch.arange(K, device=device)[None] < na[:, None]44 tr.append(((pred == gold) | ~valid).all(-1).cpu().numpy())45 f, t = np.concatenate(fin), np.concatenate(tr)46 return float(f.mean()), float(t.mean()), f, t474849def masked_loss(model: ToyTransformer, ids: torch.Tensor, loss_mask: torch.Tensor) -> torch.Tensor:50 logits = model(ids[:, :-1])51 target = ids[:, 1:]52 mask = loss_mask[:, :-1]53 ce = F.cross_entropy(logits.reshape(-1, logits.shape[-1]), target.reshape(-1), reduction="none").view_as(mask)54 return (ce * mask).sum() / mask.sum()555657def train(model_cfg: ToyModelConfig, train_cfg: TrainConfig, train_split: Split, dev_split: Split, vocab: Vocab, out_dir: Path, log: list[dict] | None = None) -> ToyTransformer:58 set_seed(train_cfg.seed)59 device = resolve_device(train_cfg.device)60 model = ToyTransformer(model_cfg, len(vocab)).to(device)61 enc = encode_split(train_split, vocab, model_cfg.max_seq_len)62 denc = encode_split(dev_split, vocab, model_cfg.max_seq_len)63 ids, lmask = enc["ids"], enc["loss_mask"]64 opt = torch.optim.AdamW(model.parameters(), lr=train_cfg.lr, weight_decay=train_cfg.weight_decay, betas=(0.9, 0.98))65 g = torch.Generator().manual_seed(train_cfg.seed)66 n = len(ids)67 t0 = time.time()68 history = log if log is not None else []69 for step in range(1, train_cfg.steps + 1):70 lr = train_cfg.lr * min(1.0, step / train_cfg.warmup_steps) * 0.5 * (1 + math.cos(math.pi * min(1.0, step / train_cfg.steps)))71 for grp in opt.param_groups:72 grp["lr"] = lr73 idx = torch.randint(0, n, (train_cfg.batch_size,), generator=g).numpy()74 x = torch.from_numpy(ids[idx]).to(device)75 m = torch.from_numpy(lmask[idx]).to(device)76 model.train()77 loss = masked_loss(model, x, m)78 opt.zero_grad(set_to_none=True)79 loss.backward()80 torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)81 opt.step()82 if step % train_cfg.eval_every == 0 or step == train_cfg.steps:83 dev_acc, dev_trace, _, _ = accuracy(model, denc, device)84 history.append({"step": step, "loss": float(loss.item()), "dev_acc": dev_acc, "dev_trace_acc": dev_trace, "lr": lr, "elapsed_s": time.time() - t0})85 out_dir.mkdir(parents=True, exist_ok=True)86 torch.save(model.state_dict(), out_dir / "model.pt")87 (out_dir / "model_config.json").write_text(model_cfg.model_dump_json(indent=2))88 (out_dir / "train_config.json").write_text(train_cfg.model_dump_json(indent=2))89 (out_dir / "vocab.json").write_text(vocab.to_json())90 (out_dir / "history.json").write_text(json.dumps(history, indent=2))91 return model929394def load(out_dir: Path, device: torch.device | None = None) -> tuple[ToyTransformer, Vocab]:95 cfg = ToyModelConfig.model_validate_json((out_dir / "model_config.json").read_text())96 vocab = Vocab.from_json((out_dir / "vocab.json").read_text())97 model = ToyTransformer(cfg, len(vocab))98 model.load_state_dict(torch.load(out_dir / "model.pt", map_location="cpu"))99 model.eval()100 if device is not None:101 model.to(device)102 return model, vocab103