summaryrefslogtreecommitdiff
path: root/src/config/train_gpt2.py
diff options
context:
space:
mode:
authorVoid Agent <void@jayrup.hermes>2026-07-29 17:42:47 +0100
committerVoid Agent <void@jayrup.hermes>2026-07-29 17:42:47 +0100
commit66f99ee30087a5f28ad852e581a0334c7f556091 (patch)
treeb83da26e340a926e6eb46f1b304b66f693a15e81 /src/config/train_gpt2.py
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.
Diffstat (limited to 'src/config/train_gpt2.py')
-rw-r--r--src/config/train_gpt2.py25
1 files changed, 25 insertions, 0 deletions
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