S2 backtracking sweep launcher

latent_multi_output_cell_policy/backtracking/launch_s2_sweep.sh

@@ -0,0 +1,78 @@
+#!/usr/bin/env bash
+# Stage-2 latent backtracking sweep: 8 combinations, one per GPU, from the
+# Stage-1 GRPO checkpoint. Compares no-backtrack controls vs backtracking
+# (rehearsal of stage-1 consistency with k=1) and warm-up on/off.
+set -uo pipefail
+
+PY=/opt/pytorch/bin/python
+REPO=/home/ubuntu/curriculum-cot-code
+SCRIPT="${REPO}/latent_multi_output_cell_policy/sft_latent_multi_output_train.py"
+S1_GRPO=/home/ubuntu/hf_checkpoints/latent_stages/stage01_latent_grpo_i1_20empty_latent_recurrent_hidden
+TRAIN="${REPO}/data/sudoku_t3_20empty_value_qwen_text_stage1_train.jsonl"
+EVAL="${REPO}/data/sudoku_t3_20empty_value_qwen_text_stage1_eval.jsonl"
+ROOT=/home/ubuntu/bt_runs/s2_sweep_$(date +%Y%m%d_%H%M%S)
+mkdir -p "${ROOT}/logs"
+echo "Sweep root: ${ROOT}"
+
+# Common knobs (kept small enough for an 8-way single-GPU sweep).
+COMMON=(
+  --model_name Qwen/Qwen2.5-1.5B-Instruct
+  --train_jsonl "${TRAIN}" --eval_jsonl "${EVAL}"
+  --cache_dir /home/ubuntu/.hf_cache
+  --init_adapter_dir "${S1_GRPO}"
+  --seed 0 --gpu_id 0
+  --stage_i 2 --num_cot_tokens 2 --latent_mode recurrent_hidden
+  --total_empties_hint 20
+  --per_device_train_batch_size 8 --gradient_accumulation_steps 2
+  --num_epochs 8 --learning_rate 5e-5 --enable_gradient_checkpointing
+  --logging_steps 20 --eval_steps 400 --save_steps 400
+  --eval_rows 20 --max_completion_length 24 --limit_train_rows 4000
+  --max_steps 1200
+  --eval_value_precision_stop 0 --eval_value_recall_stop 0 --eval_solve_rate_stop 0
+  --lora_r 32 --lora_alpha 64 --lora_dropout 0.05
+)
+
+# name|extra flags
+declare -a COMBOS=(
+  "ctrl_nobt|--latent_warmup_steps 0"
+  "ctrl_warm|--latent_warmup_steps 150"
+  "bt_rr03_warm|--backtrack_enable --remember_rate 0.3 --backtrack_detect_threshold 0 --latent_warmup_steps 150"
+  "bt_rr05_warm|--backtrack_enable --remember_rate 0.5 --backtrack_detect_threshold 0 --latent_warmup_steps 150"
+  "bt_rr03_adapt|--backtrack_enable --remember_rate 0.3 --backtrack_detect_threshold 0.95 --latent_warmup_steps 150"
+  "bt_rr05_adapt|--backtrack_enable --remember_rate 0.5 --backtrack_detect_threshold 0.97 --latent_warmup_steps 150"
+  "bt_rr03_nowarm|--backtrack_enable --remember_rate 0.3 --backtrack_detect_threshold 0.95 --latent_warmup_steps 0"
+  "bt_rr02_adapt|--backtrack_enable --remember_rate 0.2 --backtrack_detect_threshold 0.97 --latent_warmup_steps 150"
+)
+
+launch() {
+  local gpu="$1" name="$2" extra="$3"
+  local out="${ROOT}/${name}"
+  local log="${ROOT}/logs/${name}.log"
+  mkdir -p "${out}"
+  echo "[gpu ${gpu}] ${name} :: ${extra}"
+  CUDA_VISIBLE_DEVICES="${gpu}" nohup "${PY}" "${SCRIPT}" "${COMMON[@]}" \
+    --output_dir "${out}" --wandb_mode disabled ${extra} > "${log}" 2>&1 &
+  echo "$!" > "${ROOT}/logs/${name}.pid"
+}
+
+# Launch combo 0 first and wait for the shared stage-2 data cache + pools to be
+# built, so the other 7 jobs hit a warm cache (no concurrent .tmp write race).
+IFS='|' read -r name0 extra0 <<< "${COMBOS[0]}"
+launch 0 "${name0}" "${extra0}"
+echo "Waiting for stage-2 data cache to build (combo 0)..."
+for _ in $(seq 1 120); do
+  if grep -q "latent sft train step 00001" "${ROOT}/logs/${name0}.log" 2>/dev/null; then
+    echo "Cache warm; launching remaining combos."
+    break
+  fi
+  sleep 5
+done
+
+for i in $(seq 1 7); do
+  IFS='|' read -r name extra <<< "${COMBOS[$i]}"
+  launch "${i}" "${name}" "${extra}"
+  sleep 1
+done
+
+echo "All combos launched under ${ROOT}"
+echo "${ROOT}" > /home/ubuntu/bt_runs/LATEST_S2_SWEEP