# Journey rungs → code Each labeled rung on the progress plot mapped to the **verbatim** code that implements it (from `opt_param_full2.py`; the swarm `train.py` is equivalent). GP snippet from `bo_champ.py`. --- ## ① Karpathy seed — `0.9961` The vanilla baseline shape: a plain decoder LM. Everything below is added on top of this. ```python # opt_param_full2.py : lines 100–108 @dataclass class GPTConfig: sequence_len: int = 2048 vocab_size: int = 32768 n_layer: int = 12 n_head: int = 6 n_kv_head: int = 6 n_embd: int = 768 window_pattern: str = "SSSL" ``` --- ## ② + Muon optimizer — `0.9913 (biggest single rung, −0.076)` Muon = Nesterov momentum, then **Polar-Express orthogonalization** (a Newton–Schulz iteration) of the gradient. Makes every token's update far more effective. ```python # opt_param_full2.py : lines 685–691 polar_express_coeffs = [ (8.156554524902461, -22.48329292557795, 15.878769915207462), (4.042929935166739, -2.808917465908714, 0.5000178451051316), (3.8916678022926607, -2.772484153217685, 0.5060648178503393), (3.285753657755655, -2.3681294933425376, 0.46449024233003106), (2.3465413258596377, -1.7097828382687081, 0.42323551169305323), ] ``` ```python # opt_param_full2.py : lines 719–743 def muon_step_fused( stacked_grads, stacked_params, momentum_buffer, second_momentum_buffer, momentum_t, lr_t, wd_t, beta2_t, ns_steps, red_dim, ): # Nesterov momentum momentum = momentum_t.to(stacked_grads.dtype) momentum_buffer.lerp_(stacked_grads, 1 - momentum) g = stacked_grads.lerp_(momentum_buffer, momentum) # Polar express orthogonalization X = g.bfloat16() X = X / (X.norm(dim=(-2, -1), keepdim=True) * 1.02 + 1e-6) if g.size(-2) > g.size(-1): for a, b, c in _orth_coeffs[:ns_steps]: A = X.mT @ X B = b * A + c * (A @ A) X = a * X + X @ B else: ``` --- ## ③ + value embeddings — `0.9834` A cheap **gated** per-token embedding injected straight into the attention *value* stream. ```python # opt_param_full2.py : lines 178–182 if ve is not None: ve = ve.view(B, T, self.n_kv_head, self.head_dim) gate = 2 * torch.sigmoid(self.ve_gate(x[..., : self.ve_gate_channels])) v = v + gate.unsqueeze(-1) * ve ``` --- ## ④ ★ n-gram memory — `0.9704 → 0.9559 (the "memory ≠ compute" pivot)` Local bigram/trigram statistics **hashed into huge lookup tables** (per-layer decorrelated primes) and added to the value stream. O(1) lookup, ~0 FLOPs — memorize instead of learning by gradient. ```python # opt_param_full2.py : lines 282–303 self.bigram_table_size = config.vocab_size * int(os.environ.get("NGRAM_MULT", 64)) # CROSSOVER B: 64x bigram tables self.bigram_K = 2 half_kv_dim = kv_dim // 2 # PER-LAYER DECORRELATED: completely disjoint hash prime pairs per bigram VE layer # Each layer uses entirely distinct multipliers -- zero prime reuse within bigram type # Constants from Murmur/FNV/golden-ratio family for good avalanche behavior _decorr_bigram_primes = [ [(2654435761, 2246822519), (1013904223, 6291469)], # layer 1: golden-ratio family [(374761393, 668265263), (3266489917, 104729)], # layer 3: prime family [(1640531527, 97531), (48271, 40503)], # layer 5: LCG/Knuth family [(16777619, 2166136261), (3432918353, 461845907)], # layer 7: MurmurHash3 family ] self.bigram_hash_primes_per_layer = {} self.bigram_ves = nn.ModuleDict() for j, layer_i in enumerate(ve_layers): self.bigram_ves[str(layer_i)] = nn.ModuleList([ nn.Embedding(self.bigram_table_size, half_kv_dim), nn.Embedding(self.bigram_table_size, half_kv_dim), ]) self.bigram_hash_primes_per_layer[layer_i] = _decorr_bigram_primes[j % len(_decorr_bigram_primes)] # Multi-layer factored trigram VE: K=2 half-dim tables at layers 1+5 plus layer 7. self.trigram_ve_layers = ( ``` ```python # opt_param_full2.py : lines 601–606 for layer_i in self.bigram_ve_layers: layer_bg_primes = self.bigram_hash_primes_per_layer[layer_i] bigram_indices_per_layer[layer_i] = [ ((prev_idx * p1) ^ (idx * p2)) % self.bigram_table_size for p1, p2 in layer_bg_primes ] ``` ```python # opt_param_full2.py : lines 627–632 # Factored multi-hash bigram VE: concat K=2 half-dim lookups from independent hashes (per-layer primes) if (not os.environ.get("NGRAM_OFF")) and i in self.bigram_ve_layers: layer_ves = self.bigram_ves[str(i)] layer_indices = bigram_indices_per_layer[i] bgve = torch.cat([layer_ves[k](layer_indices[k]) for k in range(self.bigram_K)], dim=-1) else: ``` --- ## ⑤ + sliding window — `0.9777` Per-layer attention windows from a pattern string (`SSSL`/`TTTL`): tiny/short/long interleaved, last layer full. Cheaper attention → more tokens/sec. ```python # opt_param_full2.py : lines 396–409 def _compute_window_sizes(self, config): pattern = config.window_pattern.upper() assert all(c in "SLT" for c in pattern) long_window = config.sequence_len short_window = long_window // int(os.environ.get("SHORT_DIV", 2)) tiny_window = long_window // int(os.environ.get("TINY_DIV", 4)) char_to_window = {"L": (long_window, 0), "S": (short_window, 0), "T": (tiny_window, 0)} window_sizes = [] for layer_idx in range(config.n_layer): char = pattern[layer_idx % len(pattern)] window_sizes.append(char_to_window[char]) window_sizes[-1] = (long_window, 0) if os.environ.get("FULL_WINDOW"): window_sizes = [(long_window, 0) for _ in window_sizes] ``` --- ## ⑥ + logit softcap — `0.9759` Tanh soft-cap on output logits (Gemma-2 style) — stabilizes the softmax so a higher LR is safe. ```python # opt_param_full2.py : lines 652–656 # Decoupled softcap in BF16: skip float() cast, halve logit tensor memory # Since model is natively BF16, softcap in BF16 should be numerically adequate logits = self.lm_head(x) if not os.environ.get("SOFTCAP_OFF"): logits = 16.5 * torch.tanh(logits / 15.0) ``` --- ## ⑦ modern arch: QK-norm + ReLU² — `—` QK-normalization **before** rotary, and a **ReLU² MLP** activation with a learned threshold τ. ```python # opt_param_full2.py : lines 196–199 # QK-norm refinement: normalize BEFORE rotary instead of after if not os.environ.get("QKNORM_OFF"): q, k = norm(q), norm(k) q, k = apply_rotary_emb(q, cos, sin), apply_rotary_emb(k, cos, sin) ``` ```python # opt_param_full2.py : lines 224–226 def forward(self, x): h = self.c_fc(x) h = F.relu(h - self.tau) if os.environ.get("RELU2_OFF") else F.relu(h - self.tau).square() ``` --- ## ⑧ compute-optimal width: dim 768 → 640 — `~0.9376 (BO-found)` The counterintuitive "go narrow": fewer params → more tokens/sec under the fixed budget. Just the `MODEL_DIM` knob, discovered by BO. ```python # opt_param_full2.py : lines 947–949 base_dim = depth * ASPECT_RATIO model_dim = int(os.environ["MODEL_DIM"]) if os.environ.get("MODEL_DIM") else ((base_dim + HEAD_DIM - 1) // HEAD_DIM) * HEAD_DIM num_heads = model_dim // HEAD_DIM ``` --- ## ⑨ joint Bayesian-opt → 0.93421 — `0.93421 (the squeeze phase)` A from-scratch numpy **GP + Expected-Improvement** over the continuous hyperparameters: RBF kernel, Cholesky solve, EI acquisition (`bo_champ.py`). ```python # bo_champ.py : lines 130–150 ymu, ysd = y.mean(), y.std() + 1e-9 yz = (y - ymu) / ysd l, noise = 0.22, 1e-3 def rbf(A, B): return np.exp(-0.5 * ((A[:, None, :] - B[None, :, :]) ** 2).sum(-1) / (l * l)) Kn = rbf(X, X) + noise * np.eye(len(X)) L = np.linalg.cholesky(Kn + 1e-6 * np.eye(len(X))) alpha = np.linalg.solve(L.T, np.linalg.solve(L, yz)) Ks = rbf(C, X) mu = Ks @ alpha V = np.linalg.solve(L, Ks.T) var = np.clip(1.0 - (V * V).sum(0), 1e-9, None) sd = np.sqrt(var) ybest = yz.min() xi = 0.001 imp = ybest - mu - xi z = imp / sd ei = np.where(sd > 0, imp * _Phi(z) + sd * _phi(z), 0.0) picks, eiw = [], ei.copy() for _ in range(k): bi = int(np.argmax(eiw)) ```