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, 57 insertions, 0 deletions

diff --git a/sbatch/fw135m_smoke.sbatch b/sbatch/fw135m_smoke.sbatch
new file mode 100644
index 0000000..8fc366f
--- /dev/null
+++ b/sbatch/fw135m_smoke.sbatch
@@ -0,0 +1,57 @@
+#!/bin/bash
+#SBATCH --job-name=fw135m-smoke
+#SBATCH --output=/orion/u/oscarwan/ept/ep_run/runs/slurm-fw135m-smoke-%j.out
+#SBATCH --error=/orion/u/oscarwan/ept/ep_run/runs/slurm-fw135m-smoke-%j.err
+#SBATCH --time=02: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
+
+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}}"
+SEED="${SEED:-1}"
+
+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}"
+echo "Git commit: $(git rev-parse --short HEAD)"
+echo "CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES:-unset}"
+echo "W&B: ${WAND_ENTITY}/${WAND_PROJECT}"
+echo "Data: ${DATA_DIR}"
+nvidia-smi
+
+"${PYTHON_BIN}" - <<'PY'
+import torch, wandb
+print("torch", torch.__version__, "cuda_available", torch.cuda.is_available(), "gpu_count", torch.cuda.device_count())
+print("wandb", wandb.__version__)
+PY
+
+"${PYTHON_BIN}" fw135m_baseline.py \
+  --mode smoke \
+  --seed "${SEED}" \
+  --wandb_project "${WAND_PROJECT}" \
+  --execute