social-data-attribution-w2/scripts/ekfac.sh

#!/usr/bin/env bash
set -euo pipefail

# EK-FAC attribution driver for the pinned Bergson (v0.9.0).
#
# Replaces the earlier scaffold that called a non-existent `bergson ekfac`
# subcommand. The pinned Bergson exposes the EK-FAC pieces as separate
# stages, so this driver chains them with the real CLI:
#
#   A. bergson build   -> raw gradient indices for value and query data
#   B. bergson hessian -> KFAC covariances + eigenvalue correction (EK-FAC factors)
#   C. apply_hessian   -> inverse-Hessian-vector product (IVHP) preconditioned grads
#   D. (hand-off)       -> existing TrackStar dot-score over IVHP query grads
#
# Stage D reuses the unchanged TrackStar scoring path, so this driver owns
# only the EK-FAC-specific work (A-C). The follow-on command is printed at
# the end.
#
# Usage:
#   bash scripts/ekfac.sh \
#       --model allenai/Olmo-3-1025-7B \
#       --value-data /path/to/value.jsonl \
#       --query-data /path/to/query.jsonl \
#       --output-dir ~/scratch/ekfac/olmo7b \
#       --gpus 8

MODEL="allenai/Olmo-3-1025-7B"
VALUE_DATA=""
QUERY_DATA=""
OUTPUT_DIR=""
PROMPT_COLUMN="text"
COMPLETION_COLUMN=""
PROJECTION_DIM=16
PRECISION="fp32"
HESSIAN_DTYPE="bf16"
TOKEN_BATCH=4096
LAMBDA_DAMP=0.1
GPUS=1
EV_CORRECTION=true
IVHP_SIDE="query"
PYTHONEXEC=""
DRY_RUN=false

while [ $# -gt 0 ]; do
    case "$1" in
        --model) MODEL="$2"; shift 2 ;;
        --value-data) VALUE_DATA="$2"; shift 2 ;;
        --query-data) QUERY_DATA="$2"; shift 2 ;;
        --output-dir) OUTPUT_DIR="$2"; shift 2 ;;
        --prompt-column) PROMPT_COLUMN="$2"; shift 2 ;;
        --completion-column) COMPLETION_COLUMN="$2"; shift 2 ;;
        --projection-dim) PROJECTION_DIM="$2"; shift 2 ;;
        --precision) PRECISION="$2"; shift 2 ;;
        --hessian-dtype) HESSIAN_DTYPE="$2"; shift 2 ;;
        --token-batch) TOKEN_BATCH="$2"; shift 2 ;;
        --lambda-damp) LAMBDA_DAMP="$2"; shift 2 ;;
        --gpus) GPUS="$2"; shift 2 ;;
        --ev-correction) EV_CORRECTION="$2"; shift 2 ;;
        --ivhp-side) IVHP_SIDE="$2"; shift 2 ;;
        --python) PYTHONEXEC="$2"; shift 2 ;;
        --dry-run) DRY_RUN=true; shift ;;
        *) echo "Unknown option: $1" >&2; exit 1 ;;
    esac
done

REPO_DIR="$(cd "$(dirname "$0")/.." && pwd)"
PYTHONEXEC="${PYTHONEXEC:-$REPO_DIR/.venv/bin/python}"

for required in MODEL VALUE_DATA QUERY_DATA OUTPUT_DIR; do
    if [ -z "${!required}" ]; then
        echo "ERROR: --${required,,} is required" >&2
        exit 1
    fi
done

case "$IVHP_SIDE" in
    query | value) ;;
    *) echo "ERROR: --ivhp-side must be query or value" >&2; exit 1 ;;
esac

