Initial project setup: J-lens implementation for nanoGPT

- Core J-lens computation (batched per-layer, per-token gradient method)
- Project README with background and experiment plan
- Sync script for meru Docker container deployment
- Upstream nanoGPT model code copied to src/

Architecture: Computes d(log_p(token))/d(residual_stream) averaged
over corpus contexts, replicating Anthropic's Jacobian Lens technique.

7 files changed, 119 insertions, 0 deletions

diff --git a/src/config/eval_gpt2.py b/src/config/eval_gpt2.py
new file mode 100644
index 0000000..53978cb
--- /dev/null
+++ b/src/config/eval_gpt2.py
@@ -0,0 +1,8 @@
+# evaluate the base gpt2
+# n_layer=12, n_head=12, n_embd=768
+# 124M parameters
+batch_size = 8
+eval_iters = 500 # use more iterations to get good estimate
+eval_only = True
+wandb_log = False
+init_from = 'gpt2'
diff --git a/src/config/eval_gpt2_large.py b/src/config/eval_gpt2_large.py
new file mode 100644
index 0000000..4cbeaef
--- /dev/null
+++ b/src/config/eval_gpt2_large.py
@@ -0,0 +1,8 @@
+# evaluate the base gpt2
+# n_layer=36, n_head=20, n_embd=1280
+# 774M parameters
+batch_size = 8
+eval_iters = 500 # use more iterations to get good estimate
+eval_only = True
+wandb_log = False
+init_from = 'gpt2-large'
diff --git a/src/config/eval_gpt2_medium.py b/src/config/eval_gpt2_medium.py
new file mode 100644
index 0000000..9d0db11
--- /dev/null
+++ b/src/config/eval_gpt2_medium.py
@@ -0,0 +1,8 @@
+# evaluate the base gpt2
+# n_layer=24, n_head=16, n_embd=1024
+# 350M parameters
+batch_size = 8
+eval_iters = 500 # use more iterations to get good estimate
+eval_only = True
+wandb_log = False
+init_from = 'gpt2-medium'
diff --git a/src/config/eval_gpt2_xl.py b/src/config/eval_gpt2_xl.py
new file mode 100644
index 0000000..1bae34f
--- /dev/null
+++ b/src/config/eval_gpt2_xl.py
@@ -0,0 +1,8 @@
+# evaluate the base gpt2
+# n_layer=48, n_head=25, n_embd=1600
+# 1558M parameters
+batch_size = 8
+eval_iters = 500 # use more iterations to get good estimate
+eval_only = True
+wandb_log = False
+init_from = 'gpt2-xl'
diff --git a/src/config/finetune_shakespeare.py b/src/config/finetune_shakespeare.py
new file mode 100644
index 0000000..148a4c4
--- /dev/null
+++ b/src/config/finetune_shakespeare.py
@@ -0,0 +1,25 @@
+import time
+
+out_dir = 'out-shakespeare'
+eval_interval = 5
+eval_iters = 40
+wandb_log = False # feel free to turn on
+wandb_project = 'shakespeare'
+wandb_run_name = 'ft-' + str(time.time())
+
+dataset = 'shakespeare'
+init_from = 'gpt2-xl' # this is the largest GPT-2 model
+
+# only save checkpoints if the validation loss improves
+always_save_checkpoint = False
+
+# the number of examples per iter:
+# 1 batch_size * 32 grad_accum * 1024 tokens = 32,768 tokens/iter
+# shakespeare has 301,966 tokens, so 1 epoch ~= 9.2 iters
+batch_size = 1
+gradient_accumulation_steps = 32
+max_iters = 20
+
+# finetune at constant LR
+learning_rate = 3e-5
+decay_lr = False
diff --git a/src/config/train_gpt2.py b/src/config/train_gpt2.py
new file mode 100644
index 0000000..8f19273
--- /dev/null
+++ b/src/config/train_gpt2.py
@@ -0,0 +1,25 @@
+# config for training GPT-2 (124M) down to very nice loss of ~2.85 on 1 node of 8X A100 40GB
+# launch as the following (e.g. in a screen session) and wait ~5 days:
+# $ torchrun --standalone --nproc_per_node=8 train.py config/train_gpt2.py
+
+wandb_log = True
+wandb_project = 'owt'
+wandb_run_name='gpt2-124M'
+
+# these make the total batch size be ~0.5M
+# 12 batch size * 1024 block size * 5 gradaccum * 8 GPUs = 491,520
+batch_size = 12
+block_size = 1024
+gradient_accumulation_steps = 5 * 8
+
+# this makes total number of tokens be 300B
+max_iters = 600000
+lr_decay_iters = 600000
+
+# eval stuff
+eval_interval = 1000
+eval_iters = 200
+log_interval = 10
+
+# weight decay
+weight_decay = 1e-1
diff --git a/src/config/train_shakespeare_char.py b/src/config/train_shakespeare_char.py
new file mode 100644
index 0000000..41c81df
--- /dev/null
+++ b/src/config/train_shakespeare_char.py
@@ -0,0 +1,37 @@
+# train a miniature character-level shakespeare model
+# good for debugging and playing on macbooks and such
+
+out_dir = 'out-shakespeare-char'
+eval_interval = 250 # keep frequent because we'll overfit
+eval_iters = 200
+log_interval = 10 # don't print too too often
+
+# we expect to overfit on this small dataset, so only save when val improves
+always_save_checkpoint = False
+
+wandb_log = False # override via command line if you like
+wandb_project = 'shakespeare-char'
+wandb_run_name = 'mini-gpt'
+
+dataset = 'shakespeare_char'
+gradient_accumulation_steps = 1
+batch_size = 64
+block_size = 256 # context of up to 256 previous characters
+
+# baby GPT model :)
+n_layer = 6
+n_head = 6
+n_embd = 384
+dropout = 0.2
+
+learning_rate = 1e-3 # with baby networks can afford to go a bit higher
+max_iters = 5000
+lr_decay_iters = 5000 # make equal to max_iters usually
+min_lr = 1e-4 # learning_rate / 10 usually
+beta2 = 0.99 # make a bit bigger because number of tokens per iter is small
+
+warmup_iters = 100 # not super necessary potentially
+
+# on macbook also add
+# device = 'cpu'  # run on cpu only
+# compile = False # do not torch compile the model