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/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):