social-data-attribution-w2/scripts/slurm/unlearning/sweep_ngdiff.sh

#!/usr/bin/env bash
# Hyperparameter sweep: forget set size × fallback LR × forget:retain ratio
# Topics: science_math_and_technology, software_development, social_life, politics, software
# Each job: 300 steps, auto_lr disabled, followed by fast eval (200 samples/task)
#
# Usage: bash scripts/slurm/unlearning/sweep_ngdiff.sh [--dry_run]

set -euo pipefail

cd ~/dev/data-attribution-soc103
mkdir -p logs/unlearn

DRY_RUN=false
[[ "${1:-}" == "--dry_run" ]] && DRY_RUN=true

TOPICS=(
    science_math_and_technology
    software_development
    social_life
    politics
    software
)

# Sweep matrix: forget retain lr tag
# Three axes varied one at a time; center point = forget=500, retain=4500 (1:9), lr=1e-5
CONDITIONS=(
    # Forget set size sweep (ratio=1:9, lr=1e-5)
    "100  900   1e-5  f100_r900_lr1e-5"
    "300  2700  1e-5  f300_r2700_lr1e-5"
    "500  4500  1e-5  f500_r4500_lr1e-5"   # center / reference
    "1000 9000  1e-5  f1000_r9000_lr1e-5"
    # LR sweep (forget=500, ratio=1:9)
    "500  4500  1e-6  f500_r4500_lr1e-6"
    "500  4500  1e-4  f500_r4500_lr1e-4"
    # Ratio sweep (forget=500, lr=1e-5)
    "500  500   1e-5  f500_r500_lr1e-5"    # 1:1
    "500  1500  1e-5  f500_r1500_lr1e-5"   # 1:3  (previous default)
    "500  3000  1e-5  f500_r3000_lr1e-5"   # 1:6
    # 1:9 already covered above (center)
)

echo "===== NGDiff Hparam Sweep ====="
echo "Topics:     ${TOPICS[*]}"
echo "Conditions: ${#CONDITIONS[@]} per topic"
echo "Total jobs: $((${#TOPICS[@]} * ${#CONDITIONS[@]} * 2)) (unlearn + fast eval each)"
echo ""

TOTAL=0
for topic in "${TOPICS[@]}"; do
    for cond in "${CONDITIONS[@]}"; do
        read -r forget retain lr tag <<< "$cond"
        outdir="runs/unlearn/sweep/${topic}/${tag}"

        if $DRY_RUN; then
            echo "[DRY] topic=${topic} forget=${forget} retain=${retain} lr=${lr} -> ${outdir}"
        else
            # Submit unlearn job (save_steps=100 ensures adapter saved within 300 steps)
            unlearn_jid=$(sbatch --parsable \
                scripts/slurm/unlearning/unlearn_ngdiff.sh \
                --topic_bin       "${topic}" \
                --max_steps       300 \
                --save_steps      100 \
                --max_docs        "${forget}" \
                --max_retain_docs "${retain}" \
                --lr              "${lr}" \
                --no_auto_lr \
                --output_dir      "${outdir}")

            # Submit fast eval chained on unlearn completion
            # adapter/ is saved by train.py after training completes
            eval_jid=$(sbatch --parsable \
                --dependency=afterok:${unlearn_jid} \
                scripts/slurm/unlearning/eval_fast.sh \
                --topic_bin   "${topic}" \
                --adapter_dir "${outdir}/adapter" \
                --output_json "${outdir}/eval_fast_results.json" \
                --samples     200)

            echo "topic=${topic} tag=${tag}  unlearn=${unlearn_jid} eval=${eval_jid}"
            TOTAL=$((TOTAL + 2))
        fi
    done
done

if ! $DRY_RUN; then
    echo ""
    echo "Submitted ${TOTAL} jobs total."
    echo "Results will land in runs/unlearn/sweep/<topic>/<tag>/eval_fast_results.json"
fi