Add QK-norm toggle, TinyStories prep and automatic run logging

nanoGPT gets a qk_norm flag (off = upstream behaviour), a seed option and
a resolved-config dump. run.py logs each run's git commit, config, data
checksums, environment and metrics, and stops diverged runs.

Claude-Session: https://claude.ai/code/session_01S8KN7NuUCm8nXLmy6eDuVv

06-experiments/nanogpt/README-thesis.md

@@ -0,0 +1,20 @@
+# nanoGPT (vendored)
+
+`model.py`, `train.py`, `configurator.py` and `LICENSE` were copied unmodified from
+[karpathy/nanoGPT](https://github.com/karpathy/nanoGPT) at commit
+`3adf61e154c3fe3fca428ad6bc3818b27a3b8291` (MIT licence), in their own commit.
+
+Every change since then is marked `[thesis]` in the code; `git log -p -- 06-experiments/nanogpt`
+shows them all. Summary:
+
+| File | Change | Why |
+|---|---|---|
+| `model.py` | `qk_norm` config flag; when on, LayerNorm (no bias by default) over each head's queries and keys before the dot product. `1/sqrt(d)` scaling kept | the independent variable. Off = exactly upstream |
+| `train.py` | `qk_norm`, `seed`, `data_root` config options | QK-norm toggle; seeds for repeat runs (upstream hard-codes 1337); data path |
+| `train.py` | `configurator.py` loaded relative to the script | so runs can launch from any directory |
+| `train.py` | writes resolved `config.json` into `out_dir` | run logging |
+| `train.py` | resume reads `qk_norm` from checkpoint | so resumed runs keep the same architecture |
+
+**Design choice to confirm:** QK-norm here is LayerNorm-style, as in Dehghani et al. (2023) and
+Wortsman et al. (2023). Henry et al. (2020) used L2-normalisation with a learnable scale instead,
+and OLMo 2 uses RMSNorm. `[Your call: which variant the thesis studies.]`

06-experiments/nanogpt/model.py

@@ -41,6 +41,14 @@ class CausalSelfAttention(nn.Module):
         self.n_head = config.n_head
         self.n_embd = config.n_embd
         self.dropout = config.dropout
+        # [thesis] QK-norm: LayerNorm over each head's query and key vectors before the
+        # dot product (the "qk-layernorm" of Dehghani et al. 2023 / Wortsman et al. 2023).
+        # Off by default, so qk_norm=False is exactly upstream nanoGPT.
+        self.qk_norm = config.qk_norm
+        if self.qk_norm:
+            head_dim = config.n_embd // config.n_head
+            self.q_norm = LayerNorm(head_dim, bias=config.bias)
+            self.k_norm = LayerNorm(head_dim, bias=config.bias)
         # 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:
@@ -57,6 +65,9 @@ class CausalSelfAttention(nn.Module):
         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)
+        if self.qk_norm:  # [thesis] normalise over the head dimension
+            q = self.q_norm(q)
+            k = self.k_norm(k)
 
         # causal self-attention; Self-attend: (B, nh, T, hs) x (B, nh, hs, T) -> (B, nh, T, T)
         if self.flash:
@@ -114,6 +125,7 @@ class GPTConfig:
     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
+    qk_norm: bool = False # [thesis] LayerNorm on per-head queries and keys
 
 class GPT(nn.Module):
 

06-experiments/nanogpt/train.py

@@ -54,6 +54,9 @@ 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?
+qk_norm = False # [thesis] QK-norm on/off: the independent variable
+seed = 1337 # [thesis] base RNG seed (upstream hard-codes 1337)
+data_root = 'data' # [thesis] where dataset folders live
 # adamw optimizer
 learning_rate = 6e-4 # max learning rate
 max_iters = 600000 # total number of training iterations
@@ -74,7 +77,7 @@ dtype = 'bfloat16' if torch.cuda.is_available() and torch.cuda.is_bf16_supported
 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
+exec(open(os.path.join(os.path.dirname(os.path.abspath(__file__)), 'configurator.py')).read()) # overrides from command line or config file
 config = {k: globals()[k] for k in config_keys} # will be useful for logging
 # -----------------------------------------------------------------------------
 
@@ -103,7 +106,9 @@ 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)
+    import json  # [thesis] record the fully resolved config for the run log
+    with open(os.path.join(out_dir, 'config.json'), 'w') as f: json.dump(config, f, indent=2)
+torch.manual_seed(seed + 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
@@ -112,7 +117,7 @@ ptdtype = {'float32': torch.float32, 'bfloat16': torch.bfloat16, 'float16': torc
 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)
+data_dir = os.path.join(data_root, 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
@@ -145,7 +150,7 @@ if os.path.exists(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
+                  bias=bias, vocab_size=None, dropout=dropout, qk_norm=qk_norm) # start with model_args from command line
 if init_from == 'scratch':
     # init a new model from scratch
     print("Initializing a new model from scratch")
@@ -163,8 +168,8 @@ elif init_from == 'resume':
     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]
+    for k in ['n_layer', 'n_head', 'n_embd', 'block_size', 'bias', 'vocab_size', 'qk_norm']:
+        model_args[k] = checkpoint_model_args.get(k, False)
     # create the model
     gptconf = GPTConfig(**model_args)
     model = GPT(gptconf)