File size: 4,734 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 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 | import torch as t
import argparse
from typing import List, Tuple, Dict, Any, Union, Callable, cast
from sae.SAE_Trainer import train, TrainSAEConfig, DataConfig
from sae.Load_Data import load_data
from itertools import product
from functools import partial
from transformer_lens.hook_points import HookPoint
from transformer_lens import HookedTransformer
from sae.Training_Utils import *
def preproc(
toks: Tuple[t.Tensor],
input_hook: str,
output_hook: str,
transformer: HookedTransformer,
device: t.device | str = "cpu",
) -> Tuple[t.Tensor, t.Tensor]:
act_cache = {}
def hook_fn(inputs: t.Tensor, hook: HookPoint):
act_cache[hook.name] = inputs.detach()
transformer.run_with_hooks(
toks[0],
fwd_hooks=[(lambda name: name == input_hook or name == output_hook, hook_fn)],
)
return act_cache[input_hook], act_cache[output_hook]
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="Train Downstream SAE.")
parser.add_argument('--model_name', type=str, default="gpt2-small", help="The model to train the SAE on.")
parser.add_argument('--tok_name', type=str, default="gpt2", help="The tokenizer to encode the data.")
parser.add_argument('--batch_size', type=int, default=32, help='Batch size for training')
parser.add_argument('--max_tokens', type=int, default=1_000_000_000, help='Maximum tokens for training')
parser.add_argument('--layer', type=int, default=11, help='Layer to use for metric calculation')
parser.add_argument('--input_hook', type=str, required=True, help='Input hook name to use in training')
parser.add_argument('--output_hook', type=str, required=True, help='Output hook name to use in training')
parser.add_argument('--sae_name', type=str, default="topk_transcoder", help='SAE name')
parser.add_argument('--lr', type=float, default=1e-4, help='Learning rate')
parser.add_argument('--expansion_factor', type=float, default=0.5, help='Expansion factor')
parser.add_argument('--k', type=int, default=8, help='K value')
parser.add_argument('--auxk', type=int, default=32, help='Auxiliary K value')
parser.add_argument('--auxk_coef', type=float, default=0.05, help='Auxiliary K coefficient')
parser.add_argument('--use_loss_var', type=str_to_bool, default=True, help='Use variance in the loss.')
parser.add_argument('--device_id', type=int, default=7, help='Device IDs')
parser.add_argument('--max_epochs', type=int, default=5, help='Maximum number of epochs')
parser.add_argument('--dead_tokens_threshold', type=int, default=10_000_000, help='Dead tokens threshold')
parser.add_argument('--log_every_n_steps', type=int, default=32, help='Log every N steps')
parser.add_argument('--patience', type=int, default=1, help='patience')
parser.add_argument('--num_workers', type=int, default=63, help='Number of workers')
args = parser.parse_args()
device = t.device(f"cuda:{int(args.device_id)}" if t.cuda.is_available() else "cpu")
model = HookedTransformer.from_pretrained(args.model_name, device=device)
# default config
data_config = DataConfig(
batch_size=args.batch_size,
max_tokens=args.max_tokens,
max_length=1024, # the tokenizer of gpt2 only allow max 1024 tokens
num_workers=args.num_workers,
)
hook_names = [args.output_hook, args.input_hook]
default_config = TrainSAEConfig(
sae_name=args.sae_name,
lr=args.lr,
hook_names=hook_names,
expansion_factor=args.expansion_factor,
k=args.k,
auxk=args.auxk,
auxk_coef=args.auxk_coef,
device_id=[args.device_id],
max_epochs=args.max_epochs,
dead_tokens_threshold=args.dead_tokens_threshold,
log_every_n_steps=args.log_every_n_steps,
patience=args.patience,
use_loss_var=args.use_loss_var,
monitor="loss/total",
)
default_config.save_file_name = "{}_{}_{}_{}_{}_{}_{}_{}".format(
default_config.sae_name,
default_config.expansion_factor,
default_config.k,
hook_names[0],
default_config.lr,
default_config.auxk,
default_config.auxk_coef,
default_config.seed
)
# load data
train_loader, val_loader = load_data(args.tok_name, data_config)
preprocess = partial(
preproc,
input_hook=args.input_hook,
output_hook=args.output_hook,
transformer=model,
device=device
)
train(
default_config,
data_config,
model,
train_loader,
val_loader,
preprocess=preprocess,
train_mode="transcoder",
)
|