Commit 9bf71c9
authored
Stream probe training to eliminate hidden-state corpus retention (#57)
The previous probe pipeline extracted all-layer hidden states for every
train/val position up front (~fp32 CPU tensors, one per layer) before
running any probe training. On a 512-ctx base model with a 4096-game
train set this came to ~33 GB of CPU RAM, OOMing 27 GB local systems
well before training began.
Rewrite train_all_probes as a streaming trainer: for each epoch, iterate
over game batches, run one forward pass, train every probe on the
batch's valid positions, then discard activations. Validation follows
the same pattern and accumulates per-layer loss / accuracy / R² / MAE in
a single sweep. Peak RSS for the same 1024/256-game run drops from 26+
GB (OOM) to ~8 GB.
Also compact the probe-data cache: boards now int8 (was int64), side/is-
check bool, ep_square int8, halfmove_clock uint8, legal_move_counts
uint16. Promotion happens at use sites in get_probe_targets.
scripts/eval_probes.py: bound the probe-data cache to exactly one
(max_ply, prepend_outcome) entry, and load checkpoints on CPU before
moving the final model to the target device to avoid a transient 2x
VRAM peak during load.
The legacy helpers (_extract_hidden_states, _extract_all_hidden_states,
_train_probe_all_layers, train_single_probe) are intact for tests; they
are no longer called by the main pipeline.1 parent 40b7195 commit 9bf71c9
2 files changed
Lines changed: 411 additions & 110 deletions
0 commit comments