#!/usr/bin/env bash
# Run one EP-matched 135M BP LR-sweep candidate.
#
# Usage:
#   ./runs/fw135m_bp_sweep.sh <7e-4|1e-3|1.4e-3> [seed]
# Set GPUS=4 to launch a four-A6000 NCCL data-parallel run with torchrun.

set -euo pipefail

LR="${1:?Usage: $0 <7e-4|1e-3|1.4e-3> [seed]}"
SEED="${2:-1}"
case "${LR}" in
  7e-4|1e-3|1.4e-3) ;;
  *) echo "Unsupported LR: ${LR}" >&2; exit 2 ;;
esac

RUN_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
PYTHON_BIN="${PYTHON_BIN:-python}"
WAND_PROJECT="${WAND_PROJECT:-ept-fineweb-135M}"
GPUS="${GPUS:-1}"
export EPT_DATA_ROOT="${EPT_DATA_ROOT:-${RUN_DIR}/data}"

LR_TAG="${LR//-/m}"
LR_TAG="${LR_TAG//./p}"
TAG="fw135m_bp_lr${LR_TAG}_s${SEED}"

cd "${RUN_DIR}"
if (( GPUS > 1 )); then
  LAUNCH=(torchrun --standalone --nproc_per_node="${GPUS}")
else
  LAUNCH=("${PYTHON_BIN}")
fi

exec "${LAUNCH[@]}" 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}"