Add portable 135M BP sweep handoff

Match the active EP configuration and support four-GPU BP candidates from collaborators' data locations.

Co-authored-by: Cursor <cursoragent@cursor.com>

1 files changed, 54 insertions, 0 deletions

diff --git a/sbatch/fw135m_full.sbatch b/sbatch/fw135m_full.sbatch
new file mode 100644
index 0000000..f0a8b79
--- /dev/null
+++ b/sbatch/fw135m_full.sbatch
@@ -0,0 +1,54 @@
+#!/bin/bash
+#SBATCH --job-name=fw135m-full
+#SBATCH --output=/orion/u/oscarwan/ept/ep_run/runs/slurm-fw135m-full-%A_%a.out
+#SBATCH --error=/orion/u/oscarwan/ept/ep_run/runs/slurm-fw135m-full-%A_%a.err
+#SBATCH --time=72:00:00
+#SBATCH --account=orion
+#SBATCH --partition=orion
+#SBATCH --nodes=1
+#SBATCH --ntasks-per-node=1
+#SBATCH --gres=gpu:a6000:1
+#SBATCH --cpus-per-task=16
+#SBATCH --mem=64G
+#SBATCH --array=1-2
+
+set -euo pipefail
+
+ROOT_DIR="${ROOT_DIR:-/orion/u/oscarwan/ept}"
+RUN_DIR="${ROOT_DIR}/ep_run"
+VENV_DIR="${VENV_DIR:-${ROOT_DIR}/.venv}"
+PYTHON_BIN="${PYTHON_BIN:-python}"
+WAND_PROJECT="ept-fineweb-135M"
+WAND_ENTITY="eqprop-llm-training"
+DATA_DIR="${DATA_DIR:-${RUN_DIR}/data/fineweb_edu}"
+export EPT_DATA_ROOT="${EPT_DATA_ROOT:-${DATA_DIR%/fineweb_edu}}"
+SELECTED_LR="${SELECTED_LR:?Submit with --export=ALL,SELECTED_LR=<winning-LR>}"
+SEED="${SLURM_ARRAY_TASK_ID:?This script must be submitted as an array job}"
+
+cd "${RUN_DIR}"
+mkdir -p runs
+
+if [[ -f "${VENV_DIR}/bin/activate" ]]; then
+  source "${VENV_DIR}/bin/activate"
+fi
+
+if [[ ! -f "${DATA_DIR}/meta.pkl" || ! -f "${DATA_DIR}/train.bin" || ! -f "${DATA_DIR}/val.bin" ]]; then
+  echo "FineWeb-Edu data is missing at ${DATA_DIR}." >&2
+  echo "Run prepare_fineweb.py first to create the dataset." >&2
+  exit 1
+fi
+
+echo "Host: $(hostname)"
+echo "Job ID: ${SLURM_JOB_ID:-local}; seed: ${SEED}"
+echo "Git commit: $(git rev-parse --short HEAD)"
+echo "CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES:-unset}"
+echo "W&B: ${WAND_ENTITY}/${WAND_PROJECT}"
+echo "Selected Adam-side LR: ${SELECTED_LR}; data: ${DATA_DIR}"
+nvidia-smi
+
+"${PYTHON_BIN}" fw135m_baseline.py \
+  --mode full \
+  --lr "${SELECTED_LR}" \
+  --seed "${SEED}" \
+  --wandb_project "${WAND_PROJECT}" \
+  --execute