1"""Shared training utilities: run ids + manifests, seeding, device selection, token data, checkpoints, logging."""2from __future__ import annotations34import hashlib, json, os, platform, random, subprocess, sys, time5from dataclasses import dataclass6from pathlib import Path7from typing import Any, Iterator, Optional89import numpy as np10import torch11import yaml121314# ----------------------------------------------------------------------------- environment15def git_sha() -> str:16 try:17 return subprocess.run(["git", "rev-parse", "HEAD"], capture_output=True, text=True, timeout=5).stdout.strip() or "unknown"18 except Exception: # noqa: BLE00119 return "unknown"202122def pick_device(pref: str = "auto") -> torch.device:23 if pref != "auto":24 return torch.device(pref)25 if torch.cuda.is_available():26 return torch.device("cuda")27 if torch.backends.mps.is_available():28 return torch.device("mps")29 return torch.device("cpu")303132def seed_all(seed: int) -> None:33 random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)34 if torch.cuda.is_available():35 torch.cuda.manual_seed_all(seed)363738def sync(device: torch.device) -> None:39 if device.type == "mps":40 torch.mps.synchronize()41 elif device.type == "cuda":42 torch.cuda.synchronize()434445def peak_memory_bytes(device: torch.device) -> Optional[int]:46 try:47 if device.type == "mps":48 return int(torch.mps.driver_allocated_memory())49 if device.type == "cuda":50 return int(torch.cuda.max_memory_allocated())51 except Exception: # noqa: BLE00152 return None53 try:54 import resource55 return int(resource.getrusage(resource.RUSAGE_SELF).ru_maxrss)56 except Exception: # noqa: BLE00157 return None585960def sha256_file(p: str | Path) -> str:61 h = hashlib.sha256()62 with open(p, "rb") as f:63 for chunk in iter(lambda: f.read(1 << 20), b""):64 h.update(chunk)65 return h.hexdigest()666768def load_yaml(path: str | Path) -> dict:69 with open(path) as f:70 return yaml.safe_load(f)717273# ----------------------------------------------------------------------------- run bookkeeping74class Run:75 """Immutable run id + manifest + jsonl metrics log under experiments/runs/<run_id>/."""7677 def __init__(self, name: str, config: dict, root: str = "experiments/runs", seed: int = 0, resume_id: str | None = None):78 ts = time.strftime("%Y%m%d%H%M%S")79 self.run_id = resume_id or f"{name}_seed{seed:02d}_{ts}"80 self.dir = Path(root) / self.run_id81 self.dir.mkdir(parents=True, exist_ok=True)82 self.config = config83 self.t0 = time.time()84 self.manifest = {85 "run_id": self.run_id, "name": name, "seed": seed, "code_commit": git_sha(),86 "config_sha256": hashlib.sha256(json.dumps(config, sort_keys=True, default=str).encode()).hexdigest(),87 "started_at": time.strftime("%Y-%m-%dT%H:%M:%S%z"), "host": platform.node(), "python": sys.version.split()[0],88 "torch": torch.__version__, "command": " ".join(sys.argv), "state": "running",89 }90 (self.dir / "config.yaml").write_text(yaml.safe_dump(config, sort_keys=False))91 self._log = open(self.dir / "metrics.jsonl", "a")92 self.write_manifest()9394 def write_manifest(self, **extra: Any) -> None:95 self.manifest.update(extra)96 self.manifest["wall_seconds"] = round(time.time() - self.t0, 2)97 (self.dir / "manifest.json").write_text(json.dumps(self.manifest, indent=2, default=str))9899 def log(self, **kv: Any) -> None:100 kv["t"] = round(time.time() - self.t0, 2)101 self._log.write(json.dumps(kv, default=float) + "\n"); self._log.flush()102103 def finish(self, state: str = "completed", **extra: Any) -> None:104 self.write_manifest(state=state, finished_at=time.strftime("%Y-%m-%dT%H:%M:%S%z"), **extra)105 self._log.close()106107108# ----------------------------------------------------------------------------- token data109class TokenBin:110 """Memory-mapped uint16 token stream (documents separated by EOT). Random contiguous windows of length T+1."""111112 def __init__(self, path: str | Path, seq_len: int):113 self.path = Path(path)114 self.data = np.memmap(self.path, dtype=np.uint16, mode="r")115 self.seq_len = seq_len116117 def __len__(self) -> int:118 return len(self.data)119120 def batches(self, batch_size: int, seed: int, device: torch.device) -> Iterator[tuple[torch.Tensor, torch.Tensor]]:121 rng = np.random.default_rng(seed)122 n = len(self.data) - self.seq_len - 1123 while True:124 idx = rng.integers(0, n, size=batch_size)125 x = np.stack([self.data[i:i + self.seq_len + 1].astype(np.int64) for i in idx])126 xt = torch.from_numpy(x)127 yield xt[:, :-1].to(device), xt[:, 1:].to(device)128129 def sequential_batches(self, batch_size: int, device: torch.device, max_batches: int | None = None):130 """Deterministic non-overlapping windows for evaluation."""131 n = (len(self.data) - 1) // self.seq_len132 count = 0133 for start in range(0, n, batch_size):134 idx = [i * self.seq_len for i in range(start, min(start + batch_size, n))]135 x = np.stack([self.data[i:i + self.seq_len + 1].astype(np.int64) for i in idx])136 xt = torch.from_numpy(x)137 yield xt[:, :-1].to(device), xt[:, 1:].to(device)138 count += 1139 if max_batches is not None and count >= max_batches:140 break141142143# ----------------------------------------------------------------------------- optimisation helpers144def cosine_lr(step: int, warmup: int, total: int, lr: float, min_ratio: float = 0.1) -> float:145 if step < warmup:146 return lr * (step + 1) / warmup147 p = min(1.0, (step - warmup) / max(1, total - warmup))148 import math149 return lr * (min_ratio + (1 - min_ratio) * 0.5 * (1 + math.cos(math.pi * p)))150151152def make_adamw(params, lr: float, weight_decay: float, betas=(0.9, 0.95)):153 decay, no_decay = [], []154 for p in params:155 (decay if p.ndim >= 2 else no_decay).append(p)156 return torch.optim.AdamW([{"params": decay, "weight_decay": weight_decay}, {"params": no_decay, "weight_decay": 0.0}],157 lr=lr, betas=betas)158159160def save_checkpoint(path: str | Path, model: torch.nn.Module, cfg_dict: dict, extra: dict | None = None) -> str:161 from safetensors.torch import save_file162 path = Path(path); path.mkdir(parents=True, exist_ok=True)163 sd = {k: v.detach().to("cpu").contiguous() for k, v in model.state_dict().items()}164 save_file(sd, str(path / "model.safetensors"))165 meta = {"config": cfg_dict, "saved_at": time.strftime("%Y-%m-%dT%H:%M:%S%z"), **(extra or {})}166 meta["weights_sha256"] = sha256_file(path / "model.safetensors")167 (path / "checkpoint.json").write_text(json.dumps(meta, indent=2, default=str))168 return meta["weights_sha256"]169170171def load_checkpoint(path: str | Path, model: torch.nn.Module, strict: bool = True) -> dict:172 from safetensors.torch import load_file173 path = Path(path)174 sd = load_file(str(path / "model.safetensors"))175 own = model.state_dict()176 mismatched = [k for k, v in sd.items() if k in own and own[k].shape != v.shape]177 if mismatched and strict:178 raise RuntimeError(f"shape mismatch for {mismatched}")179 for k in mismatched: # e.g. a controller re-shaped by a feature option: left at its fresh init180 sd.pop(k)181 missing, unexpected = model.load_state_dict(sd, strict=strict)182 meta = json.loads((path / "checkpoint.json").read_text())183 meta["_missing"] = list(missing) + mismatched; meta["_unexpected"] = list(unexpected); meta["_shape_mismatch"] = mismatched184 return meta185