Repository navigation
Expand file tree
/
Copy pathtrain.py
More file actions
2322 lines (2177 loc) · 118 KB
/
Copy pathtrain.py
File metadata and controls
2322 lines (2177 loc) · 118 KB
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
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
#!/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
pr
8BD0
eceded 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.
# PyTorch 2.9 renamed the variable and warns about the old name; earlier
# versions read only the old one.
def _alloc_var():
try:
from importlib.metadata import version
major, minor = (int(v) for v in version("torch").split(".")[:2])
return "PYTORCH_ALLOC_CONF" if (major, minor) >= (2, 9) else "PYTORCH_CUDA_ALLOC_CONF"
except Exception: # noqa: BLE001
return "PYTORCH_CUDA_ALLOC_CONF"
if not any(v in os.environ for v in ("PYTORCH_ALLOC_CONF", "PYTORCH_CUDA_ALLOC_CONF")):
os.environ[_alloc_var()] = "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 admitted experts. 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):
"""
One line per training step: how many experts its window had to load.
`first` marks the chunk that opens a lane's window - the only place
the subject actually changes, so the place loads should concentrate.
"""
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 _truncate_history(path, chars):
"""Drop history rows past `chars` - they belong to a branch that was
abandoned, and every reader of this file assumes it moves forward."""
if not path or not os.path.exists(path) or chars <= 0:
return
try:
with open(path) as f:
rows = f.read().splitlines()
keep, dropped = [], 0
for line in rows:
try:
if json.loads(line).get("chars", 0) > chars:
dropped += 1
continue
except json.JSONDecodeError:
pass # keep anything unparseable
keep.append(line)
if dropped:
with open(path, "w") as f:
f.write("\n".join(keep) + ("\n" if keep else ""))
print(f" history: dropped {dropped} row(s) past "
f"{chars/1e6:.1f}M - superseded by this resume", flush=True)
except OSError as e:
print(f" (history not trimmed: {e})", flush=True)
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.
"""
# 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")
from minagi.precision import describe
print(f" computing in {describe(device)}; weights stay fp32 everywhere, "
f"Adam's moments store as bf16")
# 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. Applied
# before the first held-out check, which the halting settings change.
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)
cfg.halt_freeze = bool(args.halt_freeze)
# 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). "
+ (f"A character stops when it halts, in training and in "
f"writing alike." if cfg.halt_freeze else
f"Inference still runs all {cfg.max_steps} and lets halting "
f"decide which row each character keeps."))
_temp = float(getattr(pool, "select_temperature", 0.0) or 0.0)
if _temp > 0:
print(f" experts are DRAWN at selection temperature {_temp:g}: each row "
f"draws {cfg.pool_top_k} from the text's probabilities until the "
f"{pool.n_resident()} slots are full, and every character draws "
f"its {cfg.pool_top_k} from its own - the text itself is always "
f"read and written greedily")
else:
print(f" every character ranks all {pool.n_experts()} experts and asks "
f"for its top {cfg.pool_top_k}; each forward uses the "
f"{pool.n_resident()} its text asks for most")
if getattr(pool, "balance", 0) > 0:
print(f" while training, the router pays for the probability it puts "
f"on each expert in proportion to that expert's share of recent "
f"admissions (balance {pool.balance:g}), so the whole pool gets "
f"used")
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)
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 - not to the ones the
# held-out check above has just put on the card. What the restore brings
# for them is their step counter; before the first step every expert on
# the card is given its own moments from its own file
# (PagedPool._own_moments), over whatever the restore wrote in its slot.
# `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
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 ""))
# A newborn is safe from pruning for the survival window - the same
# number prune deletes on, so there is one 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))
# A RESUME SUPERSEDES WHATEVER THE OLD BRANCH WROTE. The history log is
# append-only and its character counter is not monotonic across a restart:
# reading on from a checkpoint leaves rows describing a model that no
# longer exists, sitting after the point being resumed from. Every reader
# then has to cope - the dashboard drew a seam, window_use reached back to
# the wrong row - and the file needs cleaning by hand every time an
# experiment is abandoned.
#
# So drop them here, once, at the point the run knows where it is
# resuming. Anything already past that mark describes a branch this run is
# not on.
_truncate_history(args.history, base_chars)
# 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,