train_sae.py File Reference

Go to the source code of this file.

Namespaces
namespace	aquin
namespace	aquin.compute
namespace	aquin.compute.train_sae

Functions
	aquin.compute.train_sae.chunk_paths (int layer)
Path	aquin.compute.train_sae.norm_cache_path (int layer)
Path	aquin.compute.train_sae.sae_save_path (int layer)
	aquin.compute.train_sae.collect_activations (model, int layer)
	aquin.compute.train_sae.compute_norm (int layer)
	aquin.compute.train_sae.train_layer (int layer)

Variables
str	aquin.compute.train_sae.MODEL_NAME = "meta-llama/Llama-3.2-1B-Instruct"
int	aquin.compute.train_sae.N_LAYERS = 16
int	aquin.compute.train_sae.N_FEATURES = 32768
int	aquin.compute.train_sae.D_MODEL = 2048
float	aquin.compute.train_sae.L1_COEFF = 10.0
int	aquin.compute.train_sae.LR = 1e-4
int	aquin.compute.train_sae.BATCH_SIZE = 4096
int	aquin.compute.train_sae.N_TOKENS = 2_000_000
int	aquin.compute.train_sae.SEQ_LEN = 64
int	aquin.compute.train_sae.COLLECT_BATCH = 32
int	aquin.compute.train_sae.CHUNK_SIZE = 100_000
int	aquin.compute.train_sae.SAVE_EVERY = 2000
str	aquin.compute.train_sae.SAE_DIR = Path(__file__).parent / "sae"
str	aquin.compute.train_sae.ACTS_DIR = Path(__file__).parent / "sae" / "acts"
	aquin.compute.train_sae.DEVICE = resolve_compute_device()
	aquin.compute.train_sae.parser = argparse.ArgumentParser()
	aquin.compute.train_sae.type
	aquin.compute.train_sae.str
	aquin.compute.train_sae.default
	aquin.compute.train_sae.help
	aquin.compute.train_sae.args = parser.parse_args()
str	aquin.compute.train_sae.layers = "all" else [int(x) for x in args.layers.split(",")]
	aquin.compute.train_sae.flush
	aquin.compute.train_sae.model = HookedTransformer.from_pretrained(MODEL_NAME)