artifacts/adamw_sweep_cmpatino-0/run_lr_sweep.sh

#!/usr/bin/env bash
# LR sweep at half-length with best WD from prior sweep (wd=0.05).
set -euo pipefail

cd "$(dirname "$0")/../../modded-nanogpt"

SWEEP_SCRIPT="../artifacts/adamw_baseline_v2_cmpatino-0/train_gpt_adamw_v2_cmpatino-0.py"
RESULTS_DIR="../artifacts/adamw_sweep_cmpatino-0"
mkdir -p "$RESULTS_DIR"

export TRAIN_STEPS=2812
export WARMUP_STEPS=125
export COOLDOWN_FRAC=0.7
export BLOCK_WD=0.05
export BLOCK_BETA1=0.9
export BLOCK_BETA2=0.95

for LR in 0.0010 0.0020 0.0030; do
  TAG="lr${LR}_wd0.05_steps${TRAIN_STEPS}_cmpatino-0"
  LOG="${RESULTS_DIR}/log_${TAG}.txt"
  echo "=== Running BLOCK_LR=${LR} -> ${LOG} ==="
  BLOCK_LR=$LR RUN_TAG="$TAG" \
    ../.venv/bin/torchrun --standalone --nproc_per_node=2 "$SWEEP_SCRIPT" 1 \
    2>&1 | tee "$LOG"
done

echo "=== LR sweep complete ==="