20 files changed, 1421 insertions, 0 deletions
diff --git a/.gitignore b/.gitignore
new file mode 100644
index 0000000..91b3024
--- /dev/null
+++ b/.gitignore
@@ -0,0 +1,14 @@
+__pycache__/
+*.pyc
+*.pt
+*.pth
+*.bin
+*.safetensors
+outputs/
+data/shakespeare/
+data/shakespeare_char/
+*.egg-info/
+.venv/
+.ipynb_checkpoints/
+upstream-nanogpt/
+*.log
diff --git a/README.md b/README.md
new file mode 100644
index 0000000..e244b24
--- /dev/null
+++ b/README.md
@@ -0,0 +1,29 @@
+# J-space on nanoGPT
+
+Replicating Anthropic's Jacobian Lens ("J-space") technique on a small character-level transformer (nanoGPT).
+
+## Hypothesis
+
+If J-space (the subspace of representations readily available for verbal report) is an **architectural/structural property of transformers** rather than an emergent feature of advanced models, it should appear at all scales — including 10M-parameter char-level models.
+
+## Background
+
+Anthropic's 2026 paper "Verbalizable Representations Form a Global Workspace in Language Models" introduces the Jacobian Lens (J-lens), which computes the average linearized effect of activations on future token probabilities, averaged over many contexts. This reveals a privileged "J-space" of representations that the model can report on, modulate, and use for reasoning.
+
+Full paper: https://transformer-circuits.pub/2026/workspace/index.html
+
+## Experiments
+
+1. **J-space visualization** — Compute J-lens vectors for all vocabulary tokens at each layer. Visualize which characters/concepts enter "verbalizable space" and when.
+2. **Ablation test** — Remove J-space components vs random directions vs full activations. Measure prediction quality impact.
+3. **Training dynamics** — Save checkpoints during training, compute J-space at each, track when it crystallizes.
+4. **Capacity measurement** — Count active J-lens tokens per position.
+
+## Setup
+
+Runs on meru's Quadro K2200 (4GB VRAM) via Docker with GPU passthrough.
+
+## References
+
+- [Verbalizable Representations Form a Global Workspace in Language Models](https://transformer-circuits.pub/2026/workspace/index.html) — Anthropic, 2026
+- [nanoGPT](https://github.com/karpathy/nanoGPT) — Andrej Karpathy
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
diff --git a/src/eval_gpt2.py b/src/eval_gpt2.py
new file mode 100644
index 0000000..53978cb
--- /dev/null
+++ b/src/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/eval_gpt2_large.py b/src/eval_gpt2_large.py
new file mode 100644
index 0000000..4cbeaef
--- /dev/null
+++ b/src/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/eval_gpt2_medium.py b/src/eval_gpt2_medium.py
new file mode 100644
index 0000000..9d0db11
--- /dev/null
+++ b/src/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/eval_gpt2_xl.py b/src/eval_gpt2_xl.py
new file mode 100644
index 0000000..1bae34f
--- /dev/null
+++ b/src/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/finetune_shakespeare.py b/src/finetune_shakespeare.py
new file mode 100644
index 0000000..148a4c4
--- /dev/null
+++ b/src/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/jlens.py b/src/jlens.py
new file mode 100644
index 0000000..515dc08
--- /dev/null
+++ b/src/jlens.py
@@ -0,0 +1,441 @@
+"""
+J-lens: Jacobian Lens for Transformer Models
+
+Replicates Anthropic's technique from:
+"Verbalizable Representations Form a Global Workspace in Language Models"
+https://transformer-circuits.pub/2026/workspace/index.html
+
+Core idea: For each token in the vocabulary, compute the average gradient
+of log p(token) with respect to the residual stream at each layer,
+averaged over many contexts. This reveals which concepts are "verbalizable"
+— readily available for the model to report on.
+
+Usage:
+    python jlens.py --model checkpoints/ckpt.pt --data data/shakespeare_char
+"""
+
+import torch
+import torch.nn as nn
+from torch.utils.data import DataLoader
+import numpy as np
+import argparse
+import json
+import os
+import pickle
+from pathlib import Path
+from collections import defaultdict
+
+
+def load_model(checkpoint_path, model_class, device='cuda'):
+    """Load a trained nanoGPT model from checkpoint."""
+    checkpoint = torch.load(checkpoint_path, map_location=device)
+    # nanoGPT stores model args, state_dict + optimizer in checkpoint
+    model_args = checkpoint['model_args']
+    
+    # Create model with saved config
+    model = model_class(model_args)
+    
+    # Fix state dict keys (nanoGPT wraps in DataParallel)
+    state_dict = checkpoint['model']
+    unwanted_prefix = '_orig_mod.'
+    for k in list(state_dict.keys()):
+        if k.startswith(unwanted_prefix):
+            state_dict[k[len(unwanted_prefix):]] = state_dict.pop(k)
+    
+    model.load_state_dict(state_dict)
+    model.to(device)
+    model.eval()
+    return model, model_args
+
+
+def compute_jlens_single_token(model, token_id, dataloader, layer_idx, device='cuda'):
+    """
+    Compute J-lens vector for a single token at a specific layer.
+    
+    J_l(token, layer) = E_x [ ∇_{resid[layer]} log p(token | x) ]
+    
+    Where the expectation is taken over all positions in the corpus.
+    """
+    vectors = []
+    
+    with torch.no_grad():
+        for batch_idx, (x, y) in enumerate(dataloader):
+            x, y = x.to(device), y.to(device)
+            B, T = x.shape
+            
+            # We need gradients, so we'll do forward passes with hooks
+            # Strategy: use torch.autograd.grad on a forward pass
+            # where we capture residual stream activations
+
+            # Register hook to capture residual stream at target layer
+            activations = {}
+            
+            def make_hook():
+                def hook(module, input, output):
+                    # output is (B, T, n_embd)
+                    # Detach then require grad so we can compute gradient through it
+                    activations['resid'] = output.detach().requires_grad_(True)
+                    return activations['resid']
+                return hook
+            
+            # Find the target layer
+            target_block = model.transformer.h[layer_idx]
+            # nanoGPT architecture: h = x + attn(ln1(x)), then x = h + mlp(ln2(h))
+            # We want the residual stream AFTER the attention + MLP of this layer
+            # which is the output of the block
+            
+            handle = target_block.register_forward_hook(make_hook())
+            
+            # Forward pass
+            logits, loss = model(x, y)
+            
+            handle.remove()
+            
+            # Now compute gradient of log p(token_id) w.r.t. residual stream
+            # log p(token_id) at each position = log_softmax(logits)[:, :, token_id]
+            log_probs = torch.nn.functional.log_softmax(logits, dim=-1)
+            token_log_probs = log_probs[:, :, token_id].sum()  # sum over B, T
+            
+            # Gradient of this sum w.r.t. the captured activations
+            grad = torch.autograd.grad(
+                token_log_probs, 
+                activations['resid'],
+                retain_graph=False
+            )[0]  # Shape: (B, T, n_embd)
+            
+            vectors.append(grad.detach().cpu())
+            
+            # Cleanup
+            del logits, loss, log_probs, activations, grad
+            torch.cuda.empty_cache()
+    
+    # Average over all positions in the corpus
+    all_vectors = torch.cat([v.reshape(-1, v.shape[-1]) for v in vectors], dim=0)
+    jlens_vector = all_vectors.mean(dim=0)  # Shape: (n_embd,)
+    
+    return jlens_vector
+
+
+def compute_jlens_all_tokens(model, dataloader, layer_idx, vocab_size, device='cuda'):
+    """
+    Compute J-lens vectors for all tokens at a specific layer.
+    
+    Returns: dict mapping token_id -> jlens_vector (n_embd,)
+    """
+    jlens_vectors = {}
+    
+    for token_id in range(vocab_size):
+        vec = compute_jlens_single_token(model, token_id, dataloader, layer_idx, device)
+        jlens_vectors[token_id] = vec
+        
+        if (token_id + 1) % 10 == 0:
+            print(f"  Token {token_id + 1}/{vocab_size} done")
+    
+    return jlens_vectors
+
+
+def compute_jlens_all_layers(model, dataloader, n_layers, vocab_size, device='cuda', 
+                              use_batched=True):
+    """
+    Compute J-lens vectors for all layers and all tokens.
+    
+    Uses batched approach: for each context, compute gradients for ALL tokens
+    at once using vector-Jacobian products. Much faster than per-token.
+    
+    Returns: dict mapping layer_idx -> {token_id: jlens_vector}
+    """
+    all_layer_vectors = defaultdict(dict)
+    
+    if use_batched:
+        # Optimized: compute all token J-lens vectors simultaneously
+        # For each context position, the gradient of log p(token) w.r.t. resid
+        # for all tokens is just the Jacobian of the unembedding layer
+        # which equals W_U^T * (one_hot(token) - softmax(logits))
+        # Wait, let me think about this more carefully...
+        
+        print("Using batched J-lens computation...")
+        
+        for layer_idx in range(n_layers):
+            print(f"\nLayer {layer_idx}/{n_layers}...")
+            
+            layer_accum = torch.zeros(vocab_size, model.config.n_embd, device='cpu')
+            token_count = torch.zeros(vocab_size, device='cpu')
+            
+            with torch.no_grad():
+                for batch_idx, (x, y) in enumerate(dataloader):
+                    x, y = x.to(device), y.to(device)
+                    B, T = x.shape
+                    
+                    # Capture residual stream at target layer
+                    resid_captured = {}
+                    
+                    def make_hook(resid_dict):
+                        def hook(module, input, output):
+                            resid_dict['val'] = output.detach().requires_grad_(True)
+                            return resid_dict['val']
+                        return hook
+                    
+                    target_block = model.transformer.h[layer_idx]
+                    handle = target_block.register_forward_hook(make_hook(resid_captured))
+                    
+                    logits, loss = model(x, y)
+                    handle.remove()
+                    
+                    # Now: for each token in vocab, we want d(logit_t)/d(resid)
+                    # This is the Jacobian of unembedding w.r.t. residual stream
+                    # Chain rule: d(logit_t)/d(resid) = W_U[t, :] * d(layer_out)/d(resid)
+                    # where layer_out is the final layer output after all remaining layers
+                    # plus the direct path through the residual stream.
+                    
+                    # Actually, since we captured resid at layer L, and the model
+                    # applies remaining layers resid_L -> ... -> resid_final -> logits,
+                    # the gradient d(logits)/d(resid_L) = d(logits)/d(resid_final) * d(resid_final)/d(resid_L)
+                    # 
+                    # We can compute this by:
+                    # 1. Get logits
+                    # 2. For EACH position, compute gradient of logit for EACH token
+                    #    w.r.t. the captured residual stream
+                    # 3. Average across positions
+                    
+                    # Vectorized approach: compute gradients for ALL tokens simultaneously
+                    # using torch.autograd.grad with list of outputs
+                    
+                    # For efficiency, compute per position, then aggregate
+                    log_probs = torch.nn.functional.log_softmax(logits, dim=-1)  # (B, T, vocab)
+                    
+                    # For each position (b, t), we need jacobian of log_probs[b,t,:] w.r.t. resid[b,t,:]
+                    # This is (vocab, n_embd) per position
+                    # We can batch by computing gradient of sum_{tokens} a_i * log_p(token_i)
+                    # where a_i cycles through standard basis vectors
+                    
+                    # Practical approach for small vocab (nanoGPT: 65 tokens):
+                    # Just loop over tokens, compute gradient, and accumulate
+                    
+                    resid = resid_captured['val']  # (B, T, n_embd)
+                    
+                    for token_id in range(vocab_size):
+                        # Gradient of log_p(token_id) summed over all positions
+                        token_log_prob = log_probs[:, :, token_id].sum()
+                        
+                        grad = torch.autograd.grad(
+                            token_log_prob, resid, retain_graph=(token_id < vocab_size - 1)
+                        )[0]  # (B, T, n_embd)
+                        
+                        # Accumulate: sum of gradients across all positions
+                        layer_accum[token_id] += grad.detach().cpu().reshape(-1, model.config.n_embd).sum(dim=0)
+                        token_count[token_id] += B * T
+                    
+                    del logits, loss, log_probs, resid
+                    del grad  # pyright: ignore[reportPossiblyUnboundVariable]
+                    torch.cuda.empty_cache()
+                    
+                    if (batch_idx + 1) % 10 == 0:
+                        print(f"  Batch {batch_idx + 1}/{len(dataloader)}")
+            
+            # Average: divide sum by count
+            for token_id in range(vocab_size):
+                if token_count[token_id] > 0:
+                    all_layer_vectors[layer_idx][token_id] = layer_accum[token_id] / token_count[token_id]
+                else:
+                    all_layer_vectors[layer_idx][token_id] = torch.zeros(model.config.n_embd)
+            
+            print(f"  Layer {layer_idx} complete. Saved {vocab_size} token vectors.")
+    
+    return dict(all_layer_vectors)
+
+
+def save_jlens(jlens_data, output_path, metadata=None):
+    """Save J-lens vectors to disk."""
+    output = {
+        'metadata': metadata or {},
+        'vectors': {
+            str(layer): {
+                str(token_id): vec.numpy() for token_id, vec in tokens.items()
+            }
+            for layer, tokens in jlens_data.items()
+        }
+    }
+    
+    os.makedirs(os.path.dirname(output_path), exist_ok=True)
+    with open(output_path, 'wb') as f:
+        pickle.dump(output, f)
+    
+    print(f"Saved J-lens data to {output_path}")
+
+
+def load_jlens(path):
+    """Load saved J-lens vectors."""
+    with open(path, 'rb') as f:
+        data = pickle.load(f)
+    
+    # Convert back to tensors
+    jlens = {}
+    for layer_str, tokens in data['vectors'].items():
+        layer = int(layer_str)
+        jlens[layer] = {}
+        for token_id_str, vec in tokens.items():
+            jlens[layer][int(token_id_str)] = torch.from_numpy(vec)
+    
+    return jlens, data['metadata']
+
+
+def analyze_jlens(jlens_data, itos, n_layers, output_dir='outputs'):
+    """Analyze and visualize J-lens vectors."""
+    os.makedirs(output_dir, exist_ok=True)
+    vocab_size = len(itos)
+    
+    print(f"\n{'='*60}")
+    print("J-LENS ANALYSIS")
+    print(f"{'='*60}")
+    
+    for layer_idx in range(n_layers):
+        if layer_idx not in jlens_data:
+            continue
+        
+        layer_vectors = jlens_data[layer_idx]
+        
+        # Compute norm of each token's J-lens vector
+        norms = {}
+        for token_id, vec in layer_vectors.items():
+            norms[token_id] = vec.norm().item()
+        
+        # Sort by norm (most "verbalizable" tokens first)
+        sorted_tokens = sorted(norms.items(), key=lambda x: x[1], reverse=True)
+        
+        print(f"\n--- Layer {layer_idx} ---")
+        print(f"Top 10 most verbalizable tokens:")
+        for token_id, norm in sorted_tokens[:10]:
+            token_str = itos[token_id].replace('\n', '\\n')
+            print(f"  '{token_str}': norm={norm:.4f}")
+        
+        print(f"Bottom 5 least verbalizable tokens:")
+        for token_id, norm in sorted_tokens[-5:]:
+            token_str = itos[token_id].replace('\n', '\\n')
+            print(f"  '{token_str}': norm={norm:.4f}")
+    
+    # Compute J-space "capacity" — how many tokens have significant norm?
+    print(f"\n--- J-space Capacity ---")
+    for layer_idx in range(n_layers):
+        if layer_idx not in jlens_data:
+            continue
+        layer_vectors = jlens_data[layer_idx]
+        norms = torch.tensor([v.norm().item() for v in layer_vectors.values()])
+        
+        # Count "active" tokens (norm > median * 2)
+        threshold = norms.median() * 2
+        active = (norms > threshold).sum().item()
+        print(f"  Layer {layer_idx}: {active}/{vocab_size} tokens active (threshold={threshold:.4f})")
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description='J-lens: Jacobian Lens for nanoGPT')
+    parser.add_argument('--checkpoint', type=str, required=True, 
+                        help='Path to model checkpoint')
+    parser.add_argument('--data_dir', type=str, default='data/shakespeare_char',
+                        help='Path to data directory')
+    parser.add_argument('--output_dir', type=str, default='outputs/jlens',
+                        help='Directory for saving outputs')
+    parser.add_argument('--batch_size', type=int, default=32,
+                        help='Batch size for processing')
+    parser.add_argument('--max_batches', type=int, default=100,
+                        help='Max batches to process (limit for speed)')
+    parser.add_argument('--layers', type=str, default=None,
+                        help='Comma-separated layer indices (default: all)')
+    parser.add_argument('--device', type=str, default='cuda',
+                        help='Device to use')
+    
+    args = parser.parse_args()
+    
+    # Import model from local src
+    import sys
+    sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'src'))
+    from model import GPT, GPTConfig
+    
+    # Load model
+    print(f"Loading model from {args.checkpoint}")
+    model, model_args = load_model(args.checkpoint, GPT, args.device)
+    print(f"Model: {model_args.n_layer} layers, {model_args.n_embd} dim, "
+          f"{model_args.n_head} heads, {model_args.vocab_size} vocab")
+    
+    # Load data
+    data_dir = Path(args.data_dir)
+    train_data = np.memmap(data_dir / 'train.bin', dtype=np.uint16, mode='r')
+    val_data = np.memmap(data_dir / 'val.bin', dtype=np.uint16, mode='r')
+    
+    # Load vocab mappings
+    meta_path = data_dir / 'meta.pkl'
+    if meta_path.exists():
+        with open(meta_path, 'rb') as f:
+            meta = pickle.load(f)
+        itos = meta['itos']
+        stoi = meta['stoi']
+    else:
+        # Default char-level vocab
+        chars = sorted(list(set(open(data_dir / 'input.txt').read())))
+        stoi = {ch: i for i, ch in enumerate(chars)}
+        itos = {i: ch for i, ch in enumerate(chars)}
+    
+    print(f"Vocabulary size: {len(itos)}")
+    print(f"Train data: {len(train_data):,} tokens")
+    
+    # Create dataloader
+    def get_batch(split):
+        data = train_data if split == 'train' else val_data
+        block_size = model_args.block_size
+        ix = torch.randint(len(data) - block_size, (args.batch_size,))
+        x = torch.stack([torch.from_numpy((data[i:i+block_size]).astype(np.int64)) 
+                         for i in ix])
+        y = torch.stack([torch.from_numpy((data[i+1:i+1+block_size]).astype(np.int64)) 
+                         for i in ix])
+        return x, y
+    
+    class SimpleDataset(torch.utils.data.IterableDataset):
+        def __iter__(self):
+            while True:
+                yield get_batch('train')
+    
+    dataset = SimpleDataset()
+    dataloader = DataLoader(dataset, batch_size=None, num_workers=0)
+    
+    # Limit to max_batches
+    limited_dataloader = []
+    for i, batch in enumerate(dataloader):
+        if i >= args.max_batches:
+            break
+        limited_dataloader.append(batch)
+    
+    print(f"Processing {len(limited_dataloader)} batches of size {args.batch_size}")
+    
+    # Determine layers to process
+    if args.layers:
+        layers_to_process = [int(l) for l in args.layers.split(',')]
+    else:
+        layers_to_process = list(range(model_args.n_layer))
+    
+    print(f"Computing J-lens for layers: {layers_to_process}")
+    
+    # Compute J-lens for selected layers
+    jlens_data = {}
+    for layer_idx in layers_to_process:
+        print(f"\nComputing J-lens for layer {layer_idx}...")
+        layer_vectors = compute_jlens_all_tokens(
+            model, limited_dataloader, layer_idx, 
+            model_args.vocab_size, args.device
+        )
+        jlens_data[layer_idx] = layer_vectors
+    
+    # Save results
+    save_path = os.path.join(args.output_dir, 'jlens_vectors.pkl')
+    metadata = {
+        'model_args': vars(model_args),
+        'num_batches': len(limited_dataloader),
+        'batch_size': args.batch_size,
+        'layers_processed': layers_to_process,
+        'vocab_size': model_args.vocab_size,
+    }
+    save_jlens(jlens_data, save_path, metadata)
+    
+    # Analyze
+    analyze_jlens(jlens_data, itos, model_args.n_layer, args.output_dir)
+    
+    print(f"\nDone! Results saved to {args.output_dir}")
diff --git a/src/model.py b/src/model.py
new file mode 100644
index 0000000..c698f8b
--- /dev/null
+++ b/src/model.py
@@ -0,0 +1,330 @@
+"""
+Full definition of a GPT Language Model, all of it in this single file.
+References:
+1) the official GPT-2 TensorFlow implementation released by OpenAI:
+https://github.com/openai/gpt-2/blob/master/src/model.py
+2) huggingface/transformers PyTorch implementation:
+https://github.com/huggingface/transformers/blob/main/src/transformers/models/gpt2/modeling_gpt2.py
+"""
+
+import math
+import inspect
+from dataclasses import dataclass
+
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+
+class LayerNorm(nn.Module):
+    """ LayerNorm but with an optional bias. PyTorch doesn't support simply bias=False """
+
+    def __init__(self, ndim, bias):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(ndim))
+        self.bias = nn.Parameter(torch.zeros(ndim)) if bias else None
+
+    def forward(self, input):
+        return F.layer_norm(input, self.weight.shape, self.weight, self.bias, 1e-5)
+
+class CausalSelfAttention(nn.Module):
+
+    def __init__(self, config):
+        super().__init__()
+        assert config.n_embd % config.n_head == 0
+        # key, query, value projections for all heads, but in a batch
+        self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd, bias=config.bias)
+        # output projection
+        self.c_proj = nn.Linear(config.n_embd, config.n_embd, bias=config.bias)
+        # regularization
+        self.attn_dropout = nn.Dropout(config.dropout)
+        self.resid_dropout = nn.Dropout(config.dropout)
+        self.n_head = config.n_head
+        self.n_embd = config.n_embd
+        self.dropout = config.dropout
+        # flash attention make GPU go brrrrr but support is only in PyTorch >= 2.0
+        self.flash = hasattr(torch.nn.functional, 'scaled_dot_product_attention')
+        if not self.flash:
+            print("WARNING: using slow attention. Flash Attention requires PyTorch >= 2.0")
+            # causal mask to ensure that attention is only applied to the left in the input sequence
+            self.register_buffer("bias", torch.tril(torch.ones(config.block_size, config.block_size))
+                                        .view(1, 1, config.block_size, config.block_size))
+
+    def forward(self, x):
+        B, T, C = x.size() # batch size, sequence length, embedding dimensionality (n_embd)
+
+        # calculate query, key, values for all heads in batch and move head forward to be the batch dim
+        q, k, v  = self.c_attn(x).split(self.n_embd, dim=2)
+        k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
+        q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
+        v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
+
+        # causal self-attention; Self-attend: (B, nh, T, hs) x (B, nh, hs, T) -> (B, nh, T, T)
+        if self.flash:
+            # efficient attention using Flash Attention CUDA kernels
+            y = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=None, dropout_p=self.dropout if self.training else 0, is_causal=True)
+        else:
+            # manual implementation of attention
+            att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))
+            att = att.masked_fill(self.bias[:,:,:T,:T] == 0, float('-inf'))
+            att = F.softmax(att, dim=-1)
+            att = self.attn_dropout(att)
+            y = att @ v # (B, nh, T, T) x (B, nh, T, hs) -> (B, nh, T, hs)
+        y = y.transpose(1, 2).contiguous().view(B, T, C) # re-assemble all head outputs side by side
+
+        # output projection
+        y = self.resid_dropout(self.c_proj(y))
+        return y
+
+class MLP(nn.Module):
+
+    def __init__(self, config):
+        super().__init__()
+        self.c_fc    = nn.Linear(config.n_embd, 4 * config.n_embd, bias=config.bias)
+        self.gelu    = nn.GELU()
+        self.c_proj  = nn.Linear(4 * config.n_embd, config.n_embd, bias=config.bias)
+        self.dropout = nn.Dropout(config.dropout)
+
+    def forward(self, x):
+        x = self.c_fc(x)
+        x = self.gelu(x)
+        x = self.c_proj(x)
+        x = self.dropout(x)
+        return x
+
+class Block(nn.Module):
+
+    def __init__(self, config):
+        super().__init__()
+        self.ln_1 = LayerNorm(config.n_embd, bias=config.bias)
+        self.attn = CausalSelfAttention(config)
+        self.ln_2 = LayerNorm(config.n_embd, bias=config.bias)
+        self.mlp = MLP(config)
+
+    def forward(self, x):
+        x = x + self.attn(self.ln_1(x))
+        x = x + self.mlp(self.ln_2(x))
+        return x
+
+@dataclass
+class GPTConfig:
+    block_size: int = 1024
+    vocab_size: int = 50304 # GPT-2 vocab_size of 50257, padded up to nearest multiple of 64 for efficiency
+    n_layer: int = 12
+    n_head: int = 12
+    n_embd: int = 768
+    dropout: float = 0.0
+    bias: bool = True # True: bias in Linears and LayerNorms, like GPT-2. False: a bit better and faster
+
+class GPT(nn.Module):
+
+    def __init__(self, config):
+        super().__init__()
+        assert config.vocab_size is not None
+        assert config.block_size is not None
+        self.config = config
+
+        self.transformer = nn.ModuleDict(dict(
+            wte = nn.Embedding(config.vocab_size, config.n_embd),
+            wpe = nn.Embedding(config.block_size, config.n_embd),
+            drop = nn.Dropout(config.dropout),
+            h = nn.ModuleList([Block(config) for _ in range(config.n_layer)]),
+            ln_f = LayerNorm(config.n_embd, bias=config.bias),
+        ))
+        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
+        # with weight tying when using torch.compile() some warnings get generated:
+        # "UserWarning: functional_call was passed multiple values for tied weights.
+        # This behavior is deprecated and will be an error in future versions"
+        # not 100% sure what this is, so far seems to be harmless. TODO investigate
+        self.transformer.wte.weight = self.lm_head.weight # https://paperswithcode.com/method/weight-tying
+
+        # init all weights
+        self.apply(self._init_weights)
+        # apply special scaled init to the residual projections, per GPT-2 paper
+        for pn, p in self.named_parameters():
+            if pn.endswith('c_proj.weight'):
+                torch.nn.init.normal_(p, mean=0.0, std=0.02/math.sqrt(2 * config.n_layer))
+
+        # report number of parameters
+        print("number of parameters: %.2fM" % (self.get_num_params()/1e6,))
+
+    def get_num_params(self, non_embedding=True):
+        """
+        Return the number of parameters in the model.
+        For non-embedding count (default), the position embeddings get subtracted.
+        The token embeddings would too, except due to the parameter sharing these
+        params are actually used as weights in the final layer, so we include them.
+        """
+        n_params = sum(p.numel() for p in self.parameters())
+        if non_embedding:
+            n_params -= self.transformer.wpe.weight.numel()
+        return n_params
+
+    def _init_weights(self, module):
+        if isinstance(module, nn.Linear):
+            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
+            if module.bias is not None:
+                torch.nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Embedding):
+            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
+
+    def forward(self, idx, targets=None):
+        device = idx.device
+        b, t = idx.size()
+        assert t <= self.config.block_size, f"Cannot forward sequence of length {t}, block size is only {self.config.block_size}"
+        pos = torch.arange(0, t, dtype=torch.long, device=device) # shape (t)
+
+        # forward the GPT model itself
+        tok_emb = self.transformer.wte(idx) # token embeddings of shape (b, t, n_embd)
+        pos_emb = self.transformer.wpe(pos) # position embeddings of shape (t, n_embd)
+        x = self.transformer.drop(tok_emb + pos_emb)
+        for block in self.transformer.h:
+            x = block(x)
+        x = self.transformer.ln_f(x)
+
+        if targets is not None:
+            # if we are given some desired targets also calculate the loss
+            logits = self.lm_head(x)
+            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1)
+        else:
+            # inference-time mini-optimization: only forward the lm_head on the very last position
+            logits = self.lm_head(x[:, [-1], :]) # note: using list [-1] to preserve the time dim
+            loss = None
+
+        return logits, loss
+
+    def crop_block_size(self, block_size):
+        # model surgery to decrease the block size if necessary
+        # e.g. we may load the GPT2 pretrained model checkpoint (block size 1024)
+        # but want to use a smaller block size for some smaller, simpler model
+        assert block_size <= self.config.block_size
+        self.config.block_size = block_size
+        self.transformer.wpe.weight = nn.Parameter(self.transformer.wpe.weight[:block_size])
+        for block in self.transformer.h:
+            if hasattr(block.attn, 'bias'):
+                block.attn.bias = block.attn.bias[:,:,:block_size,:block_size]
+
+    @classmethod
+    def from_pretrained(cls, model_type, override_args=None):
+        assert model_type in {'gpt2', 'gpt2-medium', 'gpt2-large', 'gpt2-xl'}
+        override_args = override_args or {} # default to empty dict
+        # only dropout can be overridden see more notes below
+        assert all(k == 'dropout' for k in override_args)
+        from transformers import GPT2LMHeadModel
+        print("loading weights from pretrained gpt: %s" % model_type)
+
+        # n_layer, n_head and n_embd are determined from model_type
+        config_args = {
+            'gpt2':         dict(n_layer=12, n_head=12, n_embd=768),  # 124M params
+            'gpt2-medium':  dict(n_layer=24, n_head=16, n_embd=1024), # 350M params
+            'gpt2-large':   dict(n_layer=36, n_head=20, n_embd=1280), # 774M params
+            'gpt2-xl':      dict(n_layer=48, n_head=25, n_embd=1600), # 1558M params
+        }[model_type]
+        print("forcing vocab_size=50257, block_size=1024, bias=True")
+        config_args['vocab_size'] = 50257 # always 50257 for GPT model checkpoints
+        config_args['block_size'] = 1024 # always 1024 for GPT model checkpoints
+        config_args['bias'] = True # always True for GPT model checkpoints
+        # we can override the dropout rate, if desired
+        if 'dropout' in override_args:
+            print(f"overriding dropout rate to {override_args['dropout']}")
+            config_args['dropout'] = override_args['dropout']
+        # create a from-scratch initialized minGPT model
+        config = GPTConfig(**config_args)
+        model = GPT(config)
+        sd = model.state_dict()
+        sd_keys = sd.keys()
+        sd_keys = [k for k in sd_keys if not k.endswith('.attn.bias')] # discard this mask / buffer, not a param
+
+        # init a huggingface/transformers model
+        model_hf = GPT2LMHeadModel.from_pretrained(model_type)
+        sd_hf = model_hf.state_dict()
+
+        # copy while ensuring all of the parameters are aligned and match in names and shapes
+        sd_keys_hf = sd_hf.keys()
+        sd_keys_hf = [k for k in sd_keys_hf if not k.endswith('.attn.masked_bias')] # ignore these, just a buffer
+        sd_keys_hf = [k for k in sd_keys_hf if not k.endswith('.attn.bias')] # same, just the mask (buffer)
+        transposed = ['attn.c_attn.weight', 'attn.c_proj.weight', 'mlp.c_fc.weight', 'mlp.c_proj.weight']
+        # basically the openai checkpoints use a "Conv1D" module, but we only want to use a vanilla Linear
+        # this means that we have to transpose these weights when we import them
+        assert len(sd_keys_hf) == len(sd_keys), f"mismatched keys: {len(sd_keys_hf)} != {len(sd_keys)}"
+        for k in sd_keys_hf:
+            if any(k.endswith(w) for w in transposed):
+                # special treatment for the Conv1D weights we need to transpose
+                assert sd_hf[k].shape[::-1] == sd[k].shape
+                with torch.no_grad():
+                    sd[k].copy_(sd_hf[k].t())
+            else:
+                # vanilla copy over the other parameters
+                assert sd_hf[k].shape == sd[k].shape
+                with torch.no_grad():
+                    sd[k].copy_(sd_hf[k])
+
+        return model
+
+    def configure_optimizers(self, weight_decay, learning_rate, betas, device_type):
+        # start with all of the candidate parameters
+        param_dict = {pn: p for pn, p in self.named_parameters()}
+        # filter out those that do not require grad
+        param_dict = {pn: p for pn, p in param_dict.items() if p.requires_grad}
+        # create optim groups. Any parameters that is 2D will be weight decayed, otherwise no.
+        # i.e. all weight tensors in matmuls + embeddings decay, all biases and layernorms don't.
+        decay_params = [p for n, p in param_dict.items() if p.dim() >= 2]
+        nodecay_params = [p for n, p in param_dict.items() if p.dim() < 2]
+        optim_groups = [
+            {'params': decay_params, 'weight_decay': weight_decay},
+            {'params': nodecay_params, 'weight_decay': 0.0}
+        ]
+        num_decay_params = sum(p.numel() for p in decay_params)
+        num_nodecay_params = sum(p.numel() for p in nodecay_params)
+        print(f"num decayed parameter tensors: {len(decay_params)}, with {num_decay_params:,} parameters")
+        print(f"num non-decayed parameter tensors: {len(nodecay_params)}, with {num_nodecay_params:,} parameters")
+        # Create AdamW optimizer and use the fused version if it is available
+        fused_available = 'fused' in inspect.signature(torch.optim.AdamW).parameters
+        use_fused = fused_available and device_type == 'cuda'
+        extra_args = dict(fused=True) if use_fused else dict()
+        optimizer = torch.optim.AdamW(optim_groups, lr=learning_rate, betas=betas, **extra_args)
+        print(f"using fused AdamW: {use_fused}")
+
+        return optimizer
+
+    def estimate_mfu(self, fwdbwd_per_iter, dt):
+        """ estimate model flops utilization (MFU) in units of A100 bfloat16 peak FLOPS """
+        # first estimate the number of flops we do per iteration.
+        # see PaLM paper Appendix B as ref: https://arxiv.org/abs/2204.02311
+        N = self.get_num_params()
+        cfg = self.config
+        L, H, Q, T = cfg.n_layer, cfg.n_head, cfg.n_embd//cfg.n_head, cfg.block_size
+        flops_per_token = 6*N + 12*L*H*Q*T
+        flops_per_fwdbwd = flops_per_token * T
+        flops_per_iter = flops_per_fwdbwd * fwdbwd_per_iter
+        # express our flops throughput as ratio of A100 bfloat16 peak flops
+        flops_achieved = flops_per_iter * (1.0/dt) # per second
+        flops_promised = 312e12 # A100 GPU bfloat16 peak flops is 312 TFLOPS
+        mfu = flops_achieved / flops_promised
+        return mfu
+
+    @torch.no_grad()
+    def generate(self, idx, max_new_tokens, temperature=1.0, top_k=None):
+        """
+        Take a conditioning sequence of indices idx (LongTensor of shape (b,t)) and complete
+        the sequence max_new_tokens times, feeding the predictions back into the model each time.
+        Most likely you'll want to make sure to be in model.eval() mode of operation for this.
+        """
+        for _ in range(max_new_tokens):
+            # if the sequence context is growing too long we must crop it at block_size
+            idx_cond = idx if idx.size(1) <= self.config.block_size else idx[:, -self.config.block_size:]
+            # forward the model to get the logits for the index in the sequence
+            logits, _ = self(idx_cond)
+            # pluck the logits at the final step and scale by desired temperature
+            logits = logits[:, -1, :] / temperature
+            # optionally crop the logits to only the top k options
+            if top_k is not None:
+                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
+                logits[logits < v[:, [-1]]] = -float('Inf')
+            # apply softmax to convert logits to (normalized) probabilities
+            probs = F.softmax(logits, dim=-1)
+            # sample from the distribution
+            idx_next = torch.multinomial(probs, num_samples=1)
+            # append sampled index to the running sequence and continue
+            idx = torch.cat((idx, idx_next), dim=1)
+
+        return idx
diff --git a/src/train.py b/src/train.py
new file mode 100644
index 0000000..de57850
--- /dev/null
+++ b/src/train.py
@@ -0,0 +1,336 @@
+"""
+This training script can be run both on a single gpu in debug mode,
+and also in a larger training run with distributed data parallel (ddp).
+
+To run on a single GPU, example:
+$ python train.py --batch_size=32 --compile=False
+
+To run with DDP on 4 gpus on 1 node, example:
+$ torchrun --standalone --nproc_per_node=4 train.py
+
+To run with DDP on 4 gpus across 2 nodes, example:
+- Run on the first (master) node with example IP 123.456.123.456:
+$ torchrun --nproc_per_node=8 --nnodes=2 --node_rank=0 --master_addr=123.456.123.456 --master_port=1234 train.py
+- Run on the worker node:
+$ torchrun --nproc_per_node=8 --nnodes=2 --node_rank=1 --master_addr=123.456.123.456 --master_port=1234 train.py
+(If your cluster does not have Infiniband interconnect prepend NCCL_IB_DISABLE=1)
+"""
+
+import os
+import time
+import math
+import pickle
+from contextlib import nullcontext
+
+import numpy as np
+import torch
+from torch.nn.parallel import DistributedDataParallel as DDP
+from torch.distributed import init_process_group, destroy_process_group
+
+from model import GPTConfig, GPT
+
+# -----------------------------------------------------------------------------
+# default config values designed to train a gpt2 (124M) on OpenWebText
+# I/O
+out_dir = 'out'
+eval_interval = 2000
+log_interval = 1
+eval_iters = 200
+eval_only = False # if True, script exits right after the first eval
+always_save_checkpoint = True # if True, always save a checkpoint after each eval
+init_from = 'scratch' # 'scratch' or 'resume' or 'gpt2*'
+# wandb logging
+wandb_log = False # disabled by default
+wandb_project = 'owt'
+wandb_run_name = 'gpt2' # 'run' + str(time.time())
+# data
+dataset = 'openwebtext'
+gradient_accumulation_steps = 5 * 8 # used to simulate larger batch sizes
+batch_size = 12 # if gradient_accumulation_steps > 1, this is the micro-batch size
+block_size = 1024
+# model
+n_layer = 12
+n_head = 12
+n_embd = 768
+dropout = 0.0 # for pretraining 0 is good, for finetuning try 0.1+
+bias = False # do we use bias inside LayerNorm and Linear layers?
+# adamw optimizer
+learning_rate = 6e-4 # max learning rate
+max_iters = 600000 # total number of training iterations
+weight_decay = 1e-1
+beta1 = 0.9
+beta2 = 0.95
+grad_clip = 1.0 # clip gradients at this value, or disable if == 0.0
+# learning rate decay settings
+decay_lr = True # whether to decay the learning rate
+warmup_iters = 2000 # how many steps to warm up for
+lr_decay_iters = 600000 # should be ~= max_iters per Chinchilla
+min_lr = 6e-5 # minimum learning rate, should be ~= learning_rate/10 per Chinchilla
+# DDP settings
+backend = 'nccl' # 'nccl', 'gloo', etc.
+# system
+device = 'cuda' # examples: 'cpu', 'cuda', 'cuda:0', 'cuda:1' etc., or try 'mps' on macbooks
+dtype = 'bfloat16' if torch.cuda.is_available() and torch.cuda.is_bf16_supported() else 'float16' # 'float32', 'bfloat16', or 'float16', the latter will auto implement a GradScaler
+compile = True # use PyTorch 2.0 to compile the model to be faster
+# -----------------------------------------------------------------------------
+config_keys = [k for k,v in globals().items() if not k.startswith('_') and isinstance(v, (int, float, bool, str))]
+exec(open('configurator.py').read()) # overrides from command line or config file
+config = {k: globals()[k] for k in config_keys} # will be useful for logging
+# -----------------------------------------------------------------------------
+
+# various inits, derived attributes, I/O setup
+ddp = int(os.environ.get('RANK', -1)) != -1 # is this a ddp run?
+if ddp:
+    init_process_group(backend=backend)
+    ddp_rank = int(os.environ['RANK'])
+    ddp_local_rank = int(os.environ['LOCAL_RANK'])
+    ddp_world_size = int(os.environ['WORLD_SIZE'])
+    device = f'cuda:{ddp_local_rank}'
+    torch.cuda.set_device(device)
+    master_process = ddp_rank == 0 # this process will do logging, checkpointing etc.
+    seed_offset = ddp_rank # each process gets a different seed
+    # world_size number of processes will be training simultaneously, so we can scale
+    # down the desired gradient accumulation iterations per process proportionally
+    assert gradient_accumulation_steps % ddp_world_size == 0
+    gradient_accumulation_steps //= ddp_world_size
+else:
+    # if not ddp, we are running on a single gpu, and one process
+    master_process = True
+    seed_offset = 0
+    ddp_world_size = 1
+tokens_per_iter = gradient_accumulation_steps * ddp_world_size * batch_size * block_size
+print(f"tokens per iteration will be: {tokens_per_iter:,}")
+
+if master_process:
+    os.makedirs(out_dir, exist_ok=True)
+torch.manual_seed(1337 + seed_offset)
+torch.backends.cuda.matmul.allow_tf32 = True # allow tf32 on matmul
+torch.backends.cudnn.allow_tf32 = True # allow tf32 on cudnn
+device_type = 'cuda' if 'cuda' in device else 'cpu' # for later use in torch.autocast
+# note: float16 data type will automatically use a GradScaler
+ptdtype = {'float32': torch.float32, 'bfloat16': torch.bfloat16, 'float16': torch.float16}[dtype]
+ctx = nullcontext() if device_type == 'cpu' else torch.amp.autocast(device_type=device_type, dtype=ptdtype)
+
+# poor man's data loader
+data_dir = os.path.join('data', dataset)
+def get_batch(split):
+    # We recreate np.memmap every batch to avoid a memory leak, as per
+    # https://stackoverflow.com/questions/45132940/numpy-memmap-memory-usage-want-to-iterate-once/61472122#61472122
+    if split == 'train':
+        data = np.memmap(os.path.join(data_dir, 'train.bin'), dtype=np.uint16, mode='r')
+    else:
+        data = np.memmap(os.path.join(data_dir, 'val.bin'), dtype=np.uint16, mode='r')
+    ix = torch.randint(len(data) - block_size, (batch_size,))
+    x = torch.stack([torch.from_numpy((data[i:i+block_size]).astype(np.int64)) for i in ix])
+    y = torch.stack([torch.from_numpy((data[i+1:i+1+block_size]).astype(np.int64)) for i in ix])
+    if device_type == 'cuda':
+        # pin arrays x,y, which allows us to move them to GPU asynchronously (non_blocking=True)
+        x, y = x.pin_memory().to(device, non_blocking=True), y.pin_memory().to(device, non_blocking=True)
+    else:
+        x, y = x.to(device), y.to(device)
+    return x, y
+
+# init these up here, can override if init_from='resume' (i.e. from a checkpoint)
+iter_num = 0
+best_val_loss = 1e9
+
+# attempt to derive vocab_size from the dataset
+meta_path = os.path.join(data_dir, 'meta.pkl')
+meta_vocab_size = None
+if os.path.exists(meta_path):
+    with open(meta_path, 'rb') as f:
+        meta = pickle.load(f)
+    meta_vocab_size = meta['vocab_size']
+    print(f"found vocab_size = {meta_vocab_size} (inside {meta_path})")
+
+# model init
+model_args = dict(n_layer=n_layer, n_head=n_head, n_embd=n_embd, block_size=block_size,
+                  bias=bias, vocab_size=None, dropout=dropout) # start with model_args from command line
+if init_from == 'scratch':
+    # init a new model from scratch
+    print("Initializing a new model from scratch")
+    # determine the vocab size we'll use for from-scratch training
+    if meta_vocab_size is None:
+        print("defaulting to vocab_size of GPT-2 to 50304 (50257 rounded up for efficiency)")
+    model_args['vocab_size'] = meta_vocab_size if meta_vocab_size is not None else 50304
+    gptconf = GPTConfig(**model_args)
+    model = GPT(gptconf)
+elif init_from == 'resume':
+    print(f"Resuming training from {out_dir}")
+    # resume training from a checkpoint.
+    ckpt_path = os.path.join(out_dir, 'ckpt.pt')
+    checkpoint = torch.load(ckpt_path, map_location=device)
+    checkpoint_model_args = checkpoint['model_args']
+    # force these config attributes to be equal otherwise we can't even resume training
+    # the rest of the attributes (e.g. dropout) can stay as desired from command line
+    for k in ['n_layer', 'n_head', 'n_embd', 'block_size', 'bias', 'vocab_size']:
+        model_args[k] = checkpoint_model_args[k]
+    # create the model
+    gptconf = GPTConfig(**model_args)
+    model = GPT(gptconf)
+    state_dict = checkpoint['model']
+    # fix the keys of the state dictionary :(
+    # honestly no idea how checkpoints sometimes get this prefix, have to debug more
+    unwanted_prefix = '_orig_mod.'
+    for k,v in list(state_dict.items()):
+        if k.startswith(unwanted_prefix):
+            state_dict[k[len(unwanted_prefix):]] = state_dict.pop(k)
+    model.load_state_dict(state_dict)
+    iter_num = checkpoint['iter_num']
+    best_val_loss = checkpoint['best_val_loss']
+elif init_from.startswith('gpt2'):
+    print(f"Initializing from OpenAI GPT-2 weights: {init_from}")
+    # initialize from OpenAI GPT-2 weights
+    override_args = dict(dropout=dropout)
+    model = GPT.from_pretrained(init_from, override_args)
+    # read off the created config params, so we can store them into checkpoint correctly
+    for k in ['n_layer', 'n_head', 'n_embd', 'block_size', 'bias', 'vocab_size']:
+        model_args[k] = getattr(model.config, k)
+# crop down the model block size if desired, using model surgery
+if block_size < model.config.block_size:
+    model.crop_block_size(block_size)
+    model_args['block_size'] = block_size # so that the checkpoint will have the right value
+model.to(device)
+
+# initialize a GradScaler. If enabled=False scaler is a no-op
+scaler = torch.cuda.amp.GradScaler(enabled=(dtype == 'float16'))
+
+# optimizer
+optimizer = model.configure_optimizers(weight_decay, learning_rate, (beta1, beta2), device_type)
+if init_from == 'resume':
+    optimizer.load_state_dict(checkpoint['optimizer'])
+checkpoint = None # free up memory
+
+# compile the model
+if compile:
+    print("compiling the model... (takes a ~minute)")
+    unoptimized_model = model
+    model = torch.compile(model) # requires PyTorch 2.0
+
+# wrap model into DDP container
+if ddp:
+    model = DDP(model, device_ids=[ddp_local_rank])
+
+# helps estimate an arbitrarily accurate loss over either split using many batches
+@torch.no_grad()
+def estimate_loss():
+    out = {}
+    model.eval()
+    for split in ['train', 'val']:
+        losses = torch.zeros(eval_iters)
+        for k in range(eval_iters):
+            X, Y = get_batch(split)
+            with ctx:
+                logits, loss = model(X, Y)
+            losses[k] = loss.item()
+        out[split] = losses.mean()
+    model.train()
+    return out
+
+# learning rate decay scheduler (cosine with warmup)
+def get_lr(it):
+    # 1) linear warmup for warmup_iters steps
+    if it < warmup_iters:
+        return learning_rate * (it + 1) / (warmup_iters + 1)
+    # 2) if it > lr_decay_iters, return min learning rate
+    if it > lr_decay_iters:
+        return min_lr
+    # 3) in between, use cosine decay down to min learning rate
+    decay_ratio = (it - warmup_iters) / (lr_decay_iters - warmup_iters)
+    assert 0 <= decay_ratio <= 1
+    coeff = 0.5 * (1.0 + math.cos(math.pi * decay_ratio)) # coeff ranges 0..1
+    return min_lr + coeff * (learning_rate - min_lr)
+
+# logging
+if wandb_log and master_process:
+    import wandb
+    wandb.init(project=wandb_project, name=wandb_run_name, config=config)
+
+# training loop
+X, Y = get_batch('train') # fetch the very first batch
+t0 = time.time()
+local_iter_num = 0 # number of iterations in the lifetime of this process
+raw_model = model.module if ddp else model # unwrap DDP container if needed
+running_mfu = -1.0
+while True:
+
+    # determine and set the learning rate for this iteration
+    lr = get_lr(iter_num) if decay_lr else learning_rate
+    for param_group in optimizer.param_groups:
+        param_group['lr'] = lr
+
+    # evaluate the loss on train/val sets and write checkpoints
+    if iter_num % eval_interval == 0 and master_process:
+        losses = estimate_loss()
+        print(f"step {iter_num}: train loss {losses['train']:.4f}, val loss {losses['val']:.4f}")
+        if wandb_log:
+            wandb.log({
+                "iter": iter_num,
+                "train/loss": losses['train'],
+                "val/loss": losses['val'],
+                "lr": lr,
+                "mfu": running_mfu*100, # convert to percentage
+            })
+        if losses['val'] < best_val_loss or always_save_checkpoint:
+            best_val_loss = losses['val']
+            if iter_num > 0:
+                checkpoint = {
+                    'model': raw_model.state_dict(),
+                    'optimizer': optimizer.state_dict(),
+                    'model_args': model_args,
+                    'iter_num': iter_num,
+                    'best_val_loss': best_val_loss,
+                    'config': config,
+                }
+                print(f"saving checkpoint to {out_dir}")
+                torch.save(checkpoint, os.path.join(out_dir, 'ckpt.pt'))
+    if iter_num == 0 and eval_only:
+        break
+
+    # forward backward update, with optional gradient accumulation to simulate larger batch size
+    # and using the GradScaler if data type is float16
+    for micro_step in range(gradient_accumulation_steps):
+        if ddp:
+            # in DDP training we only need to sync gradients at the last micro step.
+            # the official way to do this is with model.no_sync() context manager, but
+            # I really dislike that this bloats the code and forces us to repeat code
+            # looking at the source of that context manager, it just toggles this variable
+            model.require_backward_grad_sync = (micro_step == gradient_accumulation_steps - 1)
+        with ctx:
+            logits, loss = model(X, Y)
+            loss = loss / gradient_accumulation_steps # scale the loss to account for gradient accumulation
+        # immediately async prefetch next batch while model is doing the forward pass on the GPU
+        X, Y = get_batch('train')
+        # backward pass, with gradient scaling if training in fp16
+        scaler.scale(loss).backward()
+    # clip the gradient
+    if grad_clip != 0.0:
+        scaler.unscale_(optimizer)
+        torch.nn.utils.clip_grad_norm_(model.parameters(), grad_clip)
+    # step the optimizer and scaler if training in fp16
+    scaler.step(optimizer)
+    scaler.update()
+    # flush the gradients as soon as we can, no need for this memory anymore
+    optimizer.zero_grad(set_to_none=True)
+
+    # timing and logging
+    t1 = time.time()
+    dt = t1 - t0
+    t0 = t1
+    if iter_num % log_interval == 0 and master_process:
+        # get loss as float. note: this is a CPU-GPU sync point
+        # scale up to undo the division above, approximating the true total loss (exact would have been a sum)
+        lossf = loss.item() * gradient_accumulation_steps
+        if local_iter_num >= 5: # let the training loop settle a bit
+            mfu = raw_model.estimate_mfu(batch_size * gradient_accumulation_steps, dt)
+            running_mfu = mfu if running_mfu == -1.0 else 0.9*running_mfu + 0.1*mfu
+        print(f"iter {iter_num}: loss {lossf:.4f}, time {dt*1000:.2f}ms, mfu {running_mfu*100:.2f}%")
+    iter_num += 1
+    local_iter_num += 1
+
+    # termination conditions
+    if iter_num > max_iters:
+        break
+
+if ddp:
+    destroy_process_group()
diff --git a/src/train_gpt2.py b/src/train_gpt2.py
new file mode 100644
index 0000000..8f19273
--- /dev/null
+++ b/src/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/train_shakespeare_char.py b/src/train_shakespeare_char.py
new file mode 100644
index 0000000..41c81df
--- /dev/null
+++ b/src/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
diff --git a/sync_and_run.sh b/sync_and_run.sh
new file mode 100755
index 0000000..b716f32
--- /dev/null
+++ b/sync_and_run.sh
@@ -0,0 +1,33 @@
+#!/bin/bash
+# Sync code from voidlaptop to meru's jspace container and run commands
+# Usage: ./sync_and_run.sh [command or "shell"]
+
+set -e
+
+MACHINE="meru"
+CONTAINER="jspace-nanogpt"
+REMOTE_CODE="/mnt/appdata/jspace/code"
+LOCAL_PROJECT="$HOME/projects/jspace-nanogpt"
+
+echo "=== Syncing code to $MACHINE ==="
+rsync -avz --delete \
+    --exclude='upstream-nanogpt' \
+    --exclude='.git' \
+    --exclude='__pycache__' \
+    --exclude='*.pyc' \
+    --exclude='outputs/*' \
+    --exclude='data/*' \
+    "${LOCAL_PROJECT}/" \
+    "${MACHINE}:${REMOTE_CODE}/"
+
+if [ "$1" = "shell" ]; then
+    echo "=== Opening shell in container ==="
+    ssh -t "$MACHINE" "docker exec -it $CONTAINER bash"
+elif [ -n "$1" ]; then
+    echo "=== Running: $@ ==="
+    ssh "$MACHINE" "docker exec $CONTAINER bash -c 'cd /workspace/code && $*'"
+else
+    echo "=== Synced. No command specified. ==="
+    echo "Usage: $0 shell          # interactive bash"
+    echo "       $0 'command ...'  # run a command"
+fi