ATOM (atom-ai): code snapshot (100fae0 + batched-train), README with measured numbers, curated checkpoints (part 4)
b9301f1 verified | # Controlled A/B driver for the slow-mode readout (docs/SLOW_MODE_READOUT.md), derived | |
| # from scripts/train_leakspec_arm.sh (same recipe, seed, stream skip, chunking). | |
| # Extra per-chunk measurements: 2048-transition val (eval_windows.py) and the slow-mode | |
| # probes (probe_leak_spectrum.py --slow-probes). Artifacts: docs/artifacts/slow_mode_readout. | |
| # | |
| # One ARM = the exact fieldsat recipe (scripts/train_fieldsat_branch.sh: live recipe + | |
| # --last-atom-dropout 0.5 + LR 1.5e-5 / SLR 3.6e-5, seed 20260913, stream-skip = START-496000) | |
| # resumed from START_CKPT, in CHUNK-step chunks up to END_TARGET, with its own output dir and | |
| # checkpoint suffix. The ONLY difference between arms is LEAK_FLAGS: | |
| # ARM=A LEAK_FLAGS="--field-leak 0.3" (control) | |
| # ARM=B LEAK_FLAGS="--field-leak-spectrum ... --slow-mode-readout --slow-memory-loss W ..." | |
| # Resumable: restarts from the newest ckpt of this arm (<= END_TARGET), else START_CKPT. | |
| # After each chunk: val normal / field-alone on the fixed 4876k ring, probe_leak_spectrum | |
| # (cap fraction, a_hat linear next-byte CE, prompt-memory @16/32/64), coherence chat | |
| # (3 prompts, standard params), seeded chat (seeds 1-3 x T 0.8/0.6/1.0) + prompt-dependence; | |
| # output-dependence (JS divergence of next-byte distributions across prompts at the reply | |
| # start); one row per chunk in $TSV (docs/artifacts/field_leak_spectrum/tsv_row.py). | |
| # | |
| # Env: ARM OUT SUF LEAK_FLAGS TSV [START_CKPT START_STEP END_TARGET CHUNK THREADS MEASURE_THREADS LR SLR LADROP] | |
| # THREADS=8 OMP/MKL threads for training | |
| # MEASURE_THREADS=1 threads for the (sequential, tiny-op) measurement tools; 8-thread OMP | |
| # measurement next to an 8-thread training run thrashes (seen: chat_seeds | |
| # 35 CPU-min, training 0.3 steps/s), and thread count never touches training. | |
| set -u | |
| cd /workspace/exports/atom-ai | |
| ARM=${ARM:?ARM} | |
| OUT=${OUT:?OUT} | |
| SUF=${SUF:?SUF} | |
| LEAK_FLAGS=${LEAK_FLAGS:?LEAK_FLAGS} | |
| TSV=${TSV:?TSV} | |
| START_CKPT=${START_CKPT:-checkpoints/byte_tick_fieldsat/atom_native_step_4996000_d32_ms1M_lrd_b4828_fieldsat.pt} | |
| START_STEP=${START_STEP:-4996000} | |
| END_TARGET=${END_TARGET:-5004000} | |
| CHUNK=${CHUNK:-8000} | |
| THREADS=${THREADS:-8} | |
| MEASURE_THREADS=${MEASURE_THREADS:-1} | |
| LR=${LR:-1.5e-5} | |
| SLR=${SLR:-3.6e-5} | |
| LADROP=${LADROP:-0.5} | |
| BASE=496000 | |
| A=docs/artifacts/slow_mode_readout | |
| AB=$A/arm_$ARM | |
| LOGD=logs/slowread/arm_$ARM | |
| RING=checkpoints/byte_tick_fieldsat/val_ring_4876000.pt | |
| PY=.venv/bin/python | |
| export PYTHONPATH=. OMP_NUM_THREADS=$THREADS MKL_NUM_THREADS=$THREADS OPENBLAS_NUM_THREADS=$THREADS | |
| mkdir -p "$OUT" "$AB" "$LOGD" | |
| ts() { date -Is; } | |
| exec 9>"$LOGD/lock" | |
| if ! flock -n 9; then echo "$(ts) arm $ARM already running; exiting"; exit 1; fi | |
| echo "=== $(ts) slowread arm=$ARM OUT=$OUT SUF=$SUF LEAK_FLAGS=[$LEAK_FLAGS] START=$START_STEP END=$END_TARGET THREADS=$THREADS ===" | |
| latest() { | |
| ls "$OUT"/atom_native_step_*_${SUF}.pt 2>/dev/null | sed -n 's/.*step_\([0-9]*\)_.*/\1/p' | sort -n \ | |
| | awk -v lo="$START_STEP" -v hi="$END_TARGET" '$1>lo && $1<=hi' | tail -1 | |
| } | |
| has_row() { [ -f "$TSV" ] && awk -F'\t' -v s="$1" -v a="$ARM" 'NR>1 && $1==a && $2==s {f=1} END {exit !f}' "$TSV"; } | |
| measure() { | |
| local END=$1 CK="$OUT/atom_native_step_$1_${SUF}.pt" ML="$LOGD/measure_$1.log" | |
| echo "$(ts) measure $ARM $END start (threads=$MEASURE_THREADS)" | |
| local OMP_NUM_THREADS=$MEASURE_THREADS MKL_NUM_THREADS=$MEASURE_THREADS OPENBLAS_NUM_THREADS=$MEASURE_THREADS | |
| export OMP_NUM_THREADS MKL_NUM_THREADS OPENBLAS_NUM_THREADS | |
| "$PY" -u tools/eval_last_atom_ablation.py --end-steps 4876000 --ring-cache "$RING" \ | |
| --checkpoints "$CK" --out "$AB/val_ablation_$END.json" >> "$ML" 2>&1 | |
| "$PY" -u docs/artifacts/field_leak_spectrum/eval_windows.py "$AB/eval_windows_$END.json" "$CK" >> "$ML" 2>&1 | |
| "$PY" -u tools/probe_leak_spectrum.py --checkpoint "$CK" --variants ckpt --slow-probes \ | |
| --out "$AB/probe_$END.json" >> "$ML" 2>&1 | |
| local G="$AB/gen_coherent_$END.txt"; : > "$G" | |
| for p in "Bonjour, comment ça va ?" "Qui es-tu ?" "Il était une fois"; do | |
| RESP=$("$PY" tools/chat_atom_native.py --checkpoint "$CK" --prompt "$p" --no-role-prime \ | |
| --max-packets 96 --max-length 100 --temperature 0.8 --top-k 12 2>/dev/null | sed -n 's/^Response: //p' | head -1) | |
| echo "PROMPT: $p" >> "$G"; echo "RESP: $RESP" >> "$G" | |
| done | |
| "$PY" tools/chat_seeds_probe.py --checkpoint "$CK" --seeds 1 2 3 --temps 0.8 0.6 1.0 \ | |
| --out "$AB/chat_seeds_$END.json" > "$AB/chat_seeds_$END.txt" 2>&1 | |
| "$PY" -u tools/probe_leak_spectrum.py --checkpoint "$CK" --variants ckpt --skip-episode --skip-memory \ | |
| --output-dependence --out "$AB/outdep_$END.json" >> "$ML" 2>&1 | |
| "$PY" "$A/tsv_row.py" "$ARM" "$END" "$CK" "$TSV" "$AB" "$OUT/run_metrics_${END}_${SUF}.json" | |
| echo "$(ts) measure $ARM $END done" | |
| } | |
| recover_metrics() { | |
| # ckpt produced by a run whose driver shell was stopped: recover its metrics files | |
| # (run_metrics.json in $OUT belongs to this ckpt iff its checkpoint path matches). | |
| local END=$1 NAME="atom_native_step_$1_${SUF}.pt" | |
| if [ ! -f "$AB/run_metrics_${END}.json" ] && grep -q "$NAME" "$OUT/run_metrics.json" 2>/dev/null; then | |
| cp "$OUT/run_metrics.json" "$OUT/run_metrics_${END}_${SUF}.json" | |
| cp "$OUT/run_metrics.json" "$AB/run_metrics_${END}.json" | |
| cp "$OUT/loss_trajectory.json" "$AB/loss_trajectory_${END}.json" 2>/dev/null || true | |
| echo "$(ts) recovered metrics for existing $NAME" | |
| fi | |
| } | |
| CUR=$(latest); [ -z "$CUR" ] && CUR=$START_STEP | |
| while [ "$CUR" -lt "$END_TARGET" ]; do | |
| END=$((CUR + CHUNK)); [ "$END" -gt "$END_TARGET" ] && END=$END_TARGET | |
| STEPS=$((END - CUR)); SKIP=$((CUR - BASE)) | |
| if [ "$CUR" -eq "$START_STEP" ]; then RESUME="$START_CKPT"; else RESUME="$OUT/atom_native_step_${CUR}_${SUF}.pt"; fi | |
| NAME="atom_native_step_${END}_${SUF}.pt" | |
| LOG="$LOGD/train_${CUR}_${END}.log" | |
| if [ ! -f "$OUT/$NAME" ]; then | |
| echo "$(ts) chunk $CUR->$END skip=$SKIP resume=$RESUME log=$LOG" | |
| T0=$(date +%s) | |
| # shellcheck disable=SC2086 | |
| "$PY" -u tools/run_atom_native.py \ | |
| --stream --data-glob 'data/corpus_fr_medium*.txt' --loop-shards --chunk-bytes 65536 \ | |
| --stream-skip-packets "$SKIP" --resume "$RESUME" --output-dir "$OUT" \ | |
| --checkpoint-name "$NAME" --steps "$STEPS" \ | |
| --d-model 32 --n-modes 32 --n-atoms-max 64 --max-span-bytes 1 \ | |
| --field-obligatory-hard --no-enable-merge --no-payload-copy --last-atom-readout \ | |
| --efference-every 10 --learning-rate "$LR" --surface-learning-rate "$SLR" --seed 20260913 \ | |
| --atom-flush-every 64 --episode-length 8000 --episode-reset --slow-every 1 \ | |
| --field-loss-weight 0.05 --field-contrast-weight 0 --allow-parallel-train \ | |
| --last-atom-dropout "$LADROP" $LEAK_FLAGS \ | |
| > "$LOG" 2>&1 | |
| RC=$? | |
| if [ "$RC" -ne 0 ] || [ ! -f "$OUT/$NAME" ]; then echo "$(ts) chunk FAILED rc=$RC (see $LOG)"; exit 1; fi | |
| cp "$OUT/run_metrics.json" "$OUT/run_metrics_${END}_${SUF}.json" | |
| cp "$OUT/run_metrics.json" "$AB/run_metrics_${END}.json" | |
| cp "$OUT/loss_trajectory.json" "$AB/loss_trajectory_${END}.json" 2>/dev/null || true | |
| echo "$(ts) chunk $CUR->$END saved $NAME wall=$(( $(date +%s) - T0 ))s" | |
| else | |
| recover_metrics "$END" | |
| fi | |
| has_row "$END" || measure "$END" | |
| CUR=$END | |
| done | |
| # measure any checkpoint of this arm in (START_STEP, END_TARGET] that has no TSV row yet | |
| for E in $(ls "$OUT"/atom_native_step_*_${SUF}.pt 2>/dev/null | sed -n 's/.*step_\([0-9]*\)_.*/\1/p' | sort -n); do | |
| if [ "$E" -gt "$START_STEP" ] && [ "$E" -le "$END_TARGET" ] && ! has_row "$E"; then | |
| recover_metrics "$E"; measure "$E" | |
| fi | |
| done | |
| echo "$(ts) DONE arm $ARM at $CUR" | |