File size: 2,698 Bytes
4397e12 | 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 | #!/bin/bash
# LR check (S, 20 min) -> Engram A/B (L, 60 min, equal wall-clock) -> 23h main run with the winners.
# Runs in the background; every decision is printed to $TA_DATA/logs/pipeline.log.
set -euo pipefail
cd "$(dirname "$0")/.."
source env.sh
R=$TA_DATA/runs
mkdir -p $R
log() { echo "[$(date '+%F %T')] $*" | tee -a $TA_DATA/logs/pipeline.log; }
common="--tokens 1e10 --warmup 50 --eval_every 50 --ckpt_minutes 999 --snapshot_tokens 1e12"
log "waiting for synth_grounded (part of the stable mixture)"
until ls $TA_DATA/tok/train/synth_grounded/*.bin >/dev/null 2>&1; do sleep 30; done
log "phase 1: LR check at size S (3 concurrent, 20 min)"
for lr in 1.5e-3 3e-3 6e-3; do
[ -f $R/lr_$lr/ckpt.pt ] || $TA_PY scripts/train.py --size S --engram 0 --lr $lr --batch_tokens 131072 \
--stop_minutes 20 $common --out $R/lr_$lr > $R/lr_$lr.out 2>&1 &
done
wait
$TA_PY scripts/compare_runs.py $R/lr_1.5e-3 $R/lr_3e-3 $R/lr_6e-3 | tee -a $TA_DATA/logs/pipeline.log
LR=$($TA_PY - <<EOF
import json
best=None
for lr in ["1.5e-3","3e-3","6e-3"]:
recs=[json.loads(l) for l in open("$R/lr_"+lr+"/log.jsonl") if '"val"' in l]
v=recs[-1]["val"]; m=sum(v.values())/len(v)
if best is None or m<best[0]: best=(m,lr)
print(best[1])
EOF
)
log "best short-run lr=$LR"
if [ -n "${LR_OVERRIDE:-}" ]; then LR=$LR_OVERRIDE; log "LR_OVERRIDE -> lr=$LR"; fi
log "phase 2: Engram A/B at size L (2 concurrent, 60 min, equal wall-clock)"
for e in 0 1; do
[ -f $R/ab_engram$e/ckpt.pt ] || $TA_PY scripts/train.py --size L --engram $e --lr $LR --warmup 100 \
--stop_minutes 60 --tokens 1e10 --eval_every 50 --ckpt_minutes 999 --snapshot_tokens 1e12 \
--out $R/ab_engram$e > $R/ab_engram$e.out 2>&1 &
done
wait
$TA_PY scripts/compare_runs.py $R/ab_engram0 $R/ab_engram1 | tee -a $TA_DATA/logs/pipeline.log
ENGRAM=$($TA_PY - <<EOF
import json
def last(d):
recs=[json.loads(l) for l in open(d+"/log.jsonl") if '"val"' in l]
v=recs[-1]["val"]; return sum(v.values())/len(v)
# equal wall-clock: whichever reached the lower val loss in the same 60 minutes wins
print(1 if last("$R/ab_engram1") < last("$R/ab_engram0") else 0)
EOF
)
log "picked engram=$ENGRAM"
log "waiting for synthetic sources"
n_parts() { ls $TA_DATA/tok/train/$1/*.bin 2>/dev/null | wc -l; }
until [ "$(n_parts synth_grounded)" -ge 8 ] && [ "$(n_parts synth_reasoning)" -ge 8 ]; do sleep 30; done
log "phase 3: main run, size L, engram=$ENGRAM, lr=$LR, 23h budget with automatic decay, decay at T=4096"
$TA_PY scripts/train.py --size L --engram $ENGRAM --lr $LR --tokens 1e10 --max_minutes 1380 --decay_T 4096 \
--eval_every 250 --snapshot_tokens 5e8 --out $R/main > $R/main.out 2>&1
log "main run finished"
|