ATOM (atom-ai): code snapshot (100fae0 + batched-train), README with measured numbers, curated checkpoints (part 4)
b9301f1 verified | # STEP 3 of docs/DECODE_FIX_AND_BPTT.md: long full-corpus run with truncated BPTT through the | |
| # field (--field-bptt K), resumable, 8k-step chunks, each chunk a NEW checkpoint in a NEW dir. | |
| # Re-running resumes from the newest checkpoint in $OUT (an interrupted chunk is re-run). | |
| # | |
| # Env: INIT=<ckpt to start from> START=<its step> END_TARGET CHUNK=8000 | |
| # K=64 LEAK=0.05 LEARN=1 (learnable per-mode leak) LR=1e-3 SLR=1e-3 LADROP=0.5 | |
| # OUT=checkpoints/byte_tick_bptt SUF=d32_ms1M_fieldsat_decfix_bptt MEASURE_EVERY=2 | |
| # Stream position: SKIP = (CUR - 496000) mod 1000452 (one corpus epoch = 1,000,453 bytes = 1,000,452 | |
| # byte-tick transitions): same text position as the legacy skip without replaying 4.5M packets | |
| # (the 256-period phase feature differs by a constant offset; irrelevant for this new branch). | |
| # Launch: nohup env INIT=... START=... bash scripts/train_bptt.sh >> logs/bptt/master.log 2>&1 & | |
| # | |
| # Opt-in variables (defaults = exact previous behaviour): | |
| # DATA_GLOB='data/corpus_fr_medium*.txt' corpus glob (mixed corpus: 'data/corpus_fr_mix_[0-9]*.txt') | |
| # EPOCH=1000452 corpus length in byte-tick transitions (mix: total bytes from the manifest) | |
| # SKIP_BYTES=0 1 = position the stream with --stream-skip-bytes (seek, no replay) instead of | |
| # --stream-skip-packets (replay costs ~5 min per 4.5M packets on a 7 MB corpus) | |
| # TSV_SCRIPT=$A/tsv_bptt.py progress-row script (mix: docs/artifacts/mix_long_memory/tsv_mix.py) | |
| # MEASURE_ASYNC=0 1 = run the progress row in the background (1 thread, pinned to MEASURE_CPUS), | |
| # at most MEASURE_MAX concurrent, so training does not wait for it | |
| # CPUS="" taskset core list for the training process ("" = no pinning) | |
| # EXTRA_ARGS="" extra trainer flags (e.g. a new opt-in lever) | |
| set -u | |
| cd /workspace/exports/atom-ai | |
| INIT=${INIT:?INIT ckpt}; START=${START:?START step}; END_TARGET=${END_TARGET:-$((START + 2000000))} | |
| CHUNK=${CHUNK:-8000}; K=${K:-64}; LEAK=${LEAK:-0.05}; LEARN=${LEARN:-1}; LR=${LR:-1e-3}; SLR=${SLR:-$LR} | |
| LADROP=${LADROP:-0.5}; MEASURE_EVERY=${MEASURE_EVERY:-2} | |
| OUT=${OUT:-checkpoints/byte_tick_bptt}; SUF=${SUF:-d32_ms1M_fieldsat_decfix_bptt} | |
| A=docs/artifacts/decfix_bptt; TSV=${TSV:-$A/bptt_progress.tsv}; LOGD=${LOGD:-logs/bptt} | |
| EPOCH=${EPOCH:-1000452}; BASE=${BASE:-496000}; PY=.venv/bin/python | |
| DATA_GLOB=${DATA_GLOB:-data/corpus_fr_medium*.txt}; SKIP_BYTES=${SKIP_BYTES:-0}; TSV_SCRIPT=${TSV_SCRIPT:-$A/tsv_bptt.py} | |
| MEASURE_ASYNC=${MEASURE_ASYNC:-0}; MEASURE_CPUS=${MEASURE_CPUS:-5-7}; MEASURE_MAX=${MEASURE_MAX:-2}; CPUS=${CPUS:-} | |
| EXTRA_ARGS=${EXTRA_ARGS:-} | |
| [ -x "$PY" ] || { echo "$(date -Is) $PY missing: run scripts/ensure_venv.sh"; exit 1; } | |
| PIN=(); [ -n "$CPUS" ] && PIN=(taskset -c "$CPUS") | |
| SKIP_FLAG=--stream-skip-packets; [ "$SKIP_BYTES" = "1" ] && SKIP_FLAG=--stream-skip-bytes | |
| export PYTHONPATH=. OMP_NUM_THREADS=${THREADS:-1} MKL_NUM_THREADS=${THREADS:-1} | |
| mkdir -p "$OUT" "$LOGD" | |
| exec 9>"$LOGD/driver.lock"; flock -n 9 || { echo "$(date -Is) another driver holds $LOGD/driver.lock"; exit 1; } | |
| echo "=== $(date -Is) bptt driver pid=$$ INIT=$INIT START=$START K=$K LEAK=$LEAK LEARN=$LEARN LR=$LR SLR=$SLR LADROP=$LADROP OUT=$OUT ===" | |
| LEARN_FLAG=""; [ "$LEARN" = "1" ] && LEARN_FLAG="--field-leak-learnable" | |
| latest() { ls "$OUT"/atom_native_step_*_${SUF}.pt 2>/dev/null | sed -n 's/.*step_\([0-9]*\)_.*/\1/p' | sort -n | tail -1; } | |
| CUR=$(latest); if [ -z "$CUR" ]; then CUR=$START; RESUME=$INIT; else RESUME="$OUT/atom_native_step_${CUR}_${SUF}.pt"; fi | |
| N=0 | |
| while [ "$CUR" -lt "$END_TARGET" ]; do | |
| END=$((CUR + CHUNK)); NAME="atom_native_step_${END}_${SUF}.pt"; SKIP=$(( (CUR - BASE) % EPOCH )) | |
| LOG="$LOGD/train_${CUR}_${END}_bptt.log"; T0=$(date +%s) | |
| echo "$(date -Is) chunk $CUR->$END skip=$SKIP resume=$RESUME" | |
| "${PIN[@]}" "$PY" -u tools/run_atom_native.py \ | |
| --stream --data-glob "$DATA_GLOB" --loop-shards --chunk-bytes 65536 \ | |
| $SKIP_FLAG "$SKIP" --resume "$RESUME" --output-dir "$OUT" \ | |
| --checkpoint-name "$NAME" --steps "$CHUNK" \ | |
| --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 0 --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" --field-bptt "$K" $LEARN_FLAG $EXTRA_ARGS > "$LOG" 2>&1 | |
| RC=$? | |
| if [ "$RC" -ne 0 ] || [ ! -f "$OUT/$NAME" ]; then echo "$(date -Is) chunk FAILED rc=$RC see $LOG"; exit 1; fi | |
| cp "$OUT/run_metrics.json" "$OUT/run_metrics_${END}_${SUF}.json" | |
| echo "$(date -Is) chunk $CUR->$END saved wall=$(( $(date +%s) - T0 ))s" | |
| RESUME="$OUT/$NAME"; CUR=$END; N=$((N + 1)) | |
| if [ $((N % MEASURE_EVERY)) -eq 0 ]; then | |
| if [ "$MEASURE_ASYNC" = "1" ]; then | |
| NRUN=$(pgrep -f -- "$TSV_SCRIPT $TSV" | wc -l) | |
| if [ "$NRUN" -lt "$MEASURE_MAX" ]; then | |
| OMP_NUM_THREADS=1 MKL_NUM_THREADS=1 nohup taskset -c "$MEASURE_CPUS" "$PY" -u $TSV_SCRIPT "$TSV" "$OUT/$NAME" \ | |
| "$OUT/run_metrics_${END}_${SUF}.json" >> "$LOGD/measure.log" 2>&1 < /dev/null 9>&- & | |
| else | |
| echo "$(date -Is) measure skipped at $END ($NRUN running)" | |
| fi | |
| else | |
| "$PY" -u $TSV_SCRIPT "$TSV" "$OUT/$NAME" "$OUT/run_metrics_${END}_${SUF}.json" >> "$LOGD/measure.log" 2>&1 \ | |
| && tail -n 1 "$TSV" | cut -c1-400 | |
| fi | |
| fi | |
| done | |
| echo "$(date -Is) DONE at $CUR" | |