Skip to content

Commit 9bf71c9

Browse files
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

File tree

0 commit comments

Comments
 (0)