#!/bin/bash
# fork_launch.sh — RB4 branching: resume v1_muon @280k and continue on a given data-dir.
# Reads arch+recipe from ckpt["config"]; supplies matched training hyperparams (batch/lr/opt).
# fresh-resume-arcmix (2sigma-seed) and fork-B (arcmix+edu) MUST share batch/recipe (this script).
#
# Usage: bash fork_launch.sh <DATA_DIR> <OUT_DIR> <HF_SUBDIR> <SEED>
#   DATA_DIR : dir containing train.bin + val.bin (uint16)
#   OUT_DIR  : run dir; ckpt.pt (=ckpt_280k) placed here for --resume
#   HF_SUBDIR: HF subdir under SlayerLab/gollem-v5-ckpts for autopush (e.g. forkA_arcmix_seed7)
#   SEED     : rng seed
set -euo pipefail
DATA_DIR="$1"; OUT="$2"; HFSUB="$3"; SEED="${4:-1337}"
export HF_HUB_ENABLE_HF_TRANSFER=0
export HF_TOKEN="$(cat ~/.cache/huggingface/token)"

echo "=== fork_launch: DATA_DIR=$DATA_DIR OUT=$OUT HFSUB=$HFSUB SEED=$SEED ==="
mkdir -p "$OUT"
# resume ckpt must be at $OUT/ckpt.pt
if [ ! -f "$OUT/ckpt.pt" ]; then echo "!! missing $OUT/ckpt.pt (branch-point)"; exit 1; fi

# read arch+recipe from ckpt config
CFG="$(python - "$OUT/ckpt.pt" <<'PY'
import sys, torch
c = torch.load(sys.argv[1], map_location="cpu")["config"]
s = (f"--n-layer {c['n_layer']} --n-embd {c['n_embd']} --n-head {c['n_head']} "
     f"--block {c['block']} --vocab {c['vocab']} --norm {c['norm']} --pos {c['pos']} "
     f"--rope-theta {c['rope_theta']} --ffn {c['ffn']} --ffn-mult {c['ffn_mult']}")
if c.get("value_residual"): s += " --value-residual"
if c.get("qk_norm"): s += " --qk-norm"
print(s)
PY
)"
echo "CKPT-CONFIG: $CFG"

# v1 recipe training hyperparams (matched across forks); batch=32 (5090-safe, ~33k tok/step)
nohup python -u train_gpt_ref.py \
  --data-dir "$DATA_DIR" --out-dir "$OUT" \
  $CFG \
  --batch 32 --steps 400000 --lr 6e-4 --min-lr 6e-5 --warmup 2000 --wd 0.1 \
  --optimizer muon --muon-lr 0.02 --dtype uint16 --seed "$SEED" \
  --ckpt-every 20000 --eval-every 999999 --log-every 200 \
  --events-jsonl "$OUT/events.jsonl" --compile --resume \
  > "$OUT/train.log" 2>&1 &
TRAIN_PID=$!
echo "TRAIN-PID=$TRAIN_PID  (log: $OUT/train.log)"

# autopush ckpts to HF as they land (ckpt.pt is a FILE, not dir)
sleep 5
nohup python -u ckpt_autopush.py "$OUT/train.log" "$OUT/ckpt.pt" "$HFSUB" \
  > "$OUT/autopush.log" 2>&1 &
echo "AUTOPUSH-PID=$!  (HF subdir: $HFSUB)"
echo "=== launched. tail -f $OUT/train.log ==="