# Each gradient index is nested one level under a wrapper directory. The
# TrackStar scorer's --build-dir/--query-build-dir expect a parent whose
# children are the leaf indices (discover_shard_indices iterates children),
# so the leaves cannot sit directly under the run directory.
VALUE_DIR="$OUTPUT_DIR/value_index"
VALUE_GRADS="$VALUE_DIR/shard_0000"
QUERY_DIR="$OUTPUT_DIR/query_index"
QUERY_GRADS="$QUERY_DIR/q_0000"
HESSIAN_RUN="$OUTPUT_DIR/hessian"
HESSIAN_FACTORS="$HESSIAN_RUN/kfac"
IVHP_DIR="$OUTPUT_DIR/${IVHP_SIDE}_ivhp_index"
IVHP_OUT="$IVHP_DIR/ivhp_0000"

# Projected EK-FAC (corpus-scale). apply_hessian needs a FULL-dimensional input
# index (projection_dim 0) because it multiplies by the full-dim factors, then
# projects its OUTPUT into the other side's projection_dim subspace via
# --projection_config_path. So the IVHP side is built full-dim and the other side
# stays projected; stage D dot-scores both in the same projection_dim space.
if [ "$IVHP_SIDE" = query ]; then
    VALUE_PROJ="$PROJECTION_DIM"
    QUERY_PROJ=0
    PROJ_CONFIG_DIR="$VALUE_GRADS"
else
    VALUE_PROJ=0
    QUERY_PROJ="$PROJECTION_DIM"
    PROJ_CONFIG_DIR="$QUERY_GRADS"
fi

FSDP_ARGS=()
if [ "$GPUS" -gt 1 ]; then
    FSDP_ARGS=(--fsdp)
fi

EV_ARGS=()
if [ "$EV_CORRECTION" = true ]; then
    EV_ARGS=(--ev_correction)
fi

QUERY_PROMPT_ARGS=(--prompt_column "$PROMPT_COLUMN")
if [ -n "$COMPLETION_COLUMN" ]; then
    QUERY_PROMPT_ARGS+=(--completion_column "$COMPLETION_COLUMN")
fi

echo "=============================================="
echo "  EK-FAC attribution"
echo "=============================================="
echo "Model:       $MODEL"
echo "Value data:  $VALUE_DATA"
echo "Query data:  $QUERY_DATA"
echo "Output:      $OUTPUT_DIR"
echo "ProjDim:     $PROJECTION_DIM (projected EK-FAC: IVHP input full-dim, output projected)"
echo "Precision:   $PRECISION (hessian $HESSIAN_DTYPE)"
echo "GPUs:        $GPUS"
echo "EV corr:     $EV_CORRECTION   IVHP side: $IVHP_SIDE   lambda damp: $LAMBDA_DAMP"
echo "Python:      $PYTHONEXEC"
echo "=============================================="

run() {
    echo "+ $*"
    if [ "$DRY_RUN" = false ]; then
        "$@"
    fi
}

mkdir -p "$OUTPUT_DIR" "$VALUE_DIR" "$QUERY_DIR" "$IVHP_DIR"

# Stage A: raw gradient indices for value and query data.
run "$PYTHONEXEC" -m bergson build "$VALUE_GRADS" \
    --model "$MODEL" \
    --dataset json --data_args "data_files=$VALUE_DATA" \
    --prompt_column "$PROMPT_COLUMN" \
    --projection_dim "$VALUE_PROJ" \
    --precision "$PRECISION" \
    --token_batch_size "$TOKEN_BATCH" \
    --skip_preconditioners \
    --unit_normalize --truncation --overwrite \
    ${FSDP_ARGS[@]+"${FSDP_ARGS[@]}"}

run "$PYTHONEXEC" -m bergson build "$QUERY_GRADS" \
    --model "$MODEL" \
    --dataset json --data_args "data_files=$QUERY_DATA" \
    "${QUERY_PROMPT_ARGS[@]}" \
    --projection_dim "$QUERY_PROJ" \
    --precision "$PRECISION" \
    --token_batch_size "$TOKEN_BATCH" \
    --skip_preconditioners \
    --unit_normalize --truncation --overwrite \
    ${FSDP_ARGS[@]+"${FSDP_ARGS[@]}"}

# Stage B: EK-FAC factors from the value (training) distribution.
run "$PYTHONEXEC" -m bergson hessian "$HESSIAN_RUN" \
    --model "$MODEL" \
    --dataset json --data_args "data_files=$VALUE_DATA" \
    --prompt_column "$PROMPT_COLUMN" \
    --method kfac \
    --hessian_dtype "$HESSIAN_DTYPE" \
    --projection_dim "$PROJECTION_DIM" \
    --precision "$PRECISION" \
    --token_batch_size "$TOKEN_BATCH" \
    --truncation \
    --overwrite \
    ${EV_ARGS[@]+"${EV_ARGS[@]}"} \
    ${FSDP_ARGS[@]+"${FSDP_ARGS[@]}"}

# Stage C: apply IVHP to the chosen gradient side.
if [ "$IVHP_SIDE" = query ]; then
    IVHP_GRADS="$QUERY_GRADS"
else
    IVHP_GRADS="$VALUE_GRADS"
fi

APPLY_CMD=("$PYTHONEXEC")
if [ "$GPUS" -gt 1 ]; then
    APPLY_CMD=("$PYTHONEXEC" -m torch.distributed.run --nproc_per_node="$GPUS")
fi

run "${APPLY_CMD[@]}" -m bergson.hessians.apply_hessian \
    --hessian_method_path "$HESSIAN_FACTORS" \
    --gradient_path "$IVHP_GRADS" \
    --run_path "$IVHP_OUT" \
    --projection_config_path "$PROJ_CONFIG_DIR" \
    --lambda_damp_factor "$LAMBDA_DAMP"

# The projected apply_hessian output omits the processor files the dot-score
# loader (GradientProcessor.load) needs. Copy them from PROJ_CONFIG_DIR (same
# projection) so the IVHP index is loadable for stage D, until bergson's
# apply_hessian writes a self-contained index.
# NOTE: this reuses PROJ_CONFIG_DIR's unit-normalizer, which rescales per-query
# score magnitudes (rankings preserved). The correct long-term fix is an
# identity processor on the IVHP index.
echo "+ copy processor files from $PROJ_CONFIG_DIR into $IVHP_OUT"
if [ "$DRY_RUN" = false ]; then
    cp "$PROJ_CONFIG_DIR/processor_config.json" \
       "$PROJ_CONFIG_DIR/normalizers.pth" \
       "$PROJ_CONFIG_DIR/preconditioners.pth" \
       "$PROJ_CONFIG_DIR/preconditioners_eigen.pth" \
       "$PROJ_CONFIG_DIR/preprocess_config.json" \
       "$PROJ_CONFIG_DIR/total_processed.pt" \
       "$IVHP_OUT/"
fi

if [ "$IVHP_SIDE" = query ]; then
    SCORE_BUILD_DIR="$VALUE_DIR"
    SCORE_QUERY_DIR="$IVHP_DIR"
else
    SCORE_BUILD_DIR="$IVHP_DIR"
    SCORE_QUERY_DIR="$QUERY_DIR"
fi

echo "=============================================="
echo "  EK-FAC stages A-C complete"
echo "=============================================="
echo "Value index:     $VALUE_GRADS"
echo "Query index:     $QUERY_GRADS"
echo "EK-FAC factors:  $HESSIAN_FACTORS"
echo "IVHP ($IVHP_SIDE):   $IVHP_OUT"
echo ""
echo "Stage D (score) reuses the existing TrackStar dot-score path. The"
echo "--build-dir and --query-build-dir below are wrapper directories whose"
echo "children are the leaf indices, which is what the scorer expects:"
echo "  data-attribution-trackstar-dot-score \\"
echo "    --build-dir       $SCORE_BUILD_DIR \\"
echo "    --query-build-dir $SCORE_QUERY_DIR \\"
echo "    --variant base --run-id ekfac_\$(date -u +%Y%m%dT%H%M%SZ)"