diff options
Diffstat (limited to 'sbatch')
| -rw-r--r-- | sbatch/fw135m_full.sbatch | 54 | ||||
| -rw-r--r-- | sbatch/fw135m_lr_sweep.sbatch | 67 | ||||
| -rw-r--r-- | sbatch/fw135m_smoke.sbatch | 57 |
3 files changed, 178 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 diff --git a/sbatch/fw135m_lr_sweep.sbatch b/sbatch/fw135m_lr_sweep.sbatch new file mode 100644 index 0000000..bdb403f --- /dev/null +++ b/sbatch/fw135m_lr_sweep.sbatch @@ -0,0 +1,67 @@ +#!/bin/bash +#SBATCH --job-name=fw135m-lr +#SBATCH --output=/orion/u/oscarwan/ept/ep_run/runs/slurm-fw135m-lr-%A_%a.out +#SBATCH --error=/orion/u/oscarwan/ept/ep_run/runs/slurm-fw135m-lr-%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=0-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}}" +SEED="${SEED:-1}" +LRS=(7e-4 1e-3 1.4e-3) +LR="${LRS[${SLURM_ARRAY_TASK_ID:?This script must be submitted as an array job}]}" +LR_TAG="${LR//-/m}" +LR_TAG="${LR_TAG//./p}" +TAG="fw135m_bp_lr${LR_TAG}_s${SEED}" + +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}; array task: ${SLURM_ARRAY_TASK_ID}" +echo "Git commit: $(git rev-parse --short HEAD)" +echo "CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES:-unset}" +echo "W&B: ${WAND_ENTITY}/${WAND_PROJECT}; run: ${TAG}" +echo "Adam-side LR: ${LR}; data: ${DATA_DIR}" +nvidia-smi + +"${PYTHON_BIN}" casc_bp_train.py \ + --tag "${TAG}" \ + --L 12 --C 768 --H 12 --T 256 --B 24 \ + --steps 440000 \ + --lr "${LR}" \ + --warmup 1000 \ + --amp --olmo2 \ + --wd 0.1 \ + --opt muon --muon_lr 0.02 \ + --cosine --lr_min_ratio 0.1 \ + --data fineweb_edu \ + --seed "${SEED}" \ + --save_every 5000 --log 100 \ + --wandb "${WAND_PROJECT}" \ + --wandb_run "${TAG}" 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 |