mechanistic_interp/scripts/patching.sh

#!/bin/bash
# =============================================================================
# patching.sh — activation patching: bath→toilet causal tracing, base vs lora.
#
# For each layer l: patches corrupted (bathroom-with-toilet) activations into
# the clean (bathroom-only) run, measures ΔP(toilet) = P(toilet|patch) − baseline.
# P(toilet) = P("to") × P("iel"|"to") via teacher-forced forward (single pass).
#
# Tunables:
#   DEVICE_ID       GPU id                            (default 0)
#   DTYPE           bfloat16 | float16 | float32      (default bfloat16)
#   VARIANTS        space-separated list              (default "base lora")
#   N_PAIRS         image pairs to average over       (default 50)
#   SEED            random seed for pair sampling     (default 42)
#   HOOK_TYPE       pre | mid | post                  (default post)
#   STAMP           datestamp suffix for output files (default today)
# =============================================================================
export HF_HOME="/data/caotue/hf_cache"
export HF_DATASETS_CACHE="/data/caotue/hf_cache/datasets"
export TORCH_HOME="/data/caotue/torch_cache"
export TMPDIR="/data/caotue/tmp"

DEVICE_ID="${DEVICE_ID:-0}"
DTYPE="${DTYPE:-bfloat16}"
VARIANTS="${VARIANTS:-base lora}"
N_PAIRS="${N_PAIRS:-50}"
SEED="${SEED:-42}"
HOOK_TYPE="${HOOK_TYPE:-post}"
STAMP="${STAMP:-$(date +%Y%m%d)}"

REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)"
cd "${REPO_ROOT}"
export PYTHONPATH="$(cd .. && pwd):$(pwd):${PYTHONPATH:-}"

OUT="mechanistic_interp/graph/patching_bath_toilet_${STAMP}.png"
OUT_JSON="mechanistic_interp/graph/patching_bath_toilet_${STAMP}.json"

python -m mechanistic_interp.patching \
    --device_id   "${DEVICE_ID}" \
    --dtype       "${DTYPE}" \
    --variants    ${VARIANTS} \
    --n_pairs     "${N_PAIRS}" \
    --seed        "${SEED}" \
    --hook_type   "${HOOK_TYPE}" \
    --out         "${OUT}" \
    --out_json    "${OUT_JSON}" \
    "$@"