1 files changed, 87 insertions, 0 deletions
diff --git a/research/flossing/flossing_suite/launch_trm_variant_suite.sh b/research/flossing/flossing_suite/launch_trm_variant_suite.sh
new file mode 100755
index 0000000..86c2e64
--- /dev/null
+++ b/research/flossing/flossing_suite/launch_trm_variant_suite.sh
@@ -0,0 +1,87 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+ROOT="/home/yurenh2/rrm"
+PY="/home/yurenh2/miniconda3/envs/rrm/bin/python"
+OUT_DIR="${ROOT}/research/flossing/flossing_suite/results/trm_variants"
+CKPT_ROOT="${ROOT}/trm/checkpoints/Sudoku-extreme-1k-aug-1000-ACT-torch/pretrain_mlp_t_sudoku_official_gbs768_repro"
+mkdir -p "${OUT_DIR}"
+
+GPU_TOP1="${GPU_TOP1:-0}"
+GPU_VOLUME="${GPU_VOLUME:-1}"
+GPU_KL="${GPU_KL:-3}"
+TRAIN_STEPS="${TRAIN_STEPS:-20000}"
+FLOSS_STEPS="${FLOSS_STEPS:-100}"
+INTERFLOSS_EVERY="${INTERFLOSS_EVERY:-2000}"
+INTERFLOSS_START="${INTERFLOSS_START:-2000}"
+INTERFLOSS_STOP="${INTERFLOSS_STOP:-10000}"
+EVAL_N="${EVAL_N:-1000}"
+TASK_BATCH_SIZE="${TASK_BATCH_SIZE:-32}"
+FLOSS_BATCH_SIZE="${FLOSS_BATCH_SIZE:-4}"
+
+write_and_launch() {
+  local gpu="$1"
+  local name="$2"
+  local floss_mode="$3"
+  local kl_beta="$4"
+  local cmd="${OUT_DIR}/${name}.cmd.sh"
+  local log="${OUT_DIR}/${name}.log"
+  local pid="${OUT_DIR}/${name}.pid"
+
+  cat > "${cmd}" <<EOF
+#!/usr/bin/env bash
+set -euo pipefail
+cd "${ROOT}"
+export CUDA_VISIBLE_DEVICES="${gpu}"
+export PYTHONUNBUFFERED=1
+exec "${PY}" research/flossing/step7_interfloss.py \\
+  --model trm \\
+  --ckpt-root "${CKPT_ROOT}" \\
+  --ckpt-name __random__ \\
+  --init-seed 123 \\
+  --train-steps "${TRAIN_STEPS}" \\
+  --batch-size "${TASK_BATCH_SIZE}" \\
+  --task-batch-size "${TASK_BATCH_SIZE}" \\
+  --floss-batch-size "${FLOSS_BATCH_SIZE}" \\
+  --train-lr 1e-4 \\
+  --floss-lr 1e-4 \\
+  --floss-mode "${floss_mode}" \\
+  --lambda-star 0 \\
+  --k-lyap 4 \\
+  --lyap-act-steps 4 \\
+  --seed 42 \\
+  --eval-every 1000 \\
+  --eval-n "${EVAL_N}" \\
+  --eval-batch-size 64 \\
+  --floss-log-every 10 \\
+  --train-puzzle-emb \\
+  --puzzle-emb-lr 1e-4 \\
+  --puzzle-emb-weight-decay 1.0 \\
+  --floss-steps "${FLOSS_STEPS}" \\
+  --interfloss-at "0" \\
+  --interfloss-every "${INTERFLOSS_EVERY}" \\
+  --interfloss-start "${INTERFLOSS_START}" \\
+  --interfloss-stop "${INTERFLOSS_STOP}" \\
+  --kl-beta "${kl_beta}" \\
+  --kl-replay-size 64 \\
+  --kl-batch-size 8 \\
+  --kl-temperature 1.0 \\
+  --out "${OUT_DIR}/${name}.json"
+EOF
+  chmod +x "${cmd}"
+  setsid bash "${cmd}" > "${log}" 2>&1 < /dev/null &
+  echo $! > "${pid}"
+  echo "${name}: pid $(cat "${pid}") on GPU ${gpu}"
+}
+
+write_and_launch "${GPU_TOP1}" \
+  "trm_seed123_top1_cf_periodic${INTERFLOSS_EVERY}_tb${TASK_BATCH_SIZE}_fb${FLOSS_BATCH_SIZE}_k4_${TRAIN_STEPS}" \
+  top1_cf 0
+write_and_launch "${GPU_VOLUME}" \
+  "trm_seed123_volume_cf_periodic${INTERFLOSS_EVERY}_tb${TASK_BATCH_SIZE}_fb${FLOSS_BATCH_SIZE}_k4_${TRAIN_STEPS}" \
+  volume_cf 0
+write_and_launch "${GPU_KL}" \
+  "trm_seed123_volume_cf_kl10_periodic${INTERFLOSS_EVERY}_tb${TASK_BATCH_SIZE}_fb${FLOSS_BATCH_SIZE}_k4_${TRAIN_STEPS}" \
+  volume_cf 10
+
+echo "queued TRM CF/volume variant suite in ${OUT_DIR}"