File size: 2,752 Bytes
a2ffd07 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 | #!/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"
|