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

#!/usr/bin/env bash
# SLURM job: NGDiff unlearning for one topic / topic-set
# Usage: sbatch unlearn_ngdiff.sh --topic_bin <topic> [options]
#        sbatch unlearn_ngdiff.sh --topics t1,t2,t3 --output_dir <dir> [options]
#
#SBATCH --job-name=ngdiff_unlearn
#SBATCH --partition=coe-gpu
#SBATCH --qos=coe-ice
#SBATCH --account=ic
#SBATCH --gres=gpu:1
#SBATCH --cpus-per-task=8
#SBATCH --mem=64G
#SBATCH --time=16:00:00
#SBATCH --output=logs/unlearn/ngdiff_%j_%x.out
#SBATCH --error=logs/unlearn/ngdiff_%j_%x.err

set -euo pipefail

# ---- Parse args ------------------------------------------------------- #
TOPIC_BIN=""
TOPICS=""           # comma-separated list (multi-topic mode)
MAX_STEPS=5000
MAX_FORGET_DOCS=""
MAX_RETAIN_DOCS=""
LR=""
NO_AUTO_LR=false
OUTPUT_DIR_OVERRIDE=""
SAVE_STEPS=""
RETAIN_TOPICS=""
SEED=""

while [[ $# -gt 0 ]]; do
    case $1 in
        --topic_bin)       TOPIC_BIN="$2";           shift 2;;
        --topics)          TOPICS="$2";              shift 2;;
        --max_steps)       MAX_STEPS="$2";            shift 2;;
        --max_docs)        MAX_FORGET_DOCS="$2";      shift 2;;
        --max_retain_docs) MAX_RETAIN_DOCS="$2";      shift 2;;
        --lr)              LR="$2";                   shift 2;;
        --no_auto_lr)      NO_AUTO_LR=true;           shift;;
        --output_dir)      OUTPUT_DIR_OVERRIDE="$2";  shift 2;;
        --save_steps)      SAVE_STEPS="$2";           shift 2;;
        --retain_topics)   RETAIN_TOPICS="$2";       shift 2;;
        --seed)            SEED="$2";              shift 2;;
        *) echo "Unknown arg: $1"; exit 1;;
    esac
done

# ---- Resolve topic / topics ------------------------------------------- #
# Multi-topic uses topic_bins=[t1,t2,...] (Hydra list syntax).
# Single-topic uses topic_bin=t (original behaviour).
if [[ -n "$TOPICS" ]]; then
    # Multi-topic: pass as Hydra list override
    TOPIC_HYDRA_OVERRIDE="topic_bins=[${TOPICS}]"
    TOPIC_LABEL="${TOPICS}"
elif [[ -n "$TOPIC_BIN" ]]; then
    TOPIC_HYDRA_OVERRIDE="topic_bin=${TOPIC_BIN}"
    TOPIC_LABEL="${TOPIC_BIN}"
else
    echo "ERROR: --topic_bin or --topics is required"; exit 1
fi

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

OUTPUT_DIR="${OUTPUT_DIR_OVERRIDE:-runs/unlearn/${TOPIC_LABEL}}"
mkdir -p "${OUTPUT_DIR}"

export WANDB_MODE=offline
export WANDB_PROJECT="data-attribution-unlearn"
export WANDB_RUN_NAME="ngdiff-${SLURM_JOB_ID}"
if [[ -f .env.secret ]]; then
    set -a; source .env.secret; set +a
fi
export HF_HOME="$HOME/scratch/hf_cache"
export DOLMA_CACHE="$HOME/scratch/hf_cache/datasets/dolma3_6t_filtered"
export TOKENIZERS_PARALLELISM=false
export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True

# ---- Build Hydra overrides -------------------------------------------- #
# Use an array to avoid quoting issues with spaces/commas
OVERRIDES=(
    "${TOPIC_HYDRA_OVERRIDE}"
    "trainer.max_steps=${MAX_STEPS}"
    "output_dir=${OUTPUT_DIR}"
)
[[ -n "$MAX_FORGET_DOCS"  ]] && OVERRIDES+=("data.max_forget_docs=${MAX_FORGET_DOCS}")
[[ -n "$MAX_RETAIN_DOCS"  ]] && OVERRIDES+=("data.max_retain_docs=${MAX_RETAIN_DOCS}")
[[ -n "$LR"               ]] && OVERRIDES+=("trainer.learning_rate=${LR}")
[[ -n "$SAVE_STEPS"       ]] && OVERRIDES+=("trainer.save_steps=${SAVE_STEPS}")
[[ "$NO_AUTO_LR" == true  ]] && OVERRIDES+=("trainer.auto_lr=false")
[[ -n "$RETAIN_TOPICS"   ]] && OVERRIDES+=("retain_topics=${RETAIN_TOPICS}")
[[ -n "$SEED"            ]] && OVERRIDES+=("data.seed=${SEED}")

echo "===== NGDiff Unlearning ====="
echo "Job ID:     ${SLURM_JOB_ID}"
echo "Topic(s):   ${TOPIC_LABEL}"
echo "Max steps:  ${MAX_STEPS}"
echo "Output dir: ${OUTPUT_DIR}"
[[ -n "$LR"              ]] && echo "LR:         ${LR}"
[[ -n "$MAX_FORGET_DOCS" ]] && echo "Forget:     ${MAX_FORGET_DOCS}"
[[ -n "$MAX_RETAIN_DOCS" ]] && echo "Retain:     ${MAX_RETAIN_DOCS}"
[[ -n "$RETAIN_TOPICS"   ]] && echo "Retain topics: ${RETAIN_TOPICS}"
date

# ---- Run --------------------------------------------------------------- #
PYTHONPATH=src .venv/bin/python -m unlearning.train \
    --config-dir src/unlearning/configs \
    "${OVERRIDES[@]}"

echo "===== Done ====="
date