File size: 3,483 Bytes
48c96cf | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 | #!/usr/bin/env bash
# Wave-5: push baseline 1.5B past solve=0.35.
#
# Idea: best ckpts so far cap at per-cell-exact ~0.943 (solve 0.35 = 0.943^20).
# To reach solve=0.5 we need exact ~= 0.965. That's +2.2pp of per-cell exact.
#
# 4 variants, single-GPU each, on GPUs 4..7.
# All start from the leader (pipe_m post-S3-GRPO at solve=0.35) or its S3 SFT
# ckpt, then push S3 GRPO further with different levers:
# - lower LR (escape / fine refine)
# - longer steps (3000 instead of 1500)
# - KL anchor (beta>0) to prevent regression
# - sharper rewards (mirror what worked for the latent's `s3_grpo_sharp_rwd`)
set -euo pipefail
ROOT=/home/ubuntu/curriculum_cot
SWEEP_ROOT=/home/ubuntu/curriculum_cot/_runs/baseline_1p5b_v4_20260523_184952
PIPELINE=$ROOT/_runs/baseline_1p5b_pipeline_v4.sh
# best wave-2 anchors
PIPE_M_S3GRPO_LATEST=$(ls -dt $SWEEP_ROOT/pipe_m_s3sft_from_b/s3_grpo/checkpoint-* 2>/dev/null | head -1)
PIPE_M_S3SFT_LATEST=$SWEEP_ROOT/pipe_m_s3sft_from_b/s3_sft/checkpoint-step-02400
PIPE_O_S3SFT_LATEST=$SWEEP_ROOT/pipe_o_s3sft_lr5e6/s3_sft/checkpoint-step-02400
PIPE_J_S3GRPO_LATEST=$(ls -dt $SWEEP_ROOT/pipe_j_s3sft_lr5e5_lr1e5/s3_grpo/checkpoint-* 2>/dev/null | head -1)
# Sanity
for c in "$PIPE_M_S3GRPO_LATEST" "$PIPE_M_S3SFT_LATEST" "$PIPE_O_S3SFT_LATEST" "$PIPE_J_S3GRPO_LATEST"; do
[[ -d "$c" ]] || { echo "MISSING: $c"; exit 1; }
done
CKPT_LR5E5=$ROOT/checkpoints/sudoku-9x9-20empty-baseline-1p5b-sweep/baseline_lr5e5_lowsft_v3/s2_sft_v3/checkpoint-step-03000
launch() {
local gpu="$1" variant="$2"; shift 2
local out=$SWEEP_ROOT/$variant; mkdir -p "$out"
nohup env ROOT="$ROOT" VARIANT="$variant" GPU="$gpu" S2_SFT_CKPT="$CKPT_LR5E5" \
OUTPUT_ROOT="$out" USE_WANDB=0 WANDB_MODE=offline "$@" \
bash "$PIPELINE" </dev/null >"$out/nohup.log" 2>&1 &
local pid=$!
echo "$pid $gpu $variant" >> "$SWEEP_ROOT/PIDS.txt"
disown $pid 2>/dev/null || true
printf 'GPU %s -> %s pid=%s\n' "$gpu" "$variant" "$pid"
}
# pipe_t (GPU 4): continue pipe_m's S3 GRPO with lower LR + KL anchor + longer steps.
# Keep the policy near the SFT reference to avoid the regression we saw earlier.
launch 4 pipe_t_grpo_low_kl \
START_PHASE=s3_grpo S3_GRPO_INIT="$PIPE_M_S3GRPO_LATEST" \
GRPO_LR=1e-6 GRPO_BS=32 GRPO_GA=1 GRPO_NG=8 \
GRPO_BETA=0.04 \
S3_GRPO_MAX_STEPS=3000 \
USE_GC=0
# pipe_u (GPU 5): re-run S3 GRPO from pipe_m's S3-SFT ckpt with sharper rewards
# (mirror latent `s3_grpo_sharp_rwd` recipe: bigger penalty for bad).
launch 5 pipe_u_grpo_sharp_rwd \
START_PHASE=s3_grpo S3_GRPO_INIT="$PIPE_M_S3SFT_LATEST" \
GRPO_LR=5e-6 GRPO_BS=32 GRPO_GA=1 GRPO_NG=8 \
REWARD_GOOD=1.5 PENALTY_BAD=2.0 PENALTY_MAL=4.0 \
S3_GRPO_MAX_STEPS=3000 \
USE_GC=0
# pipe_v (GPU 6): extend pipe_o's S3 SFT (the strongest pure-SFT path) with very
# low LR for 4000 more steps. Then S3 GRPO at LR=1e-6.
launch 6 pipe_v_sft_extend \
START_PHASE=s3_sft S3_SFT_INIT="$PIPE_O_S3SFT_LATEST" \
SFT_LR_S3=2e-6 SFT_BS=16 SFT_GA=1 \
S3_SFT_MAX_STEPS=4000 \
GRPO_LR=1e-6 GRPO_BS=32 GRPO_GA=1 GRPO_NG=8 \
S3_GRPO_MAX_STEPS=2000 \
USE_GC=0
# pipe_w (GPU 7): continue pipe_j's S3 GRPO with very low LR + KL anchor.
# Different lineage from pipe_m, so this gives an independent push.
launch 7 pipe_w_j_low_kl \
START_PHASE=s3_grpo S3_GRPO_INIT="$PIPE_J_S3GRPO_LATEST" \
GRPO_LR=2e-6 GRPO_BS=32 GRPO_GA=1 GRPO_NG=8 \
GRPO_BETA=0.02 \
S3_GRPO_MAX_STEPS=3000 \
USE_GC=0
echo
echo "=== launched ==="
cat "$SWEEP_ROOT/PIDS.txt" | tail -4
|