8000 mini-agi/train.py at main · ssemiya/mini-agi · GitHub
[go: up one dir, main page]
Skip to content

train.py

Latest commit

 

History

History
2277 lines (2136 loc) · 117 KB

train.py

File metadata and controls

2277 lines (2136 loc) · 117 KB
Raw
#!/usr/bin/env python3
"""
Training mini-AGI.
read point it at files and let it read them, continually. This is the
one that matters, and the one the run uses.
stream the same mechanism over a packed corpus: batch 1, a KV cache,
one chunk of characters at a time. Peak memory is set by
--chunk and nothing else, so --context is nearly free to extend.
probe what depth the model chooses, per character.
Everything here runs at batch 1 behind a cache. The fixed-window regime that
preceded it - several windows per step, one autograd graph over each whole
window - cost about three times as much at an eighth of the reach: 6857 MB at
4k characters against 3090 MB at 32k, because the retained activation graph
was bounded by the window rather than by the chunk.
"""
import os
import re
import sys
import json
import collections
import math
import time
import argparse
from dataclasses import asdict
# BEFORE torch, because the allocator reads this once at CUDA init and ignores
# it afterwards. Without it PyTorch's caching allocator keeps its free blocks in
# fixed segments, and a run that alternates between a few large short-lived
# tensors - which is exactly what a checkpointed expert dispatch does, three
# batched matmuls recomputed in the backward - strands memory it cannot hand
# back. Measured on an 8,192 window: the backward asked for 384 MiB with 394
# MiB free and 861 MiB reserved but unallocated. There was plenty of memory;
# there was no contiguous piece of it. Set it here rather than in a shell so it
# is a property of the program, not of how it happened to be launched.
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
import numpy as np
import torch
import torch.nn as nn
from minagi.recur import RecurConfig, RecurCoder, load_recur
from minagi.pool import PooledMLP, AutoGrow
from minagi.stream import StreamSet, Evaluator, detach_caches, ramp_context
from minagi.plasticity import Plasticity
from minagi.optim import GradSNR
from minagi import store as weights_store
def lr_at(step, total, base, warmup, floor_frac=0.1):
if step < warmup:
return base * (step + 1) / warmup
prog = (step - warmup) / max(1, total - warmup)
return base * (floor_frac + (1 - floor_frac) * 0.5 * (1 + math.cos(math.pi * prog)))
def _resync_opt(opt, model, args):
"""Drop parameters the model no longer has, adopt the ones it gained.
Growth appends experts and replaces the gate tensor; pruning removes them
and replaces it again. Either way the optimiser is left holding tensors the
model does not have - Adam would keep stepping them and keep their moments
alive, a leak on a pool that grows and shrinks all run.
"""
live = {id(p) for p in model.parameters()}
for g in opt.param_groups:
for p_ in g["params"]:
if id(p_) not in live:
opt.state.pop(p_, None)
g["params"] = [p for p in g["params"] if id(p) in live]
known = {id(p) for g in opt.param_groups for p in g["params"]}
fresh = [p for p in model.parameters() if id(p) not in known]
mats = [p for p in fresh if p.dim() >= 2]
vecs = [p for p in fresh if p.dim() < 2]
# a group added mid-run needs the anchor the schedule scales from, or it
# would sit at whatever rate it was born with while everything else decays
base = args.lr
if mats:
opt.add_param_group({"params": mats, "name": "pool",
"weight_decay": args.wd,
"lr": base, "base_lr": base})
if vecs:
opt.add_param_group({"params": vecs, "name": "pool",
"weight_decay": 0.0,
"lr": base, "base_lr": base})
def cmd_stream(args):
"""
Train the way the model runs: batch 1, a KV cache, one chunk at a time.
The old loop built one autograd graph over the whole window. That cost
about four times what reading the same window costs - 6857 MB against 1704
- and the difference was entirely activations retained across 26 block
applications. So the model could be served at a context it could not be
trained at, and training kept hitting OOM on a card with room to spare.
Here the graph is bounded by --chunk and the reach is bounded by the cache,
which is linear and cheap. Batch 1 gives up gradient averaging; what it
buys is roughly an order of magnitude of context on the same card.
"""
device = torch.device(args.device)
torch.manual_seed(args.seed)
rng = np.random.default_rng(args.seed)
sources = []
for part in args.mix.split(","):
name, _, w = part.partition(":")
sources.append((name.strip(), float(w) if w else 1.0))
wdir = args.weights_dir
man = None
if os.path.exists(os.path.join(wdir, "manifest.json")):
with open(os.path.join(wdir, "manifest.json")) as f:
man = json.load(f)
cfgd = dict(man["cfg"]) if man and man.get("cfg") else {}
cfgd.update(vocab_size=265, block=args.context, use_pool=True)
cfgd["pool_max"] = args.pool_max
cfg = RecurConfig(**{k: v for k, v in cfgd.items()
if k in RecurConfig.__dataclass_fields__})
# The capacity bound comes from config.yaml rather than from the
# checkpoint, because it is a property of the machine the model is running
# on - how much VRAM the dispatch may use - not of the model. A directory
# written before this existed carries no value for it and would otherwise
# get the dataclass default, which is right but silent; taking it from the
# config every load means the number in the file is the number in effect.
try:
from minagi.config import load as _load_cfg, get as _get_cfg
cfg.pool_capacity_factor = float(
_get_cfg(_load_cfg(), "pool.capacity_factor",
cfg.pool_capacity_factor))
except Exception:
pass
model = RecurCoder(cfg).to(device)
if man and int(man["n_experts"]) != model.pool.n_experts():
have, want = model.pool.n_experts(), int(man["n_experts"])
if want > have:
model.pool.add_experts(want - have, device=device)
print(f" pool resized {have} -> {want} to match the directory")
print(f"params {model.n_params()/1e6:.2f}M | {model.pool.n_experts()} "
f"experts | block-applications {cfg.n_layer_effective}")
print(f"batch 1 | chunk {args.chunk} | context up to {cfg.block:,}")
print(f" the graph spans {args.chunk} characters; the cache spans the "
f"context. only the first costs activations.")
pool_ids = set()
sites = [m for m in model.modules() if isinstance(m, PooledMLP)]
pool_ids = {id(q) for e in model.pool.experts for q in e.parameters()}
pool_ids |= {id(st.router.weight) for st in sites}
pool_ids |= {id(st.depth_emb) for st in sites}
pool_ids.add(id(model.pool.gate))
trunk = [q for q in model.parameters() if id(q) not in pool_ids]
pool = [q for q in model.parameters() if id(q) in pool_ids]
opt = torch.optim.AdamW(
[{"params": [q for q in trunk if q.dim() >= 2], "name": "trunk",
"weight_decay": args.wd},
{"params": [q for q in trunk if q.dim() < 2], "name": "trunk",
"weight_decay": 0.0},
{"params": [q for q in pool if q.dim() >= 2], "name": "pool",
"weight_decay": args.wd},
{"params": [q for q in pool if q.dim() < 2], "name": "pool",
"weight_decay": 0.0}],
lr=args.lr, betas=(0.9, 0.95), fused=(device.type == "cuda"))
print(f"trunk learns at {args.trunk_lr_mult:g}x the pool's rate "
f"({sum(q.numel() for q in trunk)/1e6:.2f}M trunk, "
f"{sum(q.numel() for q in pool)/1e6:.2f}M pool)")
start_step = 0
if man:
weights_store.load(model, wdir, opt=opt, device=device, verbose=True)
start_step = int(man.get("step", -1)) + 1
print(f"resuming from {wdir} at step {start_step}")
grower = None
if args.grow_k > 0:
# args.pool_max, not cfg.pool_max - the latter is the router width
grower = AutoGrow(grow_k=args.grow_k, max_experts=args.pool_max,
mem_frac_max=args.grow_mem_frac)
model.pool.born.fill_(float(start_step))
model.pool.gate_seen.copy_(model.pool.gate.data.abs())
print(f"auto-growth armed: {model.pool.n_experts()} experts, ceiling "
f"{args.pool_max}, VRAM brake at {100*args.grow_mem_frac:.0f}%")
streams = StreamSet(model, sources, args.chunk, args.ctx_start, device,
rng, replay=args.replay,
replay_burst=args.replay_burst)
evaluator = Evaluator(model, [n for n, _ in sources], args.chunk,
args.ctx_start, device)
tot = sum(w for _, w in sources)
for name, w in sources:
n_ch = len(np.memmap(os.path.join(name, "train.bin"),
dtype=np.uint16, mode="r"))
print(f" {name:<22}{100*w/tot:>4.0f}% {n_ch/1e6:>7.1f}M characters")
if args.ctx_start < args.context:
print(f"context ramps {args.ctx_start:,} -> {args.context:,} over the "
f"first third of the run")
os.makedirs(args.out, exist_ok=True)
# A run keeps its own history rather than leaving it in a log. A log is one
# shell redirect away from being overwritten by the next run, and two
# processes writing one log here once produced a number that looked like a
# model bug and was not. Appended a line at a time, so an interrupted run
# keeps everything up to the interruption.
hist_path = os.path.join(args.out, "history.jsonl")
hist = open(hist_path, "a", buffering=1)
def record(kind, step=None, **kw):
hist.write(json.dumps({"kind": kind, "step": step, **kw}) + "\n")
record("start", weights_dir=wdir, context=cfg.block, chunk=args.chunk,
accum=args.accum, lr=args.lr, trunk_lr_mult=args.trunk_lr_mult,
steps=args.steps, start_step=start_step,
experts=model.pool.n_experts(), params=model.n_params(),
mix=[{"name": n, "weight": w} for n, w in sources])
# the directory's own score, not infinity - resuming must not let a worse
# state call itself the best merely because this process just started
best = weights_store.best_val(wdir)
if best < float("inf"):
print(f"best on disk: val {best:.4f} - weights/ advances only past it")
reverts = 0
t0, seen = time.time(), 0
model.train()
for step in range(start_step, args.steps):
lr = lr_at(step, args.steps, args.lr, args.warmup)
for g in opt.param_groups:
g["lr"] = lr * (args.trunk_lr_mult if g["name"] == "trunk" else 1.0)
ctx_now = ramp_context(step - start_step, max(1, args.steps - start_step),
args.ctx_start, args.context)
if ctx_now != streams.context:
streams.set_context(ctx_now)
evaluator.context = ctx_now
print(f" context -> {ctx_now:,} characters", flush=True)
# Accumulate over consecutive chunks before stepping. Each chunk's
# graph is freed by its own backward, so this costs no memory beyond
# one chunk - but the gradient the optimiser sees averages
# accum x chunk characters instead of chunk. That is the whole gap
# between this and the batch trainer: at chunk 512 a step saw 512
# characters where a batch step saw 32768, and with accumulation the
# two can be made identical without giving up the memory profile.
opt.zero_grad(set_to_none=True)
for _ in range(args.accum):
loss, src = streams.step(learn=True, aux_weight=cfg.pool_aux)
(loss / args.accum).backward()
detach_caches(streams.active.caches)
seen += args.chunk
gn = torch.nn.utils.clip_grad_norm_(model.parameters(), args.clip)
opt.step()
if step % args.log_every == 0:
el = max(time.time() - t0, 1e-9)
vram = (f" vram {torch.cuda.max_memory_allocated()/1e6:.0f}MB"
if device.type == "cuda" else "")
record("step", step=step, loss=float(loss), lr=lr,
grad_norm=float(gn), chars=seen, chars_per_s=seen / el,
context=streams.context, experts=model.pool.n_experts(),
vram_mb=(torch.cuda.max_memory_allocated() / 1e6
if device.type == "cuda" else None))
rp = (f" replay {streams.replays}" if args.replay > 0 else "")
print(f"step {step:>6}/{args.steps} loss {float(loss):.4f} "
f"lr {lr:.2e} gn {float(gn):.2f} ctx {streams.context//1024}k "
f"{seen/el/1e3:.1f}k char/s{vram}{rp}", flush=True)
if (step + 1) % args.eval_every == 0 or step + 1 == args.steps:
v = evaluator.run(args.eval_chunks)
se = v.pop("stderr", 0.0)
val = float(np.mean(list(v.values())))
record("val", step=step, val=val, stderr=se, ppl=math.exp(val),
per_domain=dict(v), experts=model.pool.n_experts())
print(f" val {val:.4f} +/-{se:.4f} ppl {math.exp(val):.2f} "
+ " ".join(f"{k.replace('data_','').replace('_char','')} "
f"{x:.3f}" for k, x in v.items()), flush=True)
if grower is not None:
gone = model.pool.prune(step, survival=args.grow_min_age)
if gone:
_resync_opt(opt, model, args)
record("pruned", step=step, n=gone,
to=model.pool.n_experts())
print(f" pruned {gone} -> {model.pool.n_experts()} experts",
flush=True)
rec = grower.step(val, model.pool, step)
if rec["grew"]:
_resync_opt(opt, model, args)
record("grew", step=step, n=rec["grew"], to=rec["experts"],
reason=rec.get("reason"))
print(f" GREW +{rec['grew']} -> {rec['experts']} experts "
f"| {rec['reason']}", flush=True)
elif rec.get("reason"):
record("held", step=step, reason=rec["reason"],
experts=model.pool.n_experts())
print(f" [pool] {rec['reason']}", flush=True)
# The directory holds the best state the model has reached, so it
# is written when a run improves on it and left alone otherwise.
# Keeping a second copy aside as a snapshot would double 236 MB
# that grows with the pool, to hold a copy of what is already here.
if val < best:
best = val
weights_store.save(model, wdir, step=step, val=val, opt=opt,
cfg=asdict(cfg), verbose=True)
record("saved", step=step, val=val)
print(f" improved - weights/ advanced to val {val:.4f}",
flush=True)
elif args.revert_factor > 0 and val > best * args.revert_factor:
# A run can diverge: extending the context too fast once took
# validation from 0.7375 to 5.4490 in a thousand steps. Without
# this the live directory would carry that away with it.
print(f" DIVERGED: val {val:.4f} is {val/best:.1f}x the best "
f"{best:.4f} - reloading weights/ and halving the "
f"learning rate", flush=True)
with open(os.path.join(wdir, "manifest.json")) as _f:
man_b = json.load(_f)
n_back = int(man_b["n_experts"])
if n_back != model.pool.n_experts():
have = model.pool.n_experts()
if n_back > have:
model.pool.add_experts(n_back - have, device=device)
else:
model.pool.experts = nn.ModuleList(
list(model.pool.experts)[:n_back])
model.pool.gate = nn.Parameter(
model.pool.gate.data[:n_back].clone())
for b in ("use", "age", "born", "gate_seen"):
setattr(model.pool, b,
getattr(model.pool, b)[:n_back].clone())
model.pool.invalidate()
_resync_opt(opt, model, args)
weights_store.load(model, wdir, opt=opt, device=device)
args.lr *= 0.5
streams.set_context(max(args.ctx_start, streams.context // 2))
evaluator.context = streams.context
reverts += 1
print(f" back to val {man_b['val']:.4f} at "
f"{model.pool.n_experts()} experts, lr now "
f"{args.lr:.2e}, context back to {streams.context:,}",
flush=True)
if reverts >= args.max_reverts:
print(f" {reverts} reverts - stopping rather than "
f"thrashing", flush=True)
break
model.train()
record("done", minutes=(time.time() - t0) / 60, best=best,
experts=model.pool.n_experts())
hist.close()
print(f"\ndone in {(time.time()-t0)/60:.1f}m best val {best:.4f}")
print(f"history: {hist_path}")
return 0
class _Tracer:
"""
Record what the real training forward did, for the first few chunks.
Not a reconstruction: this sits inside the training loop, around the same
forward whose loss is about to be backpropagated, with the same KV cache,
the same precision, the same chunk and the same working set. Whatever it
writes down is what actually happened.
The capture context is exited before backward on purpose. Gradient
checkpointing replays the forward during backward, and a capture still
open would record every call site twice.
"""
def __init__(self, path, chunks, swap_chunks=160):
self.path, self.want = path, chunks
self.want_swaps = max(swap_chunks, chunks)
self.ids, self.wts, self.meta, self.text = [], [], [], []
self.swaps = []
def full(self):
"""Routes are heavy, so only the first few chunks keep them."""
return len(self.ids) >= self.want
def done(self):
"""Swap counts are cheap and run on well past the routes."""
return self.full() and len(self.swaps) >= self.want_swaps
def add_swap(self, lane, moved, first, demand_moved=None):
"""
One line per working-set decision, whether or not routes are traced.
`first` marks the chunk that opens a lane's window - the only place
the subject actually changes. If swaps have a reason, they belong
here and nowhere else; if they are on a timer, this column will make
no difference to the count.
"""
self.swaps.append({"subject": lane.name, "moved": int(moved),
"first_of_window": bool(first)})
def add(self, got, pool, lane, moved, step, loss, seen, ids):
import numpy as _np
if self.ids and ids.shape[1] != self.text[0].shape[0]:
return # a short tail chunk would not stack
self.ids.append(_np.stack([g[0].numpy() for g in got]).astype("int16"))
self.wts.append(_np.stack([g[1].numpy() for g in got]).astype("float16"))
self.text.append(ids[0].detach().cpu().numpy().astype("int16"))
self.meta.append({
"subject": lane.name, "moved": int(moved), "step": int(step),
"loss": float(loss), "seen": int(seen),
"resident": [int(e) for e in pool.slots if e >= 0],
"experts": int(pool.n_experts()),
})
written = False
def write(self, cfg, precision, chunk):
import json as _json
import numpy as _np
self.written = True
_np.savez_compressed(
self.path,
ids=_np.stack(self.ids), weights=_np.stack(self.wts),
text=_np.stack(self.text),
meta=_json.dumps({"chunks": self.meta, "swaps": self.swaps,
"chunk": chunk,
"top_k": cfg.pool_top_k,
"max_steps": cfg.max_steps,
"n_recur": cfg.n_recur, "n_coda": cfg.n_coda,
"n_prelude": cfg.n_prelude,
"context": cfg.block,
"precision": precision}))
print(f" routing trace -> {self.path} "
f"({len(self.ids)} chunks of real training)", flush=True)
class _Lane:
"""
One subject, sampled a window at a time.
Each visit opens a random file and a random offset inside it, then reads
one context window forward from there. The window is contiguous on
purpose: a lone chunk carries no history, so attention would have nothing
to reach back over and the context window would be decoration.
"""
def __init__(self, name, paths, rng):
self.name, self.paths, self.rng = name, paths, rng
self.r = None
self.visits = 0
def open(self, model, chunk, block, device, span=None):
"""
Open a file for this visit.
`block` is the context window; `span` is how much will be read before
the reader moves to another subject, which may be several windows. The
offset has to leave room for the whole visit, not just the first
window, or the file runs out part way through and the visit is cut
short.
"""
from minagi.stream import FileReader
from minagi.ingest import as_stream
span = max(span or block, block)
for _ in range(8): # a few tries for short files
path = self.paths[int(self.rng.integers(0, len(self.paths)))]
try:
data = as_stream(path)
except OSError:
# The corpus can change under a run that lasts days. The
# self-knowledge lane in particular is REGENERATED from the
# source, which deletes and rewrites every file in it, and the
# path list was taken at startup. A vanished file is a reason
# to pick another one, not to end a week of reading.
continue
if len(data) < 8:
continue
r = FileReader(model, data, path, chunk, block, device)
if len(data) > span + 1:
r.pos = int(self.rng.integers(0, len(data) - span - 1))
self.r = r
self.visits += 1
return r
return None
def rest(self, model):
"""
Drop the KV cache between visits.
A visit is exactly one context window, which is where FileReader
throws its cache away anyway - so this costs nothing, and it is what
keeps a single window of keys and values resident rather than one per
lane. At 26 block-applications a window is 1.7 GB, so four lanes
holding their own would not fit on the card at all.
"""
if self.r is not None:
self.r.seen = 0
self.r = None
def _lanes(files, seed, roots=None, resume=0):
"""
One lane per subject, where a subject is a TOP-LEVEL folder of the corpus.
Not the immediate parent: a corpus of any size grows subdirectories - a
hundred and fifty thousand articles cannot sit in one directory - and
grouping by the parent turns every one of those into its own subject. It
also matches what a reader means by their own files: pointing at ~/notes
with folders inside should give one subject per folder, not per nested
directory.
"""
roots = [os.path.abspath(r) for r in (roots or [])]
def subject(path):
p = os.path.abspath(path)
for r in roots:
if p.startswith(r + os.sep):
rel = os.path.relpath(p, r).split(os.sep)
# the first component under the root, or the root itself for
# a file lying directly in it
return rel[0] if len(rel) > 1 else os.path.basename(r)
return os.path.basename(os.path.dirname(p)) or p
by = {}
for f in files:
by.setdefault(subject(f), []).append(f)
# `resume` is how many characters the model has already read, and leaving
# it out was the single most expensive bug in this project.
#
# A lane picks a random file and a random offset inside it, from its own
# generator. Seeded by `seed + j` alone, that generator is rebuilt
# identically at every start, so every session replayed the SAME files at
# the SAME offsets in the SAME order, from the beginning. Measured over 43
# sessions: the median one read 3.2M characters, the longest 66M, and the
# counter reported 500M - but the UNIQUE text was about the length of the
# longest session, because the short ones were all re-reading the opening
# of the same sequence.
#
# It showed up as a restart spike nobody could place. Train loss fell hard
# (0.675 -> 0.446 at the last one) because the model was being handed text
# it had already memorised, and held-out rose at the same moment (0.925 ->
# 1.105) because it was being pushed to overfit that slice further. The
# learning rate was identical across the seam, nothing was pruned, and
# nothing was added - which is what made it look like lost state rather
# than repeated data.
#
# Mixing the read count into the seed means a restart continues into text
# the model has not seen instead of starting the same tape again. It stays
# reproducible from a given checkpoint, and needs nothing persisted: the
# character count is already in the manifest.
return [_Lane(k, sorted(by[k]), np.random.default_rng([seed, j, resume]))
for j, k in enumerate(sorted(by))]
def cmd_read(args):
"""
Point the model at files and let it read them.
This is the whole interface for learning on your own data. The alphabet is
the 256 byte values, so a file needs no preparation - no tokenizer to fit,
no corpus to build, no step to re-run when the data changes. Give it a
directory and it reads what is there, in order, once or several times, and
is different afterwards.
python3 train.py read ~/notes
python3 train.py read src/ docs/ --passes 3
It is the same path `stream` uses - same chunking, same cache, same
gradient step. What differs is where the characters come from, and that a
file is read from its beginning to its end rather than sampled at random,
because a document has an order and the model should see it.
"""
# Retired spellings, checked before anything loads. `--dwell` is a strict
# prefix of `--dwell-chars` and argparse matches unambiguous prefixes, so
# it would be accepted silently and read as characters - an alias that
# keeps the name and changes the unit is the exact failure this change
# exists to remove.
if getattr(args, "dwell_retired", None) is not None:
raise SystemExit(
"--dwell counted SEGMENTS and is gone. Use --dwell-chars, in "
"characters: the old default of 4 is --dwell-chars 2048.")
# EVERY CADENCE IS CONFIGURED IN CHARACTERS AND USED IN STEPS.
#
# A step is `--chunk` characters, so anything counted in steps silently
# changes meaning the moment the chunk changes - `growth.every: 400` was
# 204,800 characters at chunk 512 and would have been 1.6 million at 4,096
# without a line of the config moving. Characters are what the model
# actually reads; steps are an implementation detail of how often we stop
# to update. Convert once, here, where the chunk is known.
def _in_steps(chars, least=1):
return max(least, int(round(int(chars) / max(1, args.chunk))))
grow_every_steps = _in_steps(args.grow_every)
survival_steps = _in_steps(args.prune_survival)
ctx_every_steps = _in_steps(args.context_every)
ctx_grow_every_steps = _in_steps(args.context_grow_every)
from minagi.ingest import collect, summarise
from minagi.stream import FolderEvaluator
from minagi.precision import set_compute_dtype
device = torch.device(args.device)
torch.manual_seed(args.seed)
set_compute_dtype(args.precision)
files = collect(args.paths)
if not files:
print(f"no readable text under {', '.join(args.paths)}", file=sys.stderr)
return 1
info = summarise(files)
print(f"{info['files']} files, {info['characters']/1e6:.2f}M characters"
+ (f", {args.passes} passes" if args.passes > 1 else ""))
wdir = args.weights_dir
# No model here yet is a model that has not been made yet, not an error.
# Everything about its shape is in config.yaml, so there is nothing to ask.
#
# The test is core.npz, NOT whether the directory exists. The repository
# ships weights/manifest.json - it records what the released model is - so
# weights/ is already there on a fresh clone, and testing the directory
# sent every new reader straight into a missing core.npz.
if not os.path.exists(os.path.join(wdir, "core.npz")):
from minagi.create import create
leftovers = ([f for f in os.listdir(wdir) if f != "manifest.json"]
if os.path.isdir(wdir) else [])
if leftovers:
# Something is in there, but not a model. Wiping it is not this
# function's decision to make.
raise SystemExit(
f"{wdir} has no core.npz but is not empty: "
f"{', '.join(sorted(leftovers)[:6])}. That is a partial or "
f"foreign weights directory - move it aside, or point "
f"--weights-dir somewhere else, and it will be created.")
if os.path.isdir(wdir):
print(f"{wdir} holds only a manifest - creating a fresh model "
f"from config.yaml (the manifest describes the released "
f"model, not this one, and is replaced)")
else:
print(f"{wdir} does not exist - creating a fresh model from "
f"config.yaml")
create(wdir, force=True)
print()
model, cfg, pool, man = build_paged(wdir, device, args.resident,
args.ram_capacity, args.context)
# the best held-out this run has seen, for the optional notifier below.
# Reads the manifest, so it cannot move above build_paged.
best_val = float(man.get("val") or float("inf"))
print(f" {pool.n_experts()} experts on disk, {pool.n_resident()} resident, "
f"{pool.vram_params()/1e6:.1f}M in VRAM of "
f"{pool.n_params()/1e6:.1f}M total")
print(f" computing in {args.precision}; weights stay fp32 everywhere, "
f"Adam's moments store as bf16")
# dwell is configured in characters and counted in segments, one per
# working-set decision, which is one per chunk
pool.margin = args.margin
pool.dwell = max(1, int(args.dwell) // max(1, args.chunk))
print(f" experts move only on demand: a candidate must be wanted "
f"{args.margin:.0%} more than the resident it displaces, and a "
f"newcomer holds its slot for {args.dwell:,} characters ({pool.dwell} chunks)")
if args.no_pool_checkpoint:
n_off = 0
for _s in model.modules():
if isinstance(_s, PooledMLP):
_s.grad_checkpoint = False
n_off += 1
print(f" pool activations kept, not recomputed ({n_off} call sites) "
f"- faster, and it needs the memory depth was using")
trunk, pool_ps = _split_trunk_pool(model)
tg = {"params": trunk, "name": "trunk", "weight_decay": args.wd,
"lr": args.lr * args.trunk_lr_mult,
"base_lr": args.lr * args.trunk_lr_mult}
pg = {"params": pool_ps, "name": "pool", "weight_decay": args.wd,
"lr": args.lr, "base_lr": args.lr}
opt = torch.optim.AdamW([tg, pg], lr=args.lr, betas=(0.9, 0.95),
fused=(device.type == "cuda"))
snr = GradSNR()
print(f" trunk learns at {args.trunk_lr_mult:g}x the pool's rate "
f"({sum(q.numel() for q in trunk)/1e6:.1f}M trunk, "
f"{sum(q.numel() for q in pool_ps)/1e6:.1f}M pool)")
# Scored at the window the model is actually reading at. A fixed
# reference window would be steadier, but it would stop describing the
# model as it grows; and because the window moves one character at a time,
# two consecutive scores are effectively at the same length anyway. The
# sample log states the window with every entry so no comparison is made
# blind.
# held-out is configured in CHARACTERS per domain and read in chunks
eval_steps = max(1, int(args.eval_chunks) // max(1, args.chunk))
sample_eval_steps = max(1, int(args.sample_eval_chunks)
// max(1, args.chunk))
ev = (FolderEvaluator(model, args.held_out, args.chunk, cfg.block, device,
segment_chunks=max(1, args.segment_chunks
// max(1, args.chunk)))
if args.held_out and os.path.isdir(args.held_out) else None)
before = None
if ev is not None:
v = ev.run(eval_steps); se = v.pop("stderr", 0.0)
before = float(np.mean(list(v.values())))
print(f" before: {before:.4f} +/-{se:.4f} on {args.held_out}")
pool.attach_optimiser(opt)
# ADAM'S MOMENTS. They are written to optim.npz at every checkpoint and
# were never read back here: this path builds a fresh optimiser and only
# ever SAVED it. So every restart threw away 108.9M moment values, 99.9% of
# them non-zero, and Adam re-estimated its second moments from a single
# gradient - which is a large, badly scaled step in a direction inferred
# from one batch, decaying over a few hundred steps. That is the held-out
# jump after a restart that then settles, on a model whose whole point is
# that it never stops training.
#
# The tell was in the file: Adam's step counter t read 2,381 at step
# 458,000. It counts optimiser steps since the last restart, not since the
# run began.
#
# The pool's w1/w2/w3 moments are slot-indexed and belong to whichever
# experts were resident when they were written, but restoring them is safe:
# swap_to overwrites a slot's moments from the expert's own file the moment
# it is paged in, and zeroes them if the file has none. `pool.gate` is
# per-expert and full width, so it restores exactly; a pool that changed
# size is caught by the shape guard in _load_optim.
try:
weights_store._load_optim(opt, model, wdir)
t_seen = max((float(st.get("step", 0))
for st in opt.state.values() if "step" in st), default=0)
print(f" optimiser moments restored ({t_seen:,.0f} steps of history)")
except Exception as e: # never let this stop a run
print(f" [note] optimiser moments not restored: {e}")
grower = AutoGrow(grow_k=args.grow_k, max_experts=10_000_000,
mem_frac_max=args.grow_mem_frac,
dying_frac_max=args.grow_dying_frac,
max_in_flight=args.max_in_flight,
keep_ratio_min=args.grow_keep_ratio,
max_disk_gb=args.max_disk_gb,
birth_gate=args.birth_gate,
recent_mult=args.recent_mult) \
if args.grow_k else None
# The checkpoint carries whatever depth policy it was trained under, but
# this is a knob about how to spend compute now, not a property of the
# weights - so config.yaml wins over the manifest every time.
cfg.train_steps_mean = float(args.train_steps_mean)
cfg.min_steps = max(1, min(int(args.min_steps), cfg.max_steps))
if args.bptt_window:
cfg.bptt_window = int(args.bptt_window)
if args.ponder_beta:
cfg.ponder_beta = float(args.ponder_beta)
if args.halt_prior:
cfg.halt_prior = float(args.halt_prior)
if args.halt_thresh:
cfg.halt_thresh = float(args.halt_thresh)
# The settings above are counted in ROWS, and the CHECKPOINT decides how
# many rows a step is - n_recur + n_coda. So a config written for one
# architecture is silently wrong on another, and the failure is not always
# loud: bptt_window 16 against max_steps 6 makes the detach condition
# n < -10, which never fires, so NOTHING is detached and the graph holds
# every block instead of the last few. That is how a config meant for a
# 24-row model put a 6-row model out of memory.
warn = []
if cfg.bptt_window > cfg.max_steps:
warn.append(f"bptt_window {cfg.bptt_window} exceeds max_steps "
f"{cfg.max_steps}, so nothing would be detached and the "
f"whole chain would be held. Clamping to {cfg.max_steps}.")
cfg.bptt_window = cfg.max_steps
if cfg.train_steps_mean and cfg.train_steps_mean >= cfg.max_steps:
warn.append(f"train_steps_mean {cfg.train_steps_mean:g} is at or above "
f"max_steps {cfg.max_steps}, so depth sampling saves "
f"nothing - every step runs the ceiling.")
if cfg.halt_prior and 1.0 / cfg.halt_prior > cfg.max_steps * 1.5:
warn.append(f"halt_prior {cfg.halt_prior:g} pulls toward "
f"{1/cfg.halt_prior:.1f} rows but the model only has "
f"{cfg.max_steps}.")
if warn:
print(f" [config] this checkpoint has {cfg.n_recur + cfg.n_coda} "
f"block(s) per row and {cfg.max_steps} rows; config.yaml looks "
f"written for a different shape:", flush=True)
for w in warn:
print(f" - {w}", flush=True)
print(f" gradient reaches back {cfg.bptt_window} of {cfg.max_steps} rows")
print(f" depth: up to {cfg.max_steps} rows, halting prior pulls toward "
f"{1/max(cfg.halt_prior,1e-9):.1f} rows, inference stops at "
f"{cfg.halt_thresh:g} cumulative")
if cfg.train_steps_mean > 0:
import torch as _t
_n = (_t.poisson(_t.full((20000,), cfg.train_steps_mean)).int() + 1
).clamp(cfg.min_steps, cfg.max_steps).float().mean()
print(f" recurrence depth is SAMPLED while training: mean {_n:.2f} of "
f"{cfg.max_steps} ({100*_n/cfg.max_steps:.0f}% of the recurrent "
f"compute). Inference still runs all {cfg.max_steps} and lets "
f"halting decide.")
pool.dying_at = args.dying_at
_PLOTS.update(on=bool(args.plots), since=args.plot_since,
log=args.sample_log, weights=args.weights_dir)
if args.plots:
print(f" graphs redrawn with every sample entry -> "
f"runs/training_progress.png"
+ (f" and runs/dashboard.png (the last "
f"{args.plot_since:g}M)" if args.plot_since else ""))
pool.explore = args.explore
# A new expert competes without the handicap until it is old enough to be
# pruned, so the moment it is judged is the moment its fair turn ends.
# One number governs both, which is the point: there is no second window
# to reason about.
pool.trial = survival_steps
from minagi.tokenizer import ByteTokenizer
tok = ByteTokenizer()
model.train()
t0 = time.time()
last_sample = time.time()
# Steps and characters are cumulative over the model's whole life, not
# per session. `born[i]` records the step an expert was grown and prune
# asks whether `step - born[i]` has passed min_age, so restarting the
# count at zero made that difference negative and quietly exempted every
# expert grown in an earlier session from ever being pruned.
step = int(man.get("step", 0) or 0)
# The clock the expert trial is read against. Set only inside the loop it
# is still 0 for the first forward of a session, and `now - born` is then
# hugely NEGATIVE for every expert grown in an earlier session - so all of
# them pass `< trial` at once and the whole pool gets the newborn's waiver,
# routing on key fit alone with the gate ignored. One batch, but it is the
# first batch after every restart.
pool.now = step
base_chars = int(man.get("read_chars", 0) or 0)
seen = 0
did = 0 # steps taken in THIS session
losses = []
if args.sample_every:
# Append. A resumed run continues the same model, so wiping the file
# would throw away the only record of how it got here; the banner is
# what separates one session from the next.
with open(args.sample_log, "a") as f:
f.write(f"\n\n{'#' * 78}\n")
f.write(f"# session started {time.strftime('%Y-%m-%d %H:%M')} "
f"at step {step:,}, {base_chars/1e6:.1f}M of "
f"{info['characters']/1e6:,.0f}M characters read "
f"({100.0 * base_chars / max(info['characters'], 1):.2f}%), "
f"{pool.n_experts()} experts\n")
f.write(f"{'#' * 78}\n")
# The window starts where the model left it, not at the ceiling. RoPE
# carries no learned parameters, so the same weights read at any length;
# what a model cannot do is jump to a length it has never seen, which is
# why this moves a character at a time.
ctx_now = int(man.get("context_now") or args.context_start)
# --context is a CEILING that binds on resume, not only on a fresh model.
# It used to be clamped to cfg.block alone, and both cfg.block and
# context_now come out of the checkpoint - so lowering model.context in
# config.yaml changed nothing at all on a resumed run. Measured the hard
# way: two runs at 16,384 and 8,192 in the config both read 32,768 from
# the manifest and died with byte-identical out-of-memory errors.
#
# The window is the one setting that has to be lowerable from outside,
# because it is what decides whether the model fits on this card, and the
# card is not a property of the checkpoint.
# ONE ceiling, used everywhere. cfg.block is how far the rotary tables
# were built; --context is how far this card can afford. Clamping only at
# startup was worse than not clamping at all: the window came down to 4,096
# and the growth path below, which clamped to cfg.block alone, immediately
# started walking it back toward 32,768 a character at a time. Measured -
# context 4,151 -> 4,492 over two hours while VRAM went 7,271 -> 7,790 MB,
# and it died on the way. That looks exactly like a leak from the outside,
# which is what it was reported as.
ctx_max = max(64, min(int(cfg.block), int(args.context)))
ctx_now = max(64, min(ctx_now, ctx_max))
if ctx_now < int(man.get("context_now") or 0):
print(f" window held at {ctx_now:,} by --context-end "
f"(the checkpoint had grown to "
f"{int(man.get('context_now')):,})")
# loss by distance into the window: the chunk index IS the distance, so
# this costs nothing to collect
by_pos = collections.defaultdict(list)
def context_gain():
"""
How much better the model predicts deep into the window than early in
it. Positive means it is using the distance it has - and would use
more. Near zero means the far end of the window is being carried for
nothing.
"""
n_buckets = max(2, ctx_now // args.chunk)
early = [l for p, ls in by_pos.items() if p < n_buckets * 0.5
for l in ls]
late = [l for p, ls in by_pos.items() if p >= n_buckets * 0.75
for l in ls]
if len(early) < 8 or len(late) < 8:
return None
return float(np.mean(early) - np.mean(late))
# One lane per subject, rotated a window at a time.
lanes = _lanes(files, args.shuffle_seed, args.paths,
resume=int(man.get("read_chars", 0) or 0))
# How long a subject is read for before the next one. It is a whole
# number of context windows - a window never spans two subjects, because
# attention across the seam where chess becomes Python teaches nothing -
# but a subject may take several windows in a row.
#
# The floor matters at small windows. Tied to the window alone, a 2,048
# context would change subject every four chunks, and the first chunk
# after a change picks its experts from the states of the subject just
# left. At four chunks a visit that is a quarter of all reading.
def visit_chunks(ctx):
"""How many chunks make up one visit to a file.
`--passage` is in characters and says how far the reader walks through
one file before moving to another subject. The window slides along it
a chunk at a time, so a visit is passage/chunk steps - and it must be
at least one window, or the window never fills and the model is asked
to read with less context than it has.
This used to be `min_visit_chunks`, in chunks, rounded up to whole
context windows - so the number in the config was not the number in
effect. At 16 with a 32,768 window it silently became 64.
"""
least = max(1, -(-ctx // args.chunk)) # one full window
want = max(1, -(-args.passage // args.chunk))
return max(least, want)
turn = visit_chunks(ctx_now)
wants_more = False # whether the evidence says to keep growing
target = info["characters"] * max(1, args.passes)
print(f" {len(lanes)} subjects, reading {turn * args.chunk:,} characters "
f"from one before moving to the next"
+ (f" ({turn * args.chunk // ctx_now} windows of {ctx_now:,})"
if turn * args.chunk > ctx_now else "")
+ ": " + ", ".join(l.name for l in lanes))
if ctx_now < ctx_max:
print(f" the window grows by {args.context_step} character(s) "
f"whenever the model is still gaining more than "
f"{args.context_gain_min} deep into it, up to {ctx_max:,}")
if base_chars:
print(f" {base_chars/1e6:.1f}M characters read in earlier sessions")
def log_history(val=None):
"""
Append what the pool looks like now, so a run leaves a trail.
The manifest only ever holds the latest state, which cannot answer
whether an expert earned its gate early and coasted or climbed all
along. One line per checkpoint, and a run of any length stays
readable afterwards.
"""
if not args.history:
return
try:
t = pool.telemetry()
row = {"t": time.time(), "step": step,
"chars": base_chars + seen, "val": val,
"experts": pool.n_experts(), "segments": t["segments"],
"gate": t["gate"], "use": t["use"], "admits": t["admits"],
"since": t["since"], "born": t["born"],
"last_seen": t.get("last_seen"), "uid": t.get("uid"),
"trial": getattr(model.pool, "trial", 0)}
os.makedirs(os.path.dirname(args.history) or ".", exist_ok=True)
with open(args.history, "a") as f:
f.write(json.dumps(row, separators=(",", ":")) + "\n")
except Exception as e: # telemetry must never stop a run
print(f" (history not written: {e})", flush=True)
def checkpoint(val=None):
# There is no reading position to record any more: windows are drawn
# at random, so what carries forward is how much has been read, not
# where the reader had got to.
weights_store.save(model, wdir, step=step, val=val, opt=opt,
cfg=asdict(cfg),
extra={"read_chars": base_chars + seen,
"read_nats": nats,
"plasticity": plast.state(),
"context_now": ctx_now})
log_history(val)
stop = False
swapped = 0
recent = collections.deque(maxlen=400) # the training loss lately
last_gap = None # held-out minus train, when known
# How much information has been read, as against how many characters were
# passed over: the loss IS the information content of what was read, so
# integrating it says what the model actually had to account for. A
# million characters of chess carries 0.744M nats and a million of
# wikipedia 1.686M - the same reading, worth 2.27x as much.
#
# Recorded because it is the honest way to describe a run, and NOT used to
# decide anything. Growth is judged on the network as it stands.
# The learning rate is governed by held-out loss, not by a horizon.
# There is no lr_origin any more and nothing counts characters: see
# minagi/plasticity.py for the two rules and why they are symmetric.
plast = Plasticity.restore((man or {}).get("plasticity"))
if args.lr_reset:
plast.scale = 1.0
print(" plasticity RESET: full learning rate restored", flush=True)
print(f" {plast.describe()}", flush=True)
nats = float(man.get("read_nats") or 0.0) if man else 0.0
if not nats and base_chars:
# A checkpoint written before this was tracked has the characters but
# not the information. Starting the counter at zero would make the
# ratio enormous and silently re-block growth for hours after every
# upgrade, so seed it from what those characters were most likely
# worth: the held-out score, which is the model's own estimate of nats
# per character on text like the text it read.
View remainder of file in raw view
0