code: model/train variants + minimal champion

code/model_canon_HEAD.py

@@ -0,0 +1,701 @@
+"""
+Full definition of a GPT Language Model, all of it in this single file.
+Includes:
+- cheap Q gate
+- Q-MLP gate (128 hidden)
+- Q-MLP gate (192 hidden) -- param-matched control for Q+A MLP
+- full-width Q gate
+- Q head-shared elementwise gate
+- full-width X gate
+- bottleneck X gate
+- paper-faithful X variants:
+  - head-specific elementwise G1
+  - head-specific headwise G1
+  - head-shared elementwise G1
+- Q+A dual-signal variants:
+  - cheap Q+A gate (shared Linear 128->64 applied per head/token)
+  - Q+A MLP gate (shared MLP 128->128->64 applied per head/token)
+  - legacy Q+A head-shared elementwise gate (shared Linear 128->64; kept for compatibility)
+  - bilinear diagonal Q+A gate (learned elementwise q*y interaction)
+  - normed Q+A gate (RMSNorm on q and y before shared 128->64 mix)
+  - lowrank Q+A gate (shared 128->16->64 bottleneck mix)
+  - soft-QA gate (independent learned soft blends for q and y)
+- per-dimension soft-QA gate (independent 64-dim learned blend scales for q and y)
+"""
+
+import math
+import inspect
+from dataclasses import dataclass
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-6):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+    def forward(self, x):
+        rms = torch.sqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
+        return x / rms * self.weight
+
+class LayerNorm(nn.Module):
+    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
+        self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd, bias=config.bias)
+        self.c_proj = nn.Linear(config.n_embd, config.n_embd, bias=config.bias)
+
+        self.qv_variant = config.qv_variant
+        hs = config.n_embd // config.n_head
+
+        # --------------------------------------------------
+        # Q-conditioned gates
+        # --------------------------------------------------
+        if self.qv_variant in ['dynamic', 'dynamic_swiglu']:
+            self.qv_gate_proj = nn.Linear(hs, hs, bias=False)
+
+        if self.qv_variant == 'dynamic_qconditioned_mlp128':
+            self.q_gate_fc1 = nn.Linear(hs, 128, bias=False)
+            self.q_gate_fc2 = nn.Linear(128, hs, bias=False)
+
+        if self.qv_variant == 'dynamic_qconditioned_mlp192':
+            self.q_gate_fc1_192 = nn.Linear(hs, 192, bias=False)
+            self.q_gate_fc2_192 = nn.Linear(192, hs, bias=False)
+
+        if self.qv_variant == 'dynamic_qconditioned_fullwidth':
+            self.q_gate_proj_full = nn.Linear(config.n_embd, config.n_embd, bias=False)
+
+        if self.qv_variant == 'dynamic_qconditioned_fullwidth_headspecific':
+            # full-width Q source (768) with separate 768->64 map for each head
+            self.q_gate_proj_full_headspecific = nn.Parameter(torch.empty(config.n_head, config.n_embd, hs))
+            nn.init.normal_(self.q_gate_proj_full_headspecific, mean=0.0, std=0.02)
+
+        if self.qv_variant == 'dynamic_q_headshared_elementwise':
+            self.q_gate_headshared = nn.Linear(hs, hs, bias=False)
+
+        # --------------------------------------------------
+        # X-conditioned gates
+        # --------------------------------------------------
+        if self.qv_variant == 'dynamic_xconditioned_g1':
+            self.x_gate_proj = nn.Linear(config.n_embd, config.n_embd, bias=False)
+
+        if self.qv_variant == 'dynamic_xconditioned_fullwidth_headspecific':
+            # full-width X source (768) with separate 768->64 map for each head
+            self.x_gate_proj_full_headspecific = nn.Parameter(torch.empty(config.n_head, config.n_embd, hs))
+            nn.init.normal_(self.x_gate_proj_full_headspecific, mean=0.0, std=0.02)
+
+        if self.qv_variant == 'dynamic_xconditioned_bottleneck':
+            self.x_gate_bottleneck = nn.Linear(config.n_embd, hs, bias=False)
+
+        if self.qv_variant == 'dynamic_x_g1_headspecific_elementwise':
+            self.x_gate_headspecific_elementwise = nn.Parameter(torch.empty(config.n_head, hs, hs))
+            nn.init.normal_(self.x_gate_headspecific_elementwise, mean=0.0, std=0.02)
+
+        if self.qv_variant == 'dynamic_x_g1_headspecific_headwise':
+            self.x_gate_headspecific_headwise = nn.Parameter(torch.empty(config.n_head, hs, 1))
+            nn.init.normal_(self.x_gate_headspecific_headwise, mean=0.0, std=0.02)
+
+        if self.qv_variant == 'dynamic_x_g1_headshared_elementwise':
+            self.x_gate_headshared_elementwise = nn.Linear(hs, hs, bias=False)
+
+        # --------------------------------------------------
+        # Q+A dual-signal gates
+        # --------------------------------------------------
+        # random gates have no learned parameters — nothing to init here
+        # dot product gates also have no learned parameters — nothing to init here
+
+        if self.qv_variant == 'dynamic_a_conditioned':
+            # A-only: per-head Linear(64->64), conditioned on y only
+            self.a_gate_proj = nn.Linear(hs, hs, bias=False)
+
+        if self.qv_variant == 'dynamic_qa_conditioned':
+            # cheap Q+A: shared Linear(128->64) applied to each head/token
+            self.qa_gate_proj = nn.Linear(hs * 2, hs, bias=False)
+
+        if self.qv_variant == 'dynamic_qa_conditioned_headspecific':
+            # true head-specific Q+A: separate 128->64 matrix for each head
+            self.qa_gate_proj_headspecific = nn.Parameter(torch.empty(config.n_head, hs * 2, hs))
+            nn.init.normal_(self.qa_gate_proj_headspecific, mean=0.0, std=0.02)
+
+        if self.qv_variant == 'dynamic_qa_conditioned_mlp128':
+            # main Q+A MLP: per-head MLP(128->128->64)
+            self.qa_gate_fc1 = nn.Linear(hs * 2, 128, bias=False)
+            self.qa_gate_fc2 = nn.Linear(128, hs, bias=False)
+
+        if self.qv_variant == 'dynamic_qa_headshared_elementwise':
+            # legacy compatibility variant: shared Linear(128->64) applied per head/token
+            self.qa_gate_headshared = nn.Linear(hs * 2, hs, bias=False)
+
+        if self.qv_variant == 'dynamic_qa_bilinear_diag':
+            # learned elementwise interaction on q*y; 64 params per layer
+            self.qa_bilinear_diag = nn.Parameter(torch.ones(hs))
+
+        if self.qv_variant == 'dynamic_qa_conditioned_normed':
+            # RMSNorm q and y separately before shared Linear(128->64)
+            self.qa_q_norm = RMSNorm(hs)
+            self.qa_y_norm = RMSNorm(hs)
+            self.qa_gate_proj_normed = nn.Linear(hs * 2, hs, bias=False)
+
+        if self.qv_variant == 'dynamic_qa_conditioned_normed_yonly':
+            # RMSNorm y only, raw q, before shared Linear(128->64)
+            self.qa_y_norm_yonly = RMSNorm(hs)
+            self.qa_gate_proj_normed_yonly = nn.Linear(hs * 2, hs, bias=False)
+
+        if self.qv_variant == 'dynamic_qa_conditioned_normed_qonly':
+            # RMSNorm q only, raw y, before shared Linear(128->64)
+            self.qa_q_norm_qonly = RMSNorm(hs)
+            self.qa_gate_proj_normed_qonly = nn.Linear(hs * 2, hs, bias=False)
+
+        if self.qv_variant == 'dynamic_qa_conditioned_softq':
+            # Learnable blend between RMSNorm(q) and raw q, with raw y
+            self.qa_q_norm_soft = RMSNorm(hs)
+            self.qa_gate_proj_softq = nn.Linear(hs * 2, hs, bias=False)
+            self.qa_q_blend_alpha = nn.Parameter(torch.tensor(0.5))
+
+        if self.qv_variant == 'dynamic_qa_conditioned_softqa':
+            # Independent learnable blends for RMSNorm(q)/raw q and RMSNorm(y)/raw y
+            self.qa_q_norm_softqa = RMSNorm(hs)
+            self.qa_y_norm_softqa = RMSNorm(hs)
+            self.qa_gate_proj_softqa = nn.Linear(hs * 2, hs, bias=False)
+            self.qa_q_blend_alpha = nn.Parameter(torch.tensor(0.5))
+            self.qa_y_blend_alpha = nn.Parameter(torch.tensor(0.5))
+
+        if self.qv_variant == 'dynamic_qa_conditioned_softqa_perlayer':
+            # Per-layer independent learnable blends for RMSNorm(q)/raw q and RMSNorm(y)/raw y
+            self.qa_q_norm_softqa_pl = RMSNorm(hs)
+            self.qa_y_norm_softqa_pl = RMSNorm(hs)
+            self.qa_gate_proj_softqa_pl = nn.Linear(hs * 2, hs, bias=False)
+            self.qa_q_blend_alpha = nn.Parameter(torch.tensor(0.0))
+            self.qa_y_blend_alpha = nn.Parameter(torch.tensor(0.0))
+
+        if self.qv_variant == 'dynamic_qa_conditioned_softqa_perdim':
+            # Per-dimension independent learnable blends for RMSNorm(q)/raw q and RMSNorm(y)/raw y
+            self.qa_q_norm_softqa_pd = RMSNorm(hs)
+            self.qa_y_norm_softqa_pd = RMSNorm(hs)
+            self.qa_gate_proj_softqa_pd = nn.Linear(hs * 2, hs, bias=False)
+            self.qa_q_dim_scale = nn.Parameter(torch.zeros(hs))
+            self.qa_y_dim_scale = nn.Parameter(torch.zeros(hs))
+
+        if self.qv_variant == 'dynamic_qa_conditioned_softqa_perdim_informed':
+            # Per-dimension independent learnable blends with informed initialization:
+            # 1.0459 = sigmoid^-1(0.74), matching the Phase 17 converged scalar value.
+            self.qa_q_norm_softqa_pd = RMSNorm(hs)
+            self.qa_y_norm_softqa_pd = RMSNorm(hs)
+            self.qa_gate_proj_softqa_pd = nn.Linear(hs * 2, hs, bias=False)
+            self.qa_q_dim_scale = nn.Parameter(torch.full((hs,), 1.0459))
+            self.qa_y_dim_scale = nn.Parameter(torch.full((hs,), 1.0459))
+
+        if self.qv_variant == 'dynamic_qa_conditioned_softqa_perdim_free':
+            # Per-dimension independent learnable blends with tiny random symmetry breaking:
+            # start near sigmoid(scale)=0.5, but each dimension differs slightly.
+            self.qa_q_norm_softqa_pd = RMSNorm(hs)
+            self.qa_y_norm_softqa_pd = RMSNorm(hs)
+            self.qa_gate_proj_softqa_pd = nn.Linear(hs * 2, hs, bias=False)
+            self.qa_q_dim_scale = nn.Parameter(torch.empty(hs))
+            self.qa_y_dim_scale = nn.Parameter(torch.empty(hs))
+            nn.init.normal_(self.qa_q_dim_scale, mean=0.0, std=0.01)
+            nn.init.normal_(self.qa_y_dim_scale, mean=0.0, std=0.01)
+
+        if self.qv_variant == 'dynamic_qa_conditioned_lowrank16':
+            # shared low-rank Q+A mixer: 128->16->64 per head/token
+            self.qa_gate_lowrank_fc1 = nn.Linear(hs * 2, 16, bias=False)
+            self.qa_gate_lowrank_fc2 = nn.Linear(16, hs, bias=False)
+
+        # --------------------------------------------------
+        # Static / normalization controls
+        # --------------------------------------------------
+        if self.qv_variant == 'post_rmsnorm_y':
+            self.post_attn_norm = RMSNorm(hs)
+        elif self.qv_variant == 'static_gate':
+            self.static_gate_param = nn.Parameter(torch.zeros(config.n_embd))
+        elif self.qv_variant == 'static_gate_prehead':
+            self.static_gate_prehead_param = nn.Parameter(torch.zeros(config.n_head, hs))
+
+        self.q_norm = None
+        self.v_norm = None
+        self.qv_modulator = None
+        variant = getattr(config, 'qv_variant', 'none')
+        if variant in ['qvnorm']:
+            self.q_norm = RMSNorm(hs)
+        if variant in ['vnorm', 'qvnorm']:
+            self.v_norm = RMSNorm(hs)
+
+        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
+        self.flash = hasattr(torch.nn.functional, 'scaled_dot_product_attention')
+        if not self.flash:
+            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, x_prenorm=None):
+        B, T, C = x.size()
+        hs = C // self.n_head
+
+        q, k, v = self.c_attn(x).split(self.n_embd, dim=2)
+        k = k.view(B, T, self.n_head, hs).transpose(1, 2)
+        q = q.view(B, T, self.n_head, hs).transpose(1, 2)
+        v = v.view(B, T, self.n_head, hs).transpose(1, 2)
+
+        if self.q_norm is not None:
+            q = self.q_norm(q)
+        if self.v_norm is not None:
+            v = self.v_norm(v)
+
+        if self.flash:
+            y = torch.nn.functional.scaled_dot_product_attention(
+                q, k, v, is_causal=True, dropout_p=self.dropout if self.training else 0
+            )
+        else:
+            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
+
+        # --------------------------------------------------
+        # Gate application — y shape: (B, n_head, T, hs)
+        # --------------------------------------------------
+
+        # Q-conditioned gates
+        if self.qv_variant in ['dynamic', 'dynamic_swiglu']:
+            gate_logit = self.qv_gate_proj(q)
+            gate = torch.sigmoid(gate_logit) if self.qv_variant == 'dynamic' else gate_logit * torch.sigmoid(gate_logit)
+            y = y * gate
+
+        elif self.qv_variant == 'dynamic_qconditioned_mlp128':
+            gate_hidden = F.silu(self.q_gate_fc1(q))
+            gate_logit = self.q_gate_fc2(gate_hidden)
+            gate = torch.sigmoid(gate_logit)
+            y = y * gate
+
+        elif self.qv_variant == 'dynamic_qconditioned_mlp192':
+            gate_hidden = F.silu(self.q_gate_fc1_192(q))
+            gate_logit = self.q_gate_fc2_192(gate_hidden)
+            gate = torch.sigmoid(gate_logit)
+            y = y * gate
+
+        elif self.qv_variant == 'dynamic_qconditioned_fullwidth':
+            q_full = q.transpose(1, 2).contiguous().view(B, T, C)
+            gate_logit = self.q_gate_proj_full(q_full)
+            gate = torch.sigmoid(gate_logit)
+            gate = gate.view(B, T, self.n_head, hs).transpose(1, 2)
+            y = y * gate
+
+        elif self.qv_variant == 'dynamic_qconditioned_fullwidth_headspecific':
+            q_full = q.transpose(1, 2).contiguous().view(B, T, C)
+            gate_logit = torch.einsum('btc,hce->bhte', q_full, self.q_gate_proj_full_headspecific)
+            gate = torch.sigmoid(gate_logit)
+            y = y * gate
+
+        elif self.qv_variant == 'dynamic_q_headshared_elementwise':
+            gate_logit = self.q_gate_headshared(q)
+            gate = torch.sigmoid(gate_logit)
+            y = y * gate
+
+        # Static / normalization controls
+        elif self.qv_variant == 'post_rmsnorm_y':
+            y = self.post_attn_norm(y)
+
+        elif self.qv_variant == 'static_gate':
+            gate = torch.sigmoid(self.static_gate_param)
+            y = y.transpose(1, 2).contiguous().view(B, T, C)
+            y = y * gate[None, None, :]
+
+        elif self.qv_variant == 'static_gate_prehead':
+            gate = torch.sigmoid(self.static_gate_prehead_param)
+            y = y * gate[None, :, None, :]
+
+        # X-conditioned gates
+        elif self.qv_variant == 'dynamic_xconditioned_g1':
+            gate_logit = self.x_gate_proj(x_prenorm)
+            gate = torch.sigmoid(gate_logit)
+            gate = gate.view(B, T, self.n_head, hs).transpose(1, 2)
+            y = y * gate
+
+        elif self.qv_variant == 'dynamic_xconditioned_fullwidth_headspecific':
+            gate_logit = torch.einsum('btc,hce->bhte', x_prenorm, self.x_gate_proj_full_headspecific)
+            gate = torch.sigmoid(gate_logit)
+            y = y * gate
+
+        elif self.qv_variant == 'dynamic_xconditioned_bottleneck':
+            z = self.x_gate_bottleneck(x_prenorm)
+            gate = torch.sigmoid(z)
+            gate = gate.unsqueeze(1)
+            y = y * gate
+
+        elif self.qv_variant == 'dynamic_x_g1_headspecific_elementwise':
+            xh = x_prenorm.view(B, T, self.n_head, hs).transpose(1, 2)
+            gate_logit = torch.einsum('bhtd,hde->bhte', xh, self.x_gate_headspecific_elementwise)
+            gate = torch.sigmoid(gate_logit)
+            y = y * gate
+
+        elif self.qv_variant == 'dynamic_x_g1_headspecific_headwise':
+            xh = x_prenorm.view(B, T, self.n_head, hs).transpose(1, 2)
+            gate_logit = torch.einsum('bhtd,hde->bhte', xh, self.x_gate_headspecific_headwise)
+            gate = torch.sigmoid(gate_logit)
+            y = y * gate
+
+        elif self.qv_variant == 'dynamic_x_g1_headshared_elementwise':
+            xh = x_prenorm.view(B, T, self.n_head, hs).transpose(1, 2)
+            gate_logit = self.x_gate_headshared_elementwise(xh)
+            gate = torch.sigmoid(gate_logit)
+            y = y * gate
+
+        # Random / ablation gates — no learned params
+        elif self.qv_variant == 'dynamic_random_gate':
+            # Uniform random [0,1] gate — tests if location matters at all
+            gate = torch.rand_like(y)
+            y = y * gate
+
+        elif self.qv_variant == 'dynamic_ones_gate':
+            # Always 1.0 — pure identity, sanity check (should == baseline)
+            pass  # y unchanged
+
+        elif self.qv_variant == 'dynamic_random_normal':
+            # Normal(0.5, 0.2) gate clamped to [0,1] — different noise shape
+            gate = torch.randn_like(y) * 0.2 + 0.5
+            gate = gate.clamp(0.0, 1.0)
+            y = y * gate
+
+        elif self.qv_variant == 'dynamic_bernoulli_gate':
+            # Binary gate: randomly zeros 50% of head dims — spicy dropout variant
+            gate = torch.bernoulli(torch.full_like(y, 0.5))
+            # Scale by 2.0 to preserve expected value (like standard dropout)
+            y = y * gate * 2.0
+
+        # A-only gate
+        elif self.qv_variant == 'dynamic_a_conditioned':
+            gate_logit = self.a_gate_proj(y)
+            gate = torch.sigmoid(gate_logit)
+            y = y * gate
+
+        # Q+A dual-signal gates
+        elif self.qv_variant == 'dynamic_qa_conditioned':
+            # cheap Q+A: cat(q, y) -> shared Linear(128->64) -> sigmoid
+            gate_in = torch.cat([q, y], dim=-1)   # (B, n_head, T, 128)
+            gate_logit = self.qa_gate_proj(gate_in)
+            gate = torch.sigmoid(gate_logit)
+            y = y * gate
+
+        elif self.qv_variant == 'dynamic_qa_conditioned_headspecific':
+            # true head-specific Q+A: per-head (128->64) -> sigmoid
+            gate_in = torch.cat([q, y], dim=-1)   # (B, n_head, T, 128)
+            gate_logit = torch.einsum('bhtd,hde->bhte', gate_in, self.qa_gate_proj_headspecific)
+            gate = torch.sigmoid(gate_logit)
+            y = y * gate
+
+        elif self.qv_variant == 'dynamic_qa_conditioned_mlp128':
+            # main Q+A MLP: cat(q, y) -> Linear(128->128) -> SiLU -> Linear(128->64) -> sigmoid
+            gate_in = torch.cat([q, y], dim=-1)   # (B, n_head, T, 128)
+            gate_hidden = F.silu(self.qa_gate_fc1(gate_in))
+            gate_logit = self.qa_gate_fc2(gate_hidden)
+            gate = torch.sigmoid(gate_logit)
+            y = y * gate
+
+        elif self.qv_variant == 'dynamic_qa_headshared_elementwise':
+            # legacy compatibility variant: cat(q, y) -> shared Linear(128->64) -> sigmoid
+            gate_in = torch.cat([q, y], dim=-1)   # (B, n_head, T, 128)
+            gate_logit = self.qa_gate_headshared(gate_in)
+            gate = torch.sigmoid(gate_logit)
+            y = y * gate
+
+        elif self.qv_variant == 'dynamic_qa_bilinear_diag':
+            # learned diagonal bilinear interaction: sigmoid((q * y * w) / sqrt(d_head))
+            gate_logit = (q * y * self.qa_bilinear_diag[None, None, None, :]) / math.sqrt(hs)
+            gate = torch.sigmoid(gate_logit)
+            y = y * gate
+
+        elif self.qv_variant == 'dynamic_qa_conditioned_normed':
+            qn = self.qa_q_norm(q)
+            yn = self.qa_y_norm(y)
+            gate_in = torch.cat([qn, yn], dim=-1)
+            gate_logit = self.qa_gate_proj_normed(gate_in)
+            gate = torch.sigmoid(gate_logit)
+            y = y * gate
+
+        elif self.qv_variant == 'dynamic_qa_conditioned_normed_yonly':
+            yn = self.qa_y_norm_yonly(y)
+            gate_in = torch.cat([q, yn], dim=-1)
+            gate_logit = self.qa_gate_proj_normed_yonly(gate_in)
+            gate = torch.sigmoid(gate_logit)
+            y = y * gate
+
+        elif self.qv_variant == 'dynamic_qa_conditioned_normed_qonly':
+            qn = self.qa_q_norm_qonly(q)
+            gate_in = torch.cat([qn, y], dim=-1)
+            gate_logit = self.qa_gate_proj_normed_qonly(gate_in)
+            gate = torch.sigmoid(gate_logit)
+            y = y * gate
+
+        elif self.qv_variant == 'dynamic_qa_conditioned_softq':
+            qn = self.qa_q_norm_soft(q)
+            alpha = torch.sigmoid(self.qa_q_blend_alpha)
+            q_blend = alpha * qn + (1.0 - alpha) * q
+            gate_in = torch.cat([q_blend, y], dim=-1)
+            gate_logit = self.qa_gate_proj_softq(gate_in)
+            gate = torch.sigmoid(gate_logit)
+            y = y * gate
+
+        elif self.qv_variant == 'dynamic_qa_conditioned_softqa':
+            qn = self.qa_q_norm_softqa(q)
+            yn = self.qa_y_norm_softqa(y)
+            alpha_q = torch.sigmoid(self.qa_q_blend_alpha)
+            alpha_y = torch.sigmoid(self.qa_y_blend_alpha)
+            q_blend = alpha_q * qn + (1.0 - alpha_q) * q
+            y_blend = alpha_y * yn + (1.0 - alpha_y) * y
+            gate_in = torch.cat([q_blend, y_blend], dim=-1)
+            gate_logit = self.qa_gate_proj_softqa(gate_in)
+            gate = torch.sigmoid(gate_logit)
+            y = y * gate
+
+        elif self.qv_variant == 'dynamic_qa_conditioned_softqa_perlayer':
+            qn = self.qa_q_norm_softqa_pl(q)
+            yn = self.qa_y_norm_softqa_pl(y)
+            alpha_q = torch.sigmoid(self.qa_q_blend_alpha)
+            alpha_y = torch.sigmoid(self.qa_y_blend_alpha)
+            q_blend = alpha_q * qn + (1.0 - alpha_q) * q
+            y_blend = alpha_y * yn + (1.0 - alpha_y) * y
+            gate_in = torch.cat([q_blend, y_blend], dim=-1)
+            gate_logit = self.qa_gate_proj_softqa_pl(gate_in)
+            gate = torch.sigmoid(gate_logit)
+            y = y * gate
+
+        elif self.qv_variant in ['dynamic_qa_conditioned_softqa_perdim', 'dynamic_qa_conditioned_softqa_perdim_informed', 'dynamic_qa_conditioned_softqa_perdim_free']:
+            qn = self.qa_q_norm_softqa_pd(q)
+            yn = self.qa_y_norm_softqa_pd(y)
+            scale_q = torch.sigmoid(self.qa_q_dim_scale).to(dtype=q.dtype, device=q.device)
+            scale_y = torch.sigmoid(self.qa_y_dim_scale).to(dtype=y.dtype, device=y.device)
+            q_blend = scale_q * qn + (1.0 - scale_q) * q
+            y_blend = scale_y * yn + (1.0 - scale_y) * y
+            gate_in = torch.cat([q_blend, y_blend], dim=-1)
+            gate_logit = self.qa_gate_proj_softqa_pd(gate_in)
+            gate = torch.sigmoid(gate_logit)
+            y = y * gate
+
+        elif self.qv_variant == 'dynamic_qa_conditioned_lowrank16':
+            gate_in = torch.cat([q, y], dim=-1)
+            gate_hidden = F.silu(self.qa_gate_lowrank_fc1(gate_in))
+            gate_logit = self.qa_gate_lowrank_fc2(gate_hidden)
+            gate = torch.sigmoid(gate_logit)
+            y = y * gate
+
+        # Zero-parameter dot product gates — raw Q·y geometry
+        elif self.qv_variant == 'dynamic_dot_scalar':
+            # scalar gate: sigmoid(sum(q*y) / sqrt(d_head))
+            # one scalar per head per token, zero params
+            hs = q.shape[-1]
+            dot = (q * y).sum(dim=-1, keepdim=True)          # (B, n_head, T, 1)
+            gate = torch.sigmoid(dot / math.sqrt(hs))         # scalar broadcast
+            y = y * gate
+
+        elif self.qv_variant == 'dynamic_dot_elementwise':
+            # elementwise gate: sigmoid(q*y / sqrt(d_head))
+            # same shape as cheap_qa gate but zero params
+            hs = q.shape[-1]
+            gate = torch.sigmoid((q * y) / math.sqrt(hs))    # (B, n_head, T, 64)
+            y = y * gate
+
+        if self.qv_variant != 'static_gate':
+            y = y.transpose(1, 2).contiguous().view(B, T, C)
+
+        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_norm = self.ln_1(x)
+        attn_out = self.attn(x_norm, x_norm)
+        x = x + attn_out
+        x = x + self.mlp(self.ln_2(x))
+        return x
+
+@dataclass
+class GPTConfig:
+    block_size: int = 1024
+    vocab_size: int = 50304
+    n_layer: int = 12
+    n_head: int = 12
+    n_embd: int = 768
+    dropout: float = 0.0
+    bias: bool = True
+    qv_variant: str = 'none'
+
+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)
+        self.transformer.wte.weight = self.lm_head.weight
+        self.apply(self._init_weights)
+        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))
+        print("number of parameters: %.2fM" % (self.get_num_params()/1e6,))
+
+    def get_num_params(self, non_embedding=True):
+        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
+        pos = torch.arange(0, t, dtype=torch.long, device=device)
+        tok_emb = self.transformer.wte(idx)
+        pos_emb = self.transformer.wpe(pos)
+        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:
+            logits = self.lm_head(x)
+            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1)
+        else:
+            logits = self.lm_head(x[:, [-1], :])
+            loss = None
+        return logits, loss
+
+    def crop_block_size(self, block_size):
+        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 {}
+        assert all(k == 'dropout' for k in override_args)
+        from transformers import GPT2LMHeadModel
+        print("loading weights from pretrained gpt: %s" % model_type)
+        config_args = {
+            'gpt2':       dict(n_layer=12, n_head=12, n_embd=768),
+            'gpt2-medium': dict(n_layer=24, n_head=16, n_embd=1024),
+            'gpt2-large': dict(n_layer=36, n_head=20, n_embd=1280),
+            'gpt2-xl':    dict(n_layer=48, n_head=25, n_embd=1600),
+        }[model_type]
+        config_args['vocab_size'] = 50257
+        config_args['block_size'] = 1024
+        config_args['bias'] = True
+        if 'dropout' in override_args:
+            config_args['dropout'] = override_args['dropout']
+        config = GPTConfig(**config_args)
+        model = GPT(config)
+        sd = model.state_dict()
+        sd_keys = [k for k in sd.keys() if not k.endswith('.attn.bias')]
+        model_hf = GPT2LMHeadModel.from_pretrained(model_type)
+        sd_hf = model_hf.state_dict()
+        sd_keys_hf = [k for k in sd_hf.keys() if not k.endswith('.attn.masked_bias') and not k.endswith('.attn.bias')]
+        transposed = ['attn.c_attn.weight', 'attn.c_proj.weight', 'mlp.c_fc.weight', 'mlp.c_proj.weight']
+        assert len(sd_keys_hf) == len(sd_keys)
+        for k in sd_keys_hf:
+            if any(k.endswith(w) for w in transposed):
+                assert sd_hf[k].shape[::-1] == sd[k].shape
+                with torch.no_grad():
+                    sd[k].copy_(sd_hf[k].t())
+            else:
+                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):
+        param_dict = {pn: p for pn, p in self.named_parameters() if p.requires_grad}
+        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")
+        fused_available = 'fused' in inspect.signature(torch.optim.AdamW).parameters
+        use_fused = fused_available and device_type == 'cuda'
+        optimizer = torch.optim.AdamW(optim_groups, lr=learning_rate, betas=betas, fused=use_fused)
+        print(f"using fused AdamW: {use_fused}")
+        return optimizer
+
+    def estimate_mfu(self, fwdbwd_per_iter, dt):
+        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
+        flops_achieved = flops_per_iter * (1.0 / dt)
+        flops_promised = 312e12
+        mfu = flops_achieved / flops_promised
+        return mfu
+
+    @torch.no_grad()
+    def generate(self, idx, max_new_tokens, temperature=1.0, top_k=None):
+        for _ in range(max_new_tokens):
+            idx_cond = idx if idx.size(1) <= self.config.block_size else idx[:, -self.config.block_size:]
+            logits, _ = self(idx_cond)
+            logits = logits[:, -1, :] / temperature
+            if top_k is not None:
+                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
+                logits[logits < v[:, [-1]]] = -float('Inf')
+            probs = F.softmax(logits, dim=-1)
+            idx_next = torch.multinomial(probs, num_samples=1)
+            idx = torch.cat((idx, idx_next), dim=1)
+        return idx

code/model_cheap_qa_minimal.py

@@ -0,0 +1,248 @@
+"""
+Minimal single-variant GPT — cheap_qa champion ONLY.
+
+Distilled from the 19-variant env_canon model.py. Every other gate branch has been
+removed; the ONLY conditioning path is the cheap_qa gate:
+
+    cheap_qa  ==  qv_variant "dynamic_qa_conditioned"
+    gate = sigmoid( Linear(2*head_dim -> head_dim) ( cat(q, y) ) )
+    y = y * gate                      # applied at G1: post-SDPA, pre-W_O (c_proj)
+
+Champion result (NanoGPT 124M, OpenWebText, 10k iters, seeds 42/1337/123):
+    baseline 4.4708 -> cheap_qa 4.3897  (-0.0811 CE, -1.81%), 98,304 params (+0.079%).
+
+>>> THE G1 SEAM IS MARKED BELOW WITH:  # === G1 SEAM ===
+    That block is the entire port surface. Resize the gate for a different attention
+    by changing head_dim (hs); the gate is always Linear(2*hs -> hs) on cat(q, y).
+"""
+
+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):
+    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
+        self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd, bias=config.bias)
+        self.c_proj = nn.Linear(config.n_embd, config.n_embd, bias=config.bias)   # <- W_O
+
+        hs = config.n_embd // config.n_head   # head_dim (64 for the 124M config)
+
+        # === cheap_qa gate parameters ===
+        # cat(q, y) has width 2*hs (=128 when hs=64); gate maps it back to hs (=64).
+        # This is the ONLY learned gate in this minimal model.
+        self.qa_gate_proj = nn.Linear(hs * 2, hs, bias=False)   # Linear(128 -> 64)
+
+        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
+        self.flash = hasattr(torch.nn.functional, 'scaled_dot_product_attention')
+        if not self.flash:
+            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()
+        hs = C // self.n_head
+
+        # Project to q, k, v and split into heads. q and v are (B, n_head, T, hs).
+        q, k, v = self.c_attn(x).split(self.n_embd, dim=2)
+        k = k.view(B, T, self.n_head, hs).transpose(1, 2)
+        q = q.view(B, T, self.n_head, hs).transpose(1, 2)
+        v = v.view(B, T, self.n_head, hs).transpose(1, 2)
+
+        # Scaled dot-product attention -> y is the attention output (the "A" in Q+A).
+        if self.flash:
+            y = torch.nn.functional.scaled_dot_product_attention(
+                q, k, v, is_causal=True, dropout_p=self.dropout if self.training else 0
+            )
+        else:
+            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
+
+        # =========================== G1 SEAM ============================
+        # cheap_qa gate. Inputs available here: q and y, both (B, n_head, T, hs).
+        # Location: strictly AFTER SDPA, strictly BEFORE the W_O projection (c_proj).
+        # To port to a different attention: keep this exact 4-line body; only hs changes.
+        gate_in = torch.cat([q, y], dim=-1)      # (B, n_head, T, 2*hs)  == cat(query, attn_output)
+        gate_logit = self.qa_gate_proj(gate_in)  # Linear(2*hs -> hs)
+        gate = torch.sigmoid(gate_logit)         # (B, n_head, T, hs) in (0, 1)
+        y = y * gate                             # elementwise gate on the attention output
+        # ========================= END G1 SEAM ==========================
+
+        # Merge heads back and apply the output projection W_O.
+        y = y.transpose(1, 2).contiguous().view(B, T, C)
+        y = self.resid_dropout(self.c_proj(y))   # c_proj == W_O
+        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   # NOTE: 50304 (GPT-2 50257 rounded up for efficiency)
+    n_layer: int = 12
+    n_head: int = 12
+    n_embd: int = 768         # -> head_dim = 768 / 12 = 64, so the gate is Linear(128 -> 64)
+    dropout: float = 0.0
+    bias: bool = True
+    # qv_variant is fixed to cheap_qa in this minimal model; no field needed.
+
+
+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)
+        self.transformer.wte.weight = self.lm_head.weight   # weight tying
+        self.apply(self._init_weights)
+        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))
+        print("number of parameters: %.2fM" % (self.get_num_params() / 1e6,))
+
+    def get_num_params(self, non_embedding=True):
+        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
+        pos = torch.arange(0, t, dtype=torch.long, device=device)
+        tok_emb = self.transformer.wte(idx)
+        pos_emb = self.transformer.wpe(pos)
+        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:
+            logits = self.lm_head(x)
+            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1)
+        else:
+            logits = self.lm_head(x[:, [-1], :])
+            loss = None
+        return logits, loss
+
+    def crop_block_size(self, block_size):
+        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]
+
+    def configure_optimizers(self, weight_decay, learning_rate, betas, device_type):
+        param_dict = {pn: p for pn, p in self.named_parameters() if p.requires_grad}
+        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")
+        fused_available = 'fused' in inspect.signature(torch.optim.AdamW).parameters
+        use_fused = fused_available and device_type == 'cuda'
+        optimizer = torch.optim.AdamW(optim_groups, lr=learning_rate, betas=betas, fused=use_fused)
+        print(f"using fused AdamW: {use_fused}")
+        return optimizer
+
+    def estimate_mfu(self, fwdbwd_per_iter, dt):
+        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
+        flops_achieved = flops_per_iter * (1.0 / dt)
+        flops_promised = 312e12
+        return flops_achieved / flops_promised
+
+    @torch.no_grad()
+    def generate(self, idx, max_new_tokens, temperature=1.0, top_k=None):
+        for _ in range(max_new_tokens):
+            idx_cond = idx if idx.size(1) <= self.config.block_size else idx[:, -self.config.block_size:]
+            logits, _ = self(idx_cond)
+            logits = logits[:, -1, :] / temperature
+            if top_k is not None:
+                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
+                logits[logits < v[:, [-1]]] = -float('Inf')
+            probs = F.softmax(logits, dim=-1)
+            idx_next = torch.multinomial(probs, num_samples=1)
+            idx = torch.cat((idx, idx_next), dim=1)
+        return idx

code/model_wslbridge.py

@@ -0,0 +1,604 @@
+"""
+Full definition of a GPT Language Model, all of it in this single file.
+Includes:
+- cheap Q gate
+- Q-MLP gate (128 hidden)
+- Q-MLP gate (192 hidden) -- param-matched control for Q+A MLP
+- full-width Q gate
+- Q head-shared elementwise gate
+- full-width X gate
+- bottleneck X gate
+- paper-faithful X variants:
+  - head-specific elementwise G1
+  - head-specific headwise G1
+  - head-shared elementwise G1
+- Q+A dual-signal variants:
+  - cheap Q+A gate (shared Linear 128->64 applied per head/token)
+  - Q+A MLP gate (shared MLP 128->128->64 applied per head/token)
+  - legacy Q+A head-shared elementwise gate (shared Linear 128->64; kept for compatibility)
+  - bilinear diagonal Q+A gate (learned elementwise q*y interaction)
+  - normed Q+A gate (RMSNorm on q and y before shared 128->64 mix)
+  - lowrank Q+A gate (shared 128->16->64 bottleneck mix)
+  - normed cheap Q gate (RMSNorm on q before shared 64->64 mix)
+  - normed Q-MLP gate (RMSNorm on q before shared 64->128->64 MLP)
+"""
+
+import math
+import inspect
+from dataclasses import dataclass
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-6):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+    def forward(self, x):
+        rms = torch.sqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
+        return x / rms * self.weight
+
+class LayerNorm(nn.Module):
+    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
+        self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd, bias=config.bias)
+        self.c_proj = nn.Linear(config.n_embd, config.n_embd, bias=config.bias)
+
+        self.qv_variant = config.qv_variant
+        hs = config.n_embd // config.n_head
+
+        # --------------------------------------------------
+        # Q-conditioned gates
+        # --------------------------------------------------
+        if self.qv_variant in ['dynamic', 'dynamic_swiglu']:
+            self.qv_gate_proj = nn.Linear(hs, hs, bias=False)
+
+        if self.qv_variant == 'dynamic_qconditioned_mlp128':
+            self.q_gate_fc1 = nn.Linear(hs, 128, bias=False)
+            self.q_gate_fc2 = nn.Linear(128, hs, bias=False)
+
+        if self.qv_variant == 'dynamic_qconditioned_normed':
+            self.q_gate_norm = RMSNorm(hs)
+            self.q_gate_proj_normed = nn.Linear(hs, hs, bias=False)
+
+        if self.qv_variant == 'dynamic_qconditioned_mlp128_normed':
+            self.q_gate_norm_mlp128 = RMSNorm(hs)
+            self.q_gate_fc1_normed = nn.Linear(hs, 128, bias=False)
+            self.q_gate_fc2_normed = nn.Linear(128, hs, bias=False)
+
+        if self.qv_variant == 'dynamic_qconditioned_mlp192':
+            self.q_gate_fc1_192 = nn.Linear(hs, 192, bias=False)
+            self.q_gate_fc2_192 = nn.Linear(192, hs, bias=False)
+
+        if self.qv_variant == 'dynamic_qconditioned_fullwidth':
+            self.q_gate_proj_full = nn.Linear(config.n_embd, config.n_embd, bias=False)
+
+        if self.qv_variant == 'dynamic_qconditioned_fullwidth_headspecific':
+            # full-width Q source (768) with separate 768->64 map for each head
+            self.q_gate_proj_full_headspecific = nn.Parameter(torch.empty(config.n_head, config.n_embd, hs))
+            nn.init.normal_(self.q_gate_proj_full_headspecific, mean=0.0, std=0.02)
+
+        if self.qv_variant == 'dynamic_q_headshared_elementwise':
+            self.q_gate_headshared = nn.Linear(hs, hs, bias=False)
+
+        # --------------------------------------------------
+        # X-conditioned gates
+        # --------------------------------------------------
+        if self.qv_variant == 'dynamic_xconditioned_g1':
+            self.x_gate_proj = nn.Linear(config.n_embd, config.n_embd, bias=False)
+
+        if self.qv_variant == 'dynamic_xconditioned_fullwidth_headspecific':
+            # full-width X source (768) with separate 768->64 map for each head
+            self.x_gate_proj_full_headspecific = nn.Parameter(torch.empty(config.n_head, config.n_embd, hs))
+            nn.init.normal_(self.x_gate_proj_full_headspecific, mean=0.0, std=0.02)
+
+        if self.qv_variant == 'dynamic_xconditioned_bottleneck':
+            self.x_gate_bottleneck = nn.Linear(config.n_embd, hs, bias=False)
+
+        if self.qv_variant == 'dynamic_x_g1_headspecific_elementwise':
+            self.x_gate_headspecific_elementwise = nn.Parameter(torch.empty(config.n_head, hs, hs))
+            nn.init.normal_(self.x_gate_headspecific_elementwise, mean=0.0, std=0.02)
+
+        if self.qv_variant == 'dynamic_x_g1_headspecific_headwise':
+            self.x_gate_headspecific_headwise = nn.Parameter(torch.empty(config.n_head, hs, 1))
+            nn.init.normal_(self.x_gate_headspecific_headwise, mean=0.0, std=0.02)
+
+        if self.qv_variant == 'dynamic_x_g1_headshared_elementwise':
+            self.x_gate_headshared_elementwise = nn.Linear(hs, hs, bias=False)
+
+        # --------------------------------------------------
+        # Q+A dual-signal gates
+        # --------------------------------------------------
+        # random gates have no learned parameters — nothing to init here
+        # dot product gates also have no learned parameters — nothing to init here
+
+        if self.qv_variant == 'dynamic_a_conditioned':
+            # A-only: per-head Linear(64->64), conditioned on y only
+            self.a_gate_proj = nn.Linear(hs, hs, bias=False)
+
+        if self.qv_variant == 'dynamic_qa_conditioned':
+            # cheap Q+A: shared Linear(128->64) applied to each head/token
+            self.qa_gate_proj = nn.Linear(hs * 2, hs, bias=False)
+
+        if self.qv_variant == 'dynamic_qa_conditioned_headspecific':
+            # true head-specific Q+A: separate 128->64 matrix for each head
+            self.qa_gate_proj_headspecific = nn.Parameter(torch.empty(config.n_head, hs * 2, hs))
+            nn.init.normal_(self.qa_gate_proj_headspecific, mean=0.0, std=0.02)
+
+        if self.qv_variant == 'dynamic_qa_conditioned_mlp128':
+            # main Q+A MLP: per-head MLP(128->128->64)
+            self.qa_gate_fc1 = nn.Linear(hs * 2, 128, bias=False)
+            self.qa_gate_fc2 = nn.Linear(128, hs, bias=False)
+
+        if self.qv_variant == 'dynamic_qa_headshared_elementwise':
+            # legacy compatibility variant: shared Linear(128->64) applied per head/token
+            self.qa_gate_headshared = nn.Linear(hs * 2, hs, bias=False)
+
+        if self.qv_variant == 'dynamic_qa_bilinear_diag':
+            # learned elementwise interaction on q*y; 64 params per layer
+            self.qa_bilinear_diag = nn.Parameter(torch.ones(hs))
+
+        if self.qv_variant == 'dynamic_qa_conditioned_normed':
+            # RMSNorm q and y separately before shared Linear(128->64)
+            self.qa_q_norm = RMSNorm(hs)
+            self.qa_y_norm = RMSNorm(hs)
+            self.qa_gate_proj_normed = nn.Linear(hs * 2, hs, bias=False)
+
+        if self.qv_variant == 'dynamic_qa_conditioned_lowrank16':
+            # shared low-rank Q+A mixer: 128->16->64 per head/token
+            self.qa_gate_lowrank_fc1 = nn.Linear(hs * 2, 16, bias=False)
+            self.qa_gate_lowrank_fc2 = nn.Linear(16, hs, bias=False)
+
+        # --------------------------------------------------
+        # Static / normalization controls
+        # --------------------------------------------------
+        if self.qv_variant == 'post_rmsnorm_y':
+            self.post_attn_norm = RMSNorm(hs)
+        elif self.qv_variant == 'static_gate':
+            self.static_gate_param = nn.Parameter(torch.zeros(config.n_embd))
+        elif self.qv_variant == 'static_gate_prehead':
+            self.static_gate_prehead_param = nn.Parameter(torch.zeros(config.n_head, hs))
+
+        self.q_norm = None
+        self.v_norm = None
+        self.qv_modulator = None
+        variant = getattr(config, 'qv_variant', 'none')
+        if variant in ['qvnorm']:
+            self.q_norm = RMSNorm(hs)
+        if variant in ['vnorm', 'qvnorm']:
+            self.v_norm = RMSNorm(hs)
+
+        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
+        self.flash = hasattr(torch.nn.functional, 'scaled_dot_product_attention')
+        if not self.flash:
+            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, x_prenorm=None):
+        B, T, C = x.size()
+        hs = C // self.n_head
+
+        q, k, v = self.c_attn(x).split(self.n_embd, dim=2)
+        k = k.view(B, T, self.n_head, hs).transpose(1, 2)
+        q = q.view(B, T, self.n_head, hs).transpose(1, 2)
+        v = v.view(B, T, self.n_head, hs).transpose(1, 2)
+
+        if self.q_norm is not None:
+            q = self.q_norm(q)
+        if self.v_norm is not None:
+            v = self.v_norm(v)
+
+        if self.flash:
+            y = torch.nn.functional.scaled_dot_product_attention(
+                q, k, v, is_causal=True, dropout_p=self.dropout if self.training else 0
+            )
+        else:
+            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
+
+        # --------------------------------------------------
+        # Gate application — y shape: (B, n_head, T, hs)
+        # --------------------------------------------------
+
+        # Q-conditioned gates
+        if self.qv_variant in ['dynamic', 'dynamic_swiglu']:
+            gate_logit = self.qv_gate_proj(q)
+            gate = torch.sigmoid(gate_logit) if self.qv_variant == 'dynamic' else gate_logit * torch.sigmoid(gate_logit)
+            y = y * gate
+
+        elif self.qv_variant == 'dynamic_qconditioned_mlp128':
+            gate_hidden = F.silu(self.q_gate_fc1(q))
+            gate_logit = self.q_gate_fc2(gate_hidden)
+            gate = torch.sigmoid(gate_logit)
+            y = y * gate
+
+        elif self.qv_variant == 'dynamic_qconditioned_normed':
+            qn = self.q_gate_norm(q)
+            gate_logit = self.q_gate_proj_normed(qn)
+            gate = torch.sigmoid(gate_logit)
+            y = y * gate
+
+        elif self.qv_variant == 'dynamic_qconditioned_mlp128_normed':
+            qn = self.q_gate_norm_mlp128(q)
+            gate_hidden = F.silu(self.q_gate_fc1_normed(qn))
+            gate_logit = self.q_gate_fc2_normed(gate_hidden)
+            gate = torch.sigmoid(gate_logit)
+            y = y * gate
+
+        elif self.qv_variant == 'dynamic_qconditioned_mlp192':
+            gate_hidden = F.silu(self.q_gate_fc1_192(q))
+            gate_logit = self.q_gate_fc2_192(gate_hidden)
+            gate = torch.sigmoid(gate_logit)
+            y = y * gate
+
+        elif self.qv_variant == 'dynamic_qconditioned_fullwidth':
+            q_full = q.transpose(1, 2).contiguous().view(B, T, C)
+            gate_logit = self.q_gate_proj_full(q_full)
+            gate = torch.sigmoid(gate_logit)
+            gate = gate.view(B, T, self.n_head, hs).transpose(1, 2)
+            y = y * gate
+
+        elif self.qv_variant == 'dynamic_qconditioned_fullwidth_headspecific':
+            q_full = q.transpose(1, 2).contiguous().view(B, T, C)
+            gate_logit = torch.einsum('btc,hce->bhte', q_full, self.q_gate_proj_full_headspecific)
+            gate = torch.sigmoid(gate_logit)
+            y = y * gate
+
+        elif self.qv_variant == 'dynamic_q_headshared_elementwise':
+            gate_logit = self.q_gate_headshared(q)
+            gate = torch.sigmoid(gate_logit)
+            y = y * gate
+
+        # Static / normalization controls
+        elif self.qv_variant == 'post_rmsnorm_y':
+            y = self.post_attn_norm(y)
+
+        elif self.qv_variant == 'static_gate':
+            gate = torch.sigmoid(self.static_gate_param)
+            y = y.transpose(1, 2).contiguous().view(B, T, C)
+            y = y * gate[None, None, :]
+
+        elif self.qv_variant == 'static_gate_prehead':
+            gate = torch.sigmoid(self.static_gate_prehead_param)
+            y = y * gate[None, :, None, :]
+
+        # X-conditioned gates
+        elif self.qv_variant == 'dynamic_xconditioned_g1':
+            gate_logit = self.x_gate_proj(x_prenorm)
+            gate = torch.sigmoid(gate_logit)
+            gate = gate.view(B, T, self.n_head, hs).transpose(1, 2)
+            y = y * gate
+
+        elif self.qv_variant == 'dynamic_xconditioned_fullwidth_headspecific':
+            gate_logit = torch.einsum('btc,hce->bhte', x_prenorm, self.x_gate_proj_full_headspecific)
+            gate = torch.sigmoid(gate_logit)
+            y = y * gate
+
+        elif self.qv_variant == 'dynamic_xconditioned_bottleneck':
+            z = self.x_gate_bottleneck(x_prenorm)
+            gate = torch.sigmoid(z)
+            gate = gate.unsqueeze(1)
+            y = y * gate
+
+        elif self.qv_variant == 'dynamic_x_g1_headspecific_elementwise':
+            xh = x_prenorm.view(B, T, self.n_head, hs).transpose(1, 2)
+            gate_logit = torch.einsum('bhtd,hde->bhte', xh, self.x_gate_headspecific_elementwise)
+            gate = torch.sigmoid(gate_logit)
+            y = y * gate
+
+        elif self.qv_variant == 'dynamic_x_g1_headspecific_headwise':
+            xh = x_prenorm.view(B, T, self.n_head, hs).transpose(1, 2)
+            gate_logit = torch.einsum('bhtd,hde->bhte', xh, self.x_gate_headspecific_headwise)
+            gate = torch.sigmoid(gate_logit)
+            y = y * gate
+
+        elif self.qv_variant == 'dynamic_x_g1_headshared_elementwise':
+            xh = x_prenorm.view(B, T, self.n_head, hs).transpose(1, 2)
+            gate_logit = self.x_gate_headshared_elementwise(xh)
+            gate = torch.sigmoid(gate_logit)
+            y = y * gate
+
+        # Random / ablation gates — no learned params
+        elif self.qv_variant == 'dynamic_random_gate':
+            # Uniform random [0,1] gate — tests if location matters at all
+            gate = torch.rand_like(y)
+            y = y * gate
+
+        elif self.qv_variant == 'dynamic_ones_gate':
+            # Always 1.0 — pure identity, sanity check (should == baseline)
+            pass  # y unchanged
+
+        elif self.qv_variant == 'dynamic_random_normal':
+            # Normal(0.5, 0.2) gate clamped to [0,1] — different noise shape
+            gate = torch.randn_like(y) * 0.2 + 0.5
+            gate = gate.clamp(0.0, 1.0)
+            y = y * gate
+
+        elif self.qv_variant == 'dynamic_bernoulli_gate':
+            # Binary gate: randomly zeros 50% of head dims — spicy dropout variant
+            gate = torch.bernoulli(torch.full_like(y, 0.5))
+            # Scale by 2.0 to preserve expected value (like standard dropout)
+            y = y * gate * 2.0
+
+        # A-only gate
+        elif self.qv_variant == 'dynamic_a_conditioned':
+            gate_logit = self.a_gate_proj(y)
+            gate = torch.sigmoid(gate_logit)
+            y = y * gate
+
+        # Q+A dual-signal gates
+        elif self.qv_variant == 'dynamic_qa_conditioned':
+            # cheap Q+A: cat(q, y) -> shared Linear(128->64) -> sigmoid
+            gate_in = torch.cat([q, y], dim=-1)   # (B, n_head, T, 128)
+            gate_logit = self.qa_gate_proj(gate_in)
+            gate = torch.sigmoid(gate_logit)
+            y = y * gate
+
+        elif self.qv_variant == 'dynamic_qa_conditioned_headspecific':
+            # true head-specific Q+A: per-head (128->64) -> sigmoid
+            gate_in = torch.cat([q, y], dim=-1)   # (B, n_head, T, 128)
+            gate_logit = torch.einsum('bhtd,hde->bhte', gate_in, self.qa_gate_proj_headspecific)
+            gate = torch.sigmoid(gate_logit)
+            y = y * gate
+
+        elif self.qv_variant == 'dynamic_qa_conditioned_mlp128':
+            # main Q+A MLP: cat(q, y) -> Linear(128->128) -> SiLU -> Linear(128->64) -> sigmoid
+            gate_in = torch.cat([q, y], dim=-1)   # (B, n_head, T, 128)
+            gate_hidden = F.silu(self.qa_gate_fc1(gate_in))
+            gate_logit = self.qa_gate_fc2(gate_hidden)
+            gate = torch.sigmoid(gate_logit)
+            y = y * gate
+
+        elif self.qv_variant == 'dynamic_qa_headshared_elementwise':
+            # legacy compatibility variant: cat(q, y) -> shared Linear(128->64) -> sigmoid
+            gate_in = torch.cat([q, y], dim=-1)   # (B, n_head, T, 128)
+            gate_logit = self.qa_gate_headshared(gate_in)
+            gate = torch.sigmoid(gate_logit)
+            y = y * gate
+
+        elif self.qv_variant == 'dynamic_qa_bilinear_diag':
+            # learned diagonal bilinear interaction: sigmoid((q * y * w) / sqrt(d_head))
+            gate_logit = (q * y * self.qa_bilinear_diag[None, None, None, :]) / math.sqrt(hs)
+            gate = torch.sigmoid(gate_logit)
+            y = y * gate
+
+        elif self.qv_variant == 'dynamic_qa_conditioned_normed':
+            qn = self.qa_q_norm(q)
+            yn = self.qa_y_norm(y)
+            gate_in = torch.cat([qn, yn], dim=-1)
+            gate_logit = self.qa_gate_proj_normed(gate_in)
+            gate = torch.sigmoid(gate_logit)
+            y = y * gate
+
+        elif self.qv_variant == 'dynamic_qa_conditioned_lowrank16':
+            gate_in = torch.cat([q, y], dim=-1)
+            gate_hidden = F.silu(self.qa_gate_lowrank_fc1(gate_in))
+            gate_logit = self.qa_gate_lowrank_fc2(gate_hidden)
+            gate = torch.sigmoid(gate_logit)
+            y = y * gate
+
+        # Zero-parameter dot product gates — raw Q·y geometry
+        elif self.qv_variant == 'dynamic_dot_scalar':
+            # scalar gate: sigmoid(sum(q*y) / sqrt(d_head))
+            # one scalar per head per token, zero params
+            hs = q.shape[-1]
+            dot = (q * y).sum(dim=-1, keepdim=True)          # (B, n_head, T, 1)
+            gate = torch.sigmoid(dot / math.sqrt(hs))         # scalar broadcast
+            y = y * gate
+
+        elif self.qv_variant == 'dynamic_dot_elementwise':
+            # elementwise gate: sigmoid(q*y / sqrt(d_head))
+            # same shape as cheap_qa gate but zero params
+            hs = q.shape[-1]
+            gate = torch.sigmoid((q * y) / math.sqrt(hs))    # (B, n_head, T, 64)
+            y = y * gate
+
+        if self.qv_variant != 'static_gate':
+            y = y.transpose(1, 2).contiguous().view(B, T, C)
+
+        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_norm = self.ln_1(x)
+        attn_out = self.attn(x_norm, x_norm)
+        x = x + attn_out
+        x = x + self.mlp(self.ln_2(x))
+        return x
+
+@dataclass
+class GPTConfig:
+    block_size: int = 1024
+    vocab_size: int = 50304
+    n_layer: int = 12
+    n_head: int = 12
+    n_embd: int = 768
+    dropout: float = 0.0
+    bias: bool = True
+    qv_variant: str = 'none'
+
+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)
+        self.transformer.wte.weight = self.lm_head.weight
+        self.apply(self._init_weights)
+        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))
+        print("number of parameters: %.2fM" % (self.get_num_params()/1e6,))
+
+    def get_num_params(self, non_embedding=True):
+        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
+        pos = torch.arange(0, t, dtype=torch.long, device=device)
+        tok_emb = self.transformer.wte(idx)
+        pos_emb = self.transformer.wpe(pos)
+        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:
+            logits = self.lm_head(x)
+            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1)
+        else:
+            logits = self.lm_head(x[:, [-1], :])
+            loss = None
+        return logits, loss
+
+    def crop_block_size(self, block_size):
+        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 {}
+        assert all(k == 'dropout' for k in override_args)
+        from transformers import GPT2LMHeadModel
+        print("loading weights from pretrained gpt: %s" % model_type)
+        config_args = {
+            'gpt2':       dict(n_layer=12, n_head=12, n_embd=768),
+            'gpt2-medium': dict(n_layer=24, n_head=16, n_embd=1024),
+            'gpt2-large': dict(n_layer=36, n_head=20, n_embd=1280),
+            'gpt2-xl':    dict(n_layer=48, n_head=25, n_embd=1600),
+        }[model_type]
+        config_args['vocab_size'] = 50257
+        config_args['block_size'] = 1024
+        config_args['bias'] = True
+        if 'dropout' in override_args:
+            config_args['dropout'] = override_args['dropout']
+        config = GPTConfig(**config_args)
+        model = GPT(config)
+        sd = model.state_dict()
+        sd_keys = [k for k in sd.keys() if not k.endswith('.attn.bias')]
+        model_hf = GPT2LMHeadModel.from_pretrained(model_type)
+        sd_hf = model_hf.state_dict()
+        sd_keys_hf = [k for k in sd_hf.keys() if not k.endswith('.attn.masked_bias') and not k.endswith('.attn.bias')]
+        transposed = ['attn.c_attn.weight', 'attn.c_proj.weight', 'mlp.c_fc.weight', 'mlp.c_proj.weight']
+        assert len(sd_keys_hf) == len(sd_keys)
+        for k in sd_keys_hf:
+            if any(k.endswith(w) for w in transposed):
+                assert sd_hf[k].shape[::-1] == sd[k].shape
+                with torch.no_grad():
+                    sd[k].copy_(sd_hf[k].t())
+            else:
+                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):
+        param_dict = {pn: p for pn, p in self.named_parameters() if p.requires_grad}
+        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")
+        fused_available = 'fused' in inspect.signature(torch.optim.AdamW).parameters
+        use_fused = fused_available and device_type == 'cuda'
+        optimizer = torch.optim.AdamW(optim_groups, lr=learning_rate, betas=betas, fused=use_fused)
+        print(f"using fused AdamW: {use_fused}")
+        return optimizer
+
+    def estimate_mfu(self, fwdbwd_per_iter, dt):
+        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
+        flops_achieved = flops_per_iter * (1.0 / dt)
+        flops_promised = 312e12
+        mfu = flops_achieved / flops_promised
+        return mfu
+
+    @torch.no_grad()
+    def generate(self, idx, max_new_tokens, temperature=1.0, top_k=None):
+        for _ in range(max_new_tokens):
+            idx_cond = idx if idx.size(1) <= self.config.block_size else idx[:, -self.config.block_size:]
+            logits, _ = self(idx_cond)
+            logits = logits[:, -1, :] / temperature
+            if top_k is not None:
+                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
+                logits[logits < v[:, [-1]]] = -float('Inf')
+            probs = F.softmax(logits, dim=-1)
+            idx_next = torch.multinomial(probs, num_samples=1)
+            idx = torch.cat((idx, idx_next), dim=1)
+        return idx

code/train_canon_HEAD.py

@@ -0,0 +1,308 @@
+"""
+Training script with qv_variant choices extended for Phase 6 Q+A variants (normed + lowrank).
+"""
+
+import os
+import time
+import math
+import pickle
+import argparse
+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
+
+out_dir = 'out-shakespeare-char'
+eval_interval = 250
+log_interval = 10
+eval_iters = 200
+eval_only = False
+always_save_checkpoint = False
+wandb_log = False
+wandb_project = 'shakespeare-char'
+wandb_run_name = 'mini-gpt'
+dataset = 'shakespeare_char'
+gradient_accumulation_steps = 1
+batch_size = 64
+block_size = 256
+n_layer = 6
+n_head = 6
+n_embd = 384
+dropout = 0.2
+bias = False
+learning_rate = 1e-3
+max_iters = 5000
+lr_decay_iters = 5000
+min_lr = 1e-4
+beta1 = 0.9
+beta2 = 0.99
+warmup_iters = 100
+device = 'cuda'
+dtype = 'bfloat16' if torch.cuda.is_available() and torch.cuda.is_bf16_supported() else 'float16'
+compile = False
+
+config_keys = [k for k, v in globals().items() if not k.startswith('_') and isinstance(v, (int, float, bool, str))]
+parser = argparse.ArgumentParser(description='Train a mini-GPT')
+parser.add_argument('--out_dir', type=str, default=out_dir)
+parser.add_argument('--eval_interval', type=int, default=eval_interval)
+parser.add_argument('--log_interval', type=int, default=log_interval)
+parser.add_argument('--eval_iters', type=int, default=eval_iters)
+parser.add_argument('--eval_only', action='store_true')
+parser.add_argument('--always_save_checkpoint', action='store_true')
+parser.add_argument('--wandb_log', action='store_true')
+parser.add_argument('--wandb_project', type=str, default=wandb_project)
+parser.add_argument('--wandb_run_name', type=str, default=wandb_run_name)
+parser.add_argument('--dataset', type=str, default=dataset)
+parser.add_argument('--gradient_accumulation_steps', type=int, default=gradient_accumulation_steps)
+parser.add_argument('--batch_size', type=int, default=batch_size)
+parser.add_argument('--block_size', type=int, default=block_size)
+parser.add_argument('--n_layer', type=int, default=n_layer)
+parser.add_argument('--n_head', type=int, default=n_head)
+parser.add_argument('--n_embd', type=int, default=n_embd)
+parser.add_argument('--dropout', type=float, default=dropout)
+parser.add_argument('--bias', action='store_true')
+parser.add_argument('--learning_rate', type=float, default=learning_rate)
+parser.add_argument('--max_iters', type=int, default=max_iters)
+parser.add_argument('--lr_decay_iters', type=int, default=lr_decay_iters)
+parser.add_argument('--min_lr', type=float, default=min_lr)
+parser.add_argument('--beta1', type=float, default=beta1)
+parser.add_argument('--beta2', type=float, default=beta2)
+parser.add_argument('--warmup_iters', type=int, default=warmup_iters)
+parser.add_argument('--device', type=str, default=device)
+parser.add_argument('--dtype', type=str, default=dtype)
+parser.add_argument('--compile', action='store_true')
+parser.add_argument('--seed', type=int, default=1337, help='random seed')
+parser.add_argument('--qv_variant', type=str, default='none',
+                    choices=[
+                        # baseline
+                        'none',
+                        # normalization variants
+                        'vnorm',
+                        'qvnorm',
+                        # static gates
+                        'static_gate',
+                        'static_gate_prehead',
+                        # post-attn normalization
+                        'post_rmsnorm_y',
+                        # Q-conditioned gates
+                        'dynamic',
+                        'dynamic_swiglu',
+                        'dynamic_qconditioned_mlp128',
+                        'dynamic_qconditioned_mlp192',
+                        'dynamic_qconditioned_fullwidth',
+                        'dynamic_qconditioned_fullwidth_headspecific',
+                        'dynamic_q_headshared_elementwise',
+                        # X-conditioned gates
+                        'dynamic_xconditioned_g1',
+                        'dynamic_xconditioned_fullwidth_headspecific',
+                        'dynamic_xconditioned_bottleneck',
+                        'dynamic_x_g1_headspecific_elementwise',
+                        'dynamic_x_g1_headspecific_headwise',
+                        'dynamic_x_g1_headshared_elementwise',
+                        # Random / ablation gates
+                        'dynamic_random_gate',
+                        'dynamic_ones_gate',
+                        'dynamic_random_normal',
+                        'dynamic_bernoulli_gate',
+                        # Zero-parameter dot product gates
+                        'dynamic_dot_scalar',
+                        'dynamic_dot_elementwise',
+                        # A-only gate
+                        'dynamic_a_conditioned',
+                        # Q+A dual-signal gates
+                        'dynamic_qa_conditioned',
+                        'dynamic_qa_conditioned_headspecific',
+                        'dynamic_qa_conditioned_mlp128',
+                        'dynamic_qa_headshared_elementwise',
+                        'dynamic_qa_bilinear_diag',
+                        'dynamic_qa_conditioned_normed',
+                        'dynamic_qa_conditioned_normed_yonly',
+                        'dynamic_qa_conditioned_normed_qonly',
+                        'dynamic_qa_conditioned_softq',
+                        'dynamic_qa_conditioned_softqa',
+                        'dynamic_qa_conditioned_softqa_perlayer',
+                        'dynamic_qa_conditioned_softqa_perdim',
+                        'dynamic_qa_conditioned_softqa_perdim_informed',
+                        'dynamic_qa_conditioned_softqa_perdim_free',
+                        'dynamic_qa_conditioned_lowrank16',
+                        # legacy
+                        'dynamic_postconcat_matched',
+                    ],
+                    help='QV experiment variant')
+
+args = parser.parse_args()
+for k in config_keys:
+    if hasattr(args, k):
+        globals()[k] = getattr(args, k)
+config = {k: globals()[k] for k in config_keys}
+
+ddp = int(os.environ.get('RANK', -1)) != -1
+if ddp:
+    init_process_group(backend='nccl')
+    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
+    seed_offset = ddp_rank
+    gradient_accumulation_steps //= ddp_world_size
+else:
+    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(args.seed + seed_offset)
+torch.backends.cuda.matmul.allow_tf32 = True
+torch.backends.cudnn.allow_tf32 = True
+
+device_type = 'cuda' if 'cuda' in device else 'cpu'
+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)
+
+data_dir = os.path.join('data', dataset)
+
+# OPEN MEMMAPS ONCE, REUSE FOREVER
+train_data = np.memmap(os.path.join(data_dir, 'train.bin'), dtype=np.uint16, mode='r')
+val_data = np.memmap(os.path.join(data_dir, 'val.bin'), dtype=np.uint16, mode='r')
+
+def get_batch(split):
+    data = train_data if split == 'train' else val_data
+    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':
+        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
+
+iter_num = 0
+best_val_loss = 1e9
+
+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_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)
+
+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)
+gptconf.qv_variant = args.qv_variant
+model = GPT(gptconf)
+model.to(device)
+
+scaler = torch.cuda.amp.GradScaler(enabled=(dtype == 'float16'))
+optimizer = model.configure_optimizers(1e-1, learning_rate, (beta1, beta2), device_type)
+
+if compile:
+    print("compiling the model... (takes a ~minute)")
+    model = torch.compile(model)
+
+@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
+
+def get_lr(it):
+    if it < warmup_iters:
+        return learning_rate * (it + 1) / (warmup_iters + 1)
+    if it > lr_decay_iters:
+        return min_lr
+    decay_ratio = (it - warmup_iters) / (lr_decay_iters - warmup_iters)
+    coeff = 0.5 * (1.0 + math.cos(math.pi * decay_ratio))
+    return min_lr + coeff * (learning_rate - min_lr)
+
+if wandb_log and master_process:
+    import wandb
+    wandb.init(project=wandb_project, name=wandb_run_name, config=config)
+
+X, Y = get_batch('train')
+t0 = time.time()
+local_iter_num = 0
+raw_model = model.module if hasattr(model, 'module') else model
+running_mfu = -1.0
+
+while True:
+    lr = get_lr(iter_num)
+    for param_group in optimizer.param_groups:
+        param_group['lr'] = lr
+
+    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 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
+
+    for micro_step in range(gradient_accumulation_steps):
+        with ctx:
+            logits, loss = model(X, Y)
+            loss = loss / gradient_accumulation_steps
+        X, Y = get_batch('train')
+        scaler.scale(loss).backward()
+
+    scaler.unscale_(optimizer)
+    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+    scaler.step(optimizer)
+    scaler.update()
+    optimizer.zero_grad(set_to_none=True)
+
+    t1 = time.time()
+    dt = t1 - t0
+    t0 = t1
+    if iter_num % log_interval == 0 and master_process:
+        lossf = loss.item() * gradient_accumulation_steps
+        if local_iter_num >= 5:
+            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
+
+    if iter_num > max_iters:
+        break
+
+if ddp:
+    destroy_process_group()

code/train_wslbridge.py

@@ -0,0 +1,302 @@
+"""
+Training script with qv_variant choices extended for Phase 7 normed Q-only variants.
+"""
+
+import os
+import time
+import math
+import pickle
+import argparse
+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
+
+out_dir = 'out-shakespeare-char'
+eval_interval = 250
+log_interval = 10
+eval_iters = 200
+eval_only = False
+always_save_checkpoint = False
+wandb_log = False
+wandb_project = 'shakespeare-char'
+wandb_run_name = 'mini-gpt'
+dataset = 'shakespeare_char'
+gradient_accumulation_steps = 1
+batch_size = 64
+block_size = 256
+n_layer = 6
+n_head = 6
+n_embd = 384
+dropout = 0.2
+bias = False
+learning_rate = 1e-3
+max_iters = 5000
+lr_decay_iters = 5000
+min_lr = 1e-4
+beta1 = 0.9
+beta2 = 0.99
+warmup_iters = 100
+device = 'cuda'
+dtype = 'bfloat16' if torch.cuda.is_available() and torch.cuda.is_bf16_supported() else 'float16'
+compile = False
+
+config_keys = [k for k, v in globals().items() if not k.startswith('_') and isinstance(v, (int, float, bool, str))]
+parser = argparse.ArgumentParser(description='Train a mini-GPT')
+parser.add_argument('--out_dir', type=str, default=out_dir)
+parser.add_argument('--eval_interval', type=int, default=eval_interval)
+parser.add_argument('--log_interval', type=int, default=log_interval)
+parser.add_argument('--eval_iters', type=int, default=eval_iters)
+parser.add_argument('--eval_only', action='store_true')
+parser.add_argument('--always_save_checkpoint', action='store_true')
+parser.add_argument('--wandb_log', action='store_true')
+parser.add_argument('--wandb_project', type=str, default=wandb_project)
+parser.add_argument('--wandb_run_name', type=str, default=wandb_run_name)
+parser.add_argument('--dataset', type=str, default=dataset)
+parser.add_argument('--gradient_accumulation_steps', type=int, default=gradient_accumulation_steps)
+parser.add_argument('--batch_size', type=int, default=batch_size)
+parser.add_argument('--block_size', type=int, default=block_size)
+parser.add_argument('--n_layer', type=int, default=n_layer)
+parser.add_argument('--n_head', type=int, default=n_head)
+parser.add_argument('--n_embd', type=int, default=n_embd)
+parser.add_argument('--dropout', type=float, default=dropout)
+parser.add_argument('--bias', action='store_true')
+parser.add_argument('--learning_rate', type=float, default=learning_rate)
+parser.add_argument('--max_iters', type=int, default=max_iters)
+parser.add_argument('--lr_decay_iters', type=int, default=lr_decay_iters)
+parser.add_argument('--min_lr', type=float, default=min_lr)
+parser.add_argument('--beta1', type=float, default=beta1)
+parser.add_argument('--beta2', type=float, default=beta2)
+parser.add_argument('--warmup_iters', type=int, default=warmup_iters)
+parser.add_argument('--device', type=str, default=device)
+parser.add_argument('--dtype', type=str, default=dtype)
+parser.add_argument('--compile', action='store_true')
+parser.add_argument('--seed', type=int, default=1337, help='random seed')
+parser.add_argument('--qv_variant', type=str, default='none',
+                    choices=[
+                        # baseline
+                        'none',
+                        # normalization variants
+                        'vnorm',
+                        'qvnorm',
+                        # static gates
+                        'static_gate',
+                        'static_gate_prehead',
+                        # post-attn normalization
+                        'post_rmsnorm_y',
+                        # Q-conditioned gates
+                        'dynamic',
+                        'dynamic_swiglu',
+                        'dynamic_qconditioned_mlp128',
+                        'dynamic_qconditioned_normed',
+                        'dynamic_qconditioned_mlp128_normed',
+                        'dynamic_qconditioned_mlp192',
+                        'dynamic_qconditioned_fullwidth',
+                        'dynamic_qconditioned_fullwidth_headspecific',
+                        'dynamic_q_headshared_elementwise',
+                        # X-conditioned gates
+                        'dynamic_xconditioned_g1',
+                        'dynamic_xconditioned_fullwidth_headspecific',
+                        'dynamic_xconditioned_bottleneck',
+                        'dynamic_x_g1_headspecific_elementwise',
+                        'dynamic_x_g1_headspecific_headwise',
+                        'dynamic_x_g1_headshared_elementwise',
+                        # Random / ablation gates
+                        'dynamic_random_gate',
+                        'dynamic_ones_gate',
+                        'dynamic_random_normal',
+                        'dynamic_bernoulli_gate',
+                        # Zero-parameter dot product gates
+                        'dynamic_dot_scalar',
+                        'dynamic_dot_elementwise',
+                        # A-only gate
+                        'dynamic_a_conditioned',
+                        # Q+A dual-signal gates
+                        'dynamic_qa_conditioned',
+                        'dynamic_qa_conditioned_headspecific',
+                        'dynamic_qa_conditioned_mlp128',
+                        'dynamic_qa_headshared_elementwise',
+                        'dynamic_qa_bilinear_diag',
+                        'dynamic_qa_conditioned_normed',
+                        'dynamic_qa_conditioned_lowrank16',
+                        # legacy
+                        'dynamic_postconcat_matched',
+                    ],
+                    help='QV experiment variant')
+
+args = parser.parse_args()
+for k in config_keys:
+    if hasattr(args, k):
+        globals()[k] = getattr(args, k)
+config = {k: globals()[k] for k in config_keys}
+
+ddp = int(os.environ.get('RANK', -1)) != -1
+if ddp:
+    init_process_group(backend='nccl')
+    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
+    seed_offset = ddp_rank
+    gradient_accumulation_steps //= ddp_world_size
+else:
+    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(args.seed + seed_offset)
+torch.backends.cuda.matmul.allow_tf32 = True
+torch.backends.cudnn.allow_tf32 = True
+
+device_type = 'cuda' if 'cuda' in device else 'cpu'
+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)
+
+data_dir = os.path.join('data', dataset)
+
+# OPEN MEMMAPS ONCE, REUSE FOREVER
+train_data = np.memmap(os.path.join(data_dir, 'train.bin'), dtype=np.uint16, mode='r')
+val_data = np.memmap(os.path.join(data_dir, 'val.bin'), dtype=np.uint16, mode='r')
+
+def get_batch(split):
+    data = train_data if split == 'train' else val_data
+    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':
+        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
+
+iter_num = 0
+best_val_loss = 1e9
+
+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_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)
+
+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)
+gptconf.qv_variant = args.qv_variant
+model = GPT(gptconf)
+model.to(device)
+
+scaler = torch.cuda.amp.GradScaler(enabled=(dtype == 'float16'))
+optimizer = model.configure_optimizers(1e-1, learning_rate, (beta1, beta2), device_type)
+
+if compile:
+    print("compiling the model... (takes a ~minute)")
+    model = torch.compile(model)
+
+@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
+
+def get_lr(it):
+    if it < warmup_iters:
+        return learning_rate * (it + 1) / (warmup_iters + 1)
+    if it > lr_decay_iters:
+        return min_lr
+    decay_ratio = (it - warmup_iters) / (lr_decay_iters - warmup_iters)
+    coeff = 0.5 * (1.0 + math.cos(math.pi * decay_ratio))
+    return min_lr + coeff * (learning_rate - min_lr)
+
+if wandb_log and master_process:
+    import wandb
+    wandb.init(project=wandb_project, name=wandb_run_name, config=config)
+
+X, Y = get_batch('train')
+t0 = time.time()
+local_iter_num = 0
+raw_model = model.module if hasattr(model, 'module') else model
+running_mfu = -1.0
+
+while True:
+    lr = get_lr(iter_num)
+    for param_group in optimizer.param_groups:
+        param_group['lr'] = lr
+
+    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 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
+
+    for micro_step in range(gradient_accumulation_steps):
+        with ctx:
+            logits, loss = model(X, Y)
+            loss = loss / gradient_accumulation_steps
+        X, Y = get_batch('train')
+        scaler.scale(loss).backward()
+
+    scaler.unscale_(optimizer)
+    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+    scaler.step(optimizer)
+    scaler.update()
+    optimizer.zero_grad(set_to_none=True)
+
+    t1 = time.time()
+    dt = t1 - t0
+    t0 = t1
+    if iter_num % log_interval == 0 and master_process:
+        lossf = loss.item() * gradient_accumulation_steps
+        if local_iter_num >= 5:
+            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
+
+    if iter_num > max_iters:
+        break
+
+if ddp:
+    destroy_process_group()