ATOM (atom-ai): code snapshot (100fae0 + batched-train), README with measured numbers, curated checkpoints (part 4)
b9301f1 verified | # Resumable driver for the field-leak branch (docs/FIELD_SATURATION.md). | |
| # | |
| # Continues the fieldsat branch in 8k chunks with the exact live recipe from | |
| # scripts/train_until_coherent.sh + --last-atom-dropout 0.5 + --field-leak 0.3. | |
| # Each chunk saves checkpoints/byte_tick_fieldsat/atom_native_step_<END>_d32_ms1M_lrd_b4828_fieldsat.pt, | |
| # so the run can be resumed or switched at any chunk boundary. Re-running the | |
| # script resumes from the newest existing branch checkpoint in [START_STEP, END_TARGET] | |
| # (an interrupted chunk is simply re-run) and measures any checkpoint that has no | |
| # row in docs/artifacts/field_saturation/branch_progress.tsv yet. | |
| # | |
| # Env: | |
| # THREADS=1 OMP/MKL threads for every python step (8 only when the live loop is stopped) | |
| # YIELD=1 1 = SIGSTOP this whole process group while the live b4828 chunk is in its | |
| # multi-threaded phase (docs/artifacts/field_saturation/yield_to_live.py); | |
| # requires launching via setsid (own process group). 0 = never pause. | |
| # START_STEP=4876000 END_TARGET=4916000 CHUNK=8000 | |
| # LR=1.5e-5 SLR=3.6e-5 LEAK=0.3 LADROP=0.5 | |
| # | |
| # Launch (background, own process group): | |
| # setsid nohup env THREADS=1 YIELD=1 bash scripts/train_fieldsat_branch.sh \ | |
| # >> logs/fieldsat_branch_master.log 2>&1 < /dev/null & | |
| # Stop cleanly (group may be SIGSTOPped, so CONT after TERM): | |
| # PG=$(cat logs/fieldsat/branch.pgid); kill -TERM -- -$PG; kill -CONT -- -$PG | |
| set -u | |
| cd /workspace/exports/atom-ai | |
| THREADS=${THREADS:-1} | |
| YIELD=${YIELD:-1} | |
| START_STEP=${START_STEP:-4876000} | |
| END_TARGET=${END_TARGET:-4916000} | |
| CHUNK=${CHUNK:-8000} | |
| LR=${LR:-1.5e-5} | |
| SLR=${SLR:-3.6e-5} | |
| LEAK=${LEAK:-0.3} | |
| LADROP=${LADROP:-0.5} | |
| BASE=496000 | |
| OUT=checkpoints/byte_tick_fieldsat | |
| SUF=d32_ms1M_lrd_b4828_fieldsat | |
| A=docs/artifacts/field_saturation | |
| AB=$A/branch | |
| TSV=$A/branch_progress.tsv | |
| LOGD=logs/fieldsat/branch | |
| RING=$OUT/val_ring_4876000.pt # fixed 256-transition stream-ring val (same window as the probe table) | |
| PY=.venv/bin/python | |
| export PYTHONPATH=. OMP_NUM_THREADS=$THREADS MKL_NUM_THREADS=$THREADS OPENBLAS_NUM_THREADS=$THREADS | |
| mkdir -p "$OUT" "$AB" "$LOGD" logs/fieldsat | |
| ts() { date -Is; } | |
| exec 9>logs/fieldsat/branch.lock | |
| if ! flock -n 9; then echo "$(ts) another train_fieldsat_branch.sh holds logs/fieldsat/branch.lock; exiting"; exit 1; fi | |
| PG=$(ps -o pgid= -p $$ | tr -d ' ') | |
| echo "$PG" > logs/fieldsat/branch.pgid | |
| echo "=== $(ts) fieldsat branch driver pid=$$ pgid=$PG THREADS=$THREADS YIELD=$YIELD START_STEP=$START_STEP END_TARGET=$END_TARGET CHUNK=$CHUNK LR=$LR SLR=$SLR LEAK=$LEAK LADROP=$LADROP ===" | |
| if [ "$YIELD" = "1" ]; then | |
| if [ "$PG" != "$$" ]; then | |
| echo "$(ts) YIELD=1 needs its own process group: launch with setsid (see header). exiting"; exit 1 | |
| fi | |
| # Helper in its own session (so it never stops itself); pauses/resumes only our pgid. | |
| setsid "$PY" "$A/yield_to_live.py" logs/fieldsat/branch.pgid 9>&- \ | |
| >> logs/fieldsat/branch_yield.err 2>&1 < /dev/null & | |
| echo "$(ts) yield helper pid=$! watching pgid=$PG" | |
| fi | |
| 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" 'NR>1 && $1==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 $END start" | |
| "$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 tools/diag_field_saturation.py --checkpoints "b$END=$CK" --episode-ticks 8000 \ | |
| --prefixes 3000 --probe --out "$AB/diag_$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" - "$END" "$CK" "$TSV" "$AB" "$OUT/run_metrics_${END}_${SUF}.json" <<'PYEOF' | |
| import json, os, sys | |
| end, ck, tsv, ab, rm_path = sys.argv[1:6] | |
| def load(p): | |
| try: | |
| return json.load(open(p)) | |
| except Exception: | |
| return {} | |
| rm = load(rm_path); tr = rm.get("train", {}); va = rm.get("validation", {}) | |
| ev = load(f"{ab}/val_ablation_{end}.json") | |
| ring = next(iter(ev.get("rings", {}).values()), {}).get("checkpoints", {}).get(ck, {}) | |
| dg = (load(f"{ab}/diag_{end}.json").get("rows") or [{}])[0] | |
| fork = (dg.get("shared_prefix_forks") or {}).get("3000", {}) | |
| chats = [] | |
| try: | |
| for line in open(f"{ab}/gen_coherent_{end}.txt", encoding="utf-8"): | |
| if line.startswith("RESP: "): | |
| chats.append(" ".join(line[6:].split())[:80].replace("\t", " ")) | |
| except FileNotFoundError: | |
| pass | |
| def f(x, n=4): | |
| return "" if x is None or x == "" else (f"{x:.{n}f}" if isinstance(x, (int, float)) else str(x)) | |
| cols = [ | |
| ("end_step", end), | |
| ("chunk_val_loss_own_ring", f(va.get("loss"))), | |
| ("val_loss_fixed_ring", f(ring.get("normal", {}).get("loss"))), | |
| ("val_byte_ce_fixed_ring", f(ring.get("normal", {}).get("byte_loss"))), | |
| ("field_alone_byte_ce", f(ring.get("la_zeroed", {}).get("byte_loss"))), | |
| ("a_hat_linear_next_byte_ce", f((dg.get("probe_next_byte_from_a_hat") or {}).get("test_ce"), 3)), | |
| ("train_frac_at_cap", f(tr.get("field_limited_fraction"), 3)), | |
| ("diag_frac_at_cap", f(dg.get("frac_at_cap"), 3)), | |
| ("fork_cos_alpha", f(fork.get("mean_pairwise_cos_alpha"), 3)), | |
| ("fork_cos_a_hat", f(fork.get("mean_pairwise_cos_a_hat"), 3)), | |
| ("train_steps_per_s", f(tr.get("transitions_per_second"), 1)), | |
| ("field_leak", f((rm.get("config") or {}).get("field_leak"), 2)), | |
| ("chat_bonjour", chats[0] if len(chats) > 0 else ""), | |
| ("chat_qui", chats[1] if len(chats) > 1 else ""), | |
| ("chat_il_etait", chats[2] if len(chats) > 2 else ""), | |
| ] | |
| new = not os.path.exists(tsv) | |
| with open(tsv, "a", encoding="utf-8") as fh: | |
| if new: | |
| fh.write("\t".join(c for c, _ in cols) + "\n") | |
| fh.write("\t".join(v for _, v in cols) + "\n") | |
| print("TSV:", "\t".join(v for _, v in cols)) | |
| PYEOF | |
| echo "$(ts) measure $END done" | |
| } | |
| measure_missing() { | |
| for ck in $(ls "$OUT"/atom_native_step_*_${SUF}.pt 2>/dev/null | sed -n 's/.*step_\([0-9]*\)_.*/\1/p' | sort -n); do | |
| if [ "$ck" -ge "$START_STEP" ] && [ "$ck" -le "$END_TARGET" ] && ! has_row "$ck"; then measure "$ck"; fi | |
| done | |
| } | |
| CUR=$(latest) | |
| if [ -z "$CUR" ]; then echo "$(ts) no branch checkpoint in [$START_STEP,$END_TARGET] under $OUT; exiting"; exit 1; fi | |
| while [ "$CUR" -lt "$END_TARGET" ]; do | |
| END=$((CUR + CHUNK)); [ "$END" -gt "$END_TARGET" ] && END=$END_TARGET | |
| STEPS=$((END - CUR)); SKIP=$((CUR - BASE)); [ "$SKIP" -lt 0 ] && SKIP=0 | |
| RESUME="$OUT/atom_native_step_${CUR}_${SUF}.pt" | |
| NAME="atom_native_step_${END}_${SUF}.pt" | |
| LOG="$LOGD/train_${CUR}_${END}_fieldsat.log" | |
| if [ ! -f "$OUT/$NAME" ]; then | |
| echo "$(ts) chunk $CUR->$END skip=$SKIP resume=$RESUME log=$LOG" | |
| T0=$(date +%s) | |
| "$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" --field-leak "$LEAK" \ | |
| > "$LOG" 2>&1 | |
| RC=$? | |
| if [ "$RC" -ne 0 ] || [ ! -f "$OUT/$NAME" ]; then | |
| echo "$(ts) chunk $CUR->$END 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}_fieldsat.json" | |
| echo "$(ts) chunk $CUR->$END saved $NAME wall=$(( $(date +%s) - T0 ))s" | |
| else | |
| echo "$(ts) skip exists $NAME" | |
| fi | |
| measure_missing | |
| CUR=$END | |
| done | |
| measure_missing | |
| echo "$(ts) DONE fieldsat branch at $CUR (END_TARGET=$END_TARGET)" | |