#!/bin/bash

# ============================================================
# Multi-GPU DDP template for training LLaVA SAEs
# ============================================================
#
# Differences from single-GPU script:
#   1. --device_id accepts multiple GPU IDs (e.g. 0 1 2 3)
#   2. Lightning auto-launches DDP processes (one per GPU)
#   3. Each process creates the model on its own GPU via LOCAL_RANK
#   4. Effective batch size = batch_size * num_gpus
#      Adjust batch_size or lr accordingly.
#
# Usage:
#   bash training/scripts/Train_Llava_SAE_DDP.sh
# ============================================================

# Define the parameters
model_name="llava-hf/llava-1.5-7b-hf"
tok_name="llava-hf/llava-1.5-7b-hf"
batch_size=2                    # images per GPU (8 GPUs = 16 images/step)
max_tokens=5000000000000
sae_name="batchtopk"
lr=1e-4
expansion_factor=16
k=8
auxk=256
auxk_coef=0.03125
max_epochs=3
dead_tokens_threshold=10000000
log_every_n_steps=50
save_every_n_training_steps=1000
save_step=false
use_loss_var=true
num_workers=4
hf_dataset="pixparse/cc3m-wds"
dtype="bfloat16"
project="llava-sae"
question="Describe this image."
max_new_tokens=128

# ----- Multi-GPU config -----
# List the GPU IDs you want to use, space-separated.
device_id="0 1 2 3 4 5 6 7"

# Data paths (change as needed)
local_train_path="../hallucination/CC3M-Dataset/preproc/CC3M-Dataset/cc3m_images/train"
local_val_path="../hallucination/CC3M-Dataset/preproc/CC3M-Dataset/cc3m_images/val"

# Navigate to repo root (parent of training/scripts/)
cd "$(dirname "$0")/../.."

export HF_HOME=~/scratch/hf_home

if [ -f .env ]; then
    set -a
    source .env
    set +a
fi

if [ -z "${WANDB_API_KEY:-}" ]; then
    echo "Error: WANDB_API_KEY is not set. Add it to .env before running." >&2
    exit 1
fi

wandb login --relogin "${WANDB_API_KEY}"
python -m training.Train_Multilayer_SAE_Llava \
    --model_name "$model_name" \
    --tok_name "$tok_name" \
    --batch_size "$batch_size" \
    --sae_name "$sae_name" \
    --lr "$lr" \
    --expansion_factor "$expansion_factor" \
    --k "$k" \
    --auxk "$auxk" \
    --max_epochs "$max_epochs" \
    --dead_tokens_threshold "$dead_tokens_threshold" \
    --log_every_n_steps "$log_every_n_steps" \
    --save_every_n_training_steps "$save_every_n_training_steps" \
    --save_step "$save_step" \
    --use_loss_var "$use_loss_var" \
    --max_tokens "$max_tokens" \
    --auxk_coef "$auxk_coef" \
    --num_workers "$num_workers" \
    --hf_dataset "$hf_dataset" \
    --local_train_path "$local_train_path" \
    --local_val_path "$local_val_path" \
    --project "$project" \
    --question "$question" \
    --max_new_tokens "$max_new_tokens" \
    --device_id $device_id \
    --dtype "$dtype"