coppola_pretrain_tiny.py

#!/usr/bin/env python3
"""Train a tiny byte-level GPT with full Coppola projection from step 0.

This is the local in-repo pretraining path for the full-Coppola experiment.
It avoids external trainer dependencies and keeps the architecture small
enough that we can iterate on the projection itself.
"""

from __future__ import annotations

import argparse
import json
import math
import random
from dataclasses import asdict, dataclass
from pathlib import Path
from typing import Dict, List, Sequence

import torch
import torch.nn as nn
import torch.nn.functional as F

ROOT = Path(__file__).resolve().parent.parent
import sys

if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

from coppola_pretraining import (  # noqa: E402
    ComponentNorms,
    CoppolaPretrainingConfig,
    CoppolaPretrainingController,
    FamilyScales,
    RankPolicy,
)

try:
    import pyarrow.parquet as pq
except ImportError:
    pq = None


def set_seed(seed: int) -> None:
    random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)


@dataclass
class GPTConfig:
    seq_len: int = 256
    vocab_size: int = 256
    n_layer: int = 8
    n_head: int = 8
    n_embd: int = 256
    mlp_mult: int = 4
    mlp_activation: str = "silu"
    dropout: float = 0.0
    loss_kind: str = "ce"
    # BCE-only regularizers (no-op when loss_kind == "ce")
    logit_rmsnorm_scale: float = 0.0  # > 0 enables F.rms_norm(logits) * scale before BCE; 0 disables. Pure RMSNorm
                                      #   (scale=1) caps BCE at the vocab-sphere floor (~80 for V=256); use scale ≈ 8.
    sparsity_lambda: float = 0.0      # penalty on mean σ(z_j) for non-target j — drives idle vocab toward nop
    # Block-identity attractor: drives specified blocks toward identity by penalizing their residual contribution
    block_attractor_indices: tuple[int, ...] = ()
    block_attractor_lambda: float = 0.0  # penalty λ * Σ_{ℓ in indices} mean((block_ℓ(x) - x)²)
    # Learnable per-block gates γ_ℓ: x_{ℓ+1} = x_ℓ + γ_ℓ * (block_ℓ(x_ℓ) - x_ℓ). Initialized to 1.
    # γ → 0 makes the block identity; L1 prior penalizes capacity utilization adaptively.
    block_gates: bool = False
    block_gate_kind: str = "linear"  # "linear" (γ = α, init 1.0) | "sigmoid" (γ = σ(α), init α=4)
                                      # | "ste" (forward γ = round(σ(α)), backward gradient via σ(α))
    block_gate_l1: float = 0.0
    # Loop-block recurrence: designated blocks are applied K times during forward.
    # `loop_block_indices` is the set of block positions to iterate; `loop_block_k`
    # is how many times each iterates. Empty/K=1 = standard feedforward (no looping).
    # The block weights see K gradient applications per training step, so AdamW
    # automatically rescales effective LR. Tests "thinking in BP space": each
    # iteration is one BP round, recurrent across the loop block.
    loop_block_indices: tuple[int, ...] = ()
    loop_block_k: int = 1
    # Loop UNIT: contiguous range of blocks applied K times together (a deeper
    # loop body than --loop-block-indices, which loops each block independently).
    # Empty = disabled. When set, indices must be a contiguous list. K is also
    # loop_block_k (same parameter). At inference, --loop-block-k-override applies.
    loop_unit_indices: tuple[int, ...] = ()
    # Heterogeneous architecture: per-block n_head override (one value per block).
    # Empty tuple = uniform cfg.n_head across all blocks. When non-empty, length must equal n_layer
    # and each entry must divide n_embd. Lets OR-dominant blocks keep multi-head while focused
    # blocks use single-head.
    n_heads_per_block: tuple[int, ...] = ()
    # Heterogeneous per-block MLP width override (one mlp_mult per block).
    # Empty tuple = uniform cfg.mlp_mult. Length must equal n_layer when non-empty.
    # Motivated by SVD analysis showing block 0 needs full MLP capacity but middle/late
    # blocks use only ~50% of effective rank — heterogeneous sizing saves ~24% of params.
    mlp_mult_per_block: tuple[int, ...] = ()
    # Heterogeneous per-block MLP activation. Empty = uniform cfg.mlp_activation.
    # Per the reverse-compiled adder (docs/coppola-bp-weight-compilation.md):
    # L0/L1 are exactly linear (R²≈0.99), L2 wants sigmoid-style threshold for
    # conjunction, L3+ need full nonlinearity. Allowed values per block:
    # {"linear", "sigmoid", "silu", "gelu"}.
    mlp_activation_per_block: tuple[str, ...] = ()
    # RMSNorm scale applied to MLP output before residual add. 0 = off (default).
    # >0 calibrates the per-MLP delta to a fixed log-odds magnitude, mirroring
    # the logit-RMSNorm that gates the BP structure on the attention/LM-head
    # side. Hypothesis: this makes MLP byte-basis diag-mass emerge at late
    # blocks (signature 5) the way logit-RMSNorm makes signatures 1-4 emerge.
    mlp_output_rmsnorm_scale: float = 0.0
    # L1 sparsity penalty on the post-activation inner MLP tensor.
    # Adds λ · mean(|act|) summed across blocks to the loss.  Pushes the model
    # toward per-token typed-channel firing: only a small fraction of inner
    # neurons should activate per token (the BP-message reading of MLPs).
    # 0 = off (default).  Typical values: 1e-3 to 1e-1.  Direct analogue of
    # `sparsity_lambda` (which acts on vocab-side σ(z_j)) but on the MLP
    # intermediate.
    mlp_act_sparsity_lambda: float = 0.0
    # Hard top-K mask on inner MLP activations per (batch, token) with
    # straight-through estimator.  K > 0 forces structural sparsity: each
    # token activates exactly K of `mlp_mult * n_embd` inner neurons; gradients
    # flow through as if the mask weren't there.  Tests whether MLPs implement
    # typed BP channels when forced to per-token sparse firing.
    # 0 = off (default).  Typical: 8, 16, 32 out of 512 (inter at mlp_mult=2, d=256).
    mlp_act_topk: int = 0
    # Direct pressure on byte-basis row-concentration of MLP linearised composite
    # M = W_lm · W_proj · W_fc · W_wte^T.  For each selected layer, computes the
    # mean top-K row mass fraction of M and subtracts (λ · concentration) from
    # the loss.  Pushes the MLP toward sparse-transition BP-readability —
    # each output byte should be predicted by few input bytes through that block.
    # Set lambda > 0 AND layers non-empty to enable.
    mlp_byte_concentration_lambda: float = 0.0
    mlp_byte_concentration_topk: int = 4
    mlp_byte_concentration_layers: tuple[int, ...] = ()
    # Attention factorization kind. "mha" = standard multi-head softmax attention.
    # "tucker" = R[i,a,j,b] decomposed as sum_{p,q} core[p,q] * attn_p[i,j] * W_OV_q[a,b]:
    #   r_p attention patterns (each with own Q,K) × r_q OV kernels (each with own V), mixed by
    #   the core. Tucker(r_p=n_head, r_q=n_head, core=I) reduces to MHA exactly.
    attn_kind: str = "mha"
    tucker_r_p: int = 4  # number of attention patterns
    tucker_r_q: int = 4  # number of OV kernels
    # Tucker tied-V/tied-O controls (Phase B Run 2): rank-vs-operator-multiplicity discriminant.
    tucker_tied_v: bool = False
    tucker_tied_o: bool = False
    # Tucker per-operator value dim. 0 = auto (d / r_q, backward compat). Setting > 0
    # decouples d_v from r_q so reader count and writer width can scale independently.
    # Required by the 'many readers one wide writer' result: d_v should stay around 32 even
    # as r_q grows so the shared W_O of size (d_v, d) doesn't become the bottleneck.
    tucker_d_v: int = 0
    # Block-diagonal-(a,b) attention: split d_model into G = d_model // block_diag_size groups
    # of size block_diag_size; each group has its own (Q_g, K_g, V_g), with no cross-group
    # mixing in c_attn. The output projection is configurable.
    block_diag_size: int = 4
    # If True (default), output projection is also block-diagonal (no cross-group mixing in
    # c_proj). If False, the output projection is full d×d, restoring cross-group output mixing
    # while keeping the input-side c_attn block-diagonal.
    block_diag_proj: bool = True
    # Shared-mixer-rotor attention: Y = sum_s O_s · A · R_s · M · x
    #   M shared (d, d) full mixer; R_s per-operator cheap specialization; O_s lifting.
    #   The architecture motivated by the friend's Run-3 analysis: full input mixing +
    #   parallel cheap operator diversity, RotorQuant-style rotations applied AFTER global
    #   mixing.
    smr_n_operators: int = 4         # number of operators (parallel feature transforms)
    smr_d_v: int = 32                # per-operator inner dim
    smr_r_kind: str = "rotor"        # "diag" | "rotor" | "full" | "none" — uniform per-op spec
    smr_rotor_block: int = 4         # if "rotor": block size for SO(b) rotor (must divide d)
    # Heterogeneous per-operator R kinds (Run 4): when non-empty, overrides smr_r_kind.
    # Tuple length is the operator count. Each entry ∈ {diag, rotor, full, none}.
    # Examples: ('full', 'diag', 'diag', 'diag') — one rich + cheap tail,
    #           ('full', 'rotor', 'rotor', 'rotor') — one rich + rotor tail.
    smr_operator_kinds: tuple[str, ...] = ()
    # 'Many readers, one writer' compressed form: when True, W_O is shared across all
    # operators (single (d_v, d) projection summed across operator outputs) instead of
    # per-operator. Implements Y = W_O · sum_q A · R_q · M · x — the user's corrected
    # architecture after the Tucker tied-O empirical finding.
    smr_shared_wo: bool = False
    # Attention score normalization (Run 5). Applies to tucker and shared_mixer_rotor.
    # 'softmax' (default): standard masked softmax — BP-foreign (forces convex combination over keys).
    # 'linear_gain': no normalization, just causal-masked scores × learnable scalar γ.
    #               Most BP-canonical (signed log-odds, unbounded).
    # 'rms_signed': scores divided by per-row RMS — keeps sign, removes partition fn,
    #               but bounds magnitudes per row.
    # 'kernel_relu2': ReLU(scores)² then row-normalize to 1.
    # 'kernel_elu_plus1': elu(scores)+1 then row-normalize to 1.
    # 'sparsemax': α=2 entmax — projection onto simplex; produces exact-zero attention weights
    #              for low-scoring keys. BP-aligned (sparse attention = "attend to specific keys").
    # 'entmax15': α=1.5 entmax — sparse like sparsemax but smoother gradient.
    # 'topk_ste': forward = hard top-k of softmax (clean discrete attention, exact zeros);
    #             backward = full softmax gradient via straight-through estimator.
    #             Most BP-aligned at inference (sparse, named keys); soft training.
    attn_norm: str = "softmax"
    topk_k: int = 4  # used by attn_norm='topk_ste': how many keys each query attends to
    # Position encoding. "learned" = absolute learned wpe (default, BP-incompatible);
    # "alibi" = no wpe; per-head linear-bias (Press et al. 2021), aligning with the BP
    # framing where positional information enters as an additive log-odds offset rather
    # than as a learned token-position vector subject to softmax smearing.
    # "rope" = rotary position encoding (Su et al. 2021).
    position_encoding: str = "learned"
    # RoPE rotation base; default 10000 covers long contexts. For short byte context
    # (seq_len ~256), lower values (e.g. 1000) make the slow-frequency dims actually
    # rotate noticeably across the window.
    rope_base: float = 10000.0
    # ALiBi slope scale: multiplies the standard slopes. Gentler slopes (e.g. 0.5) may
    # help when only a few attention "channels" are present (Tucker r_p=2).
    alibi_slope_scale: float = 1.0
    # Attention Residuals (Kimi, arXiv:2603.15031). Replace the fixed unit-weight
    # residual h_l = h_{l-1} + f(h_{l-1}) with a learned, input-independent
    # softmax/entmax aggregation over ALL preceding block outputs:
    #   h_l = Σ_{i<l} α_{i→l} v_i,  α = Φ(w_lᵀ RMSNorm(v_i)),  v_0 = embedding.
    # 'none' (default): standard additive residual — existing path unchanged.
    # 'full': block-granular Full AttnRes (depth-unit = one transformer block;
    #         M is n_layer×n_layer, directly comparable to the BP-install graph).
    # See docs/attnres-bp-research-program.md.
    res_attn: str = "none"
    # Normalization Φ for the depth-attention weights (the BP-structure axis):
    # 'softmax' = paper default (dense, full support); 'entmax15'/'sparsemax' =
    # exact-sparse — the learned depth-mixing matrix becomes a sparse BP graph.
    res_attn_norm: str = "softmax"
    # A2′ depth-router "pressure to specialise" knobs (see
    # docs/attnres-bp-research-program.md First-results redesign).
    # res_attn_temp: fixed scalar τ multiplying the depth scores w_lᵀk
    # before Φ — the depth analogue of the logit-RMSNorm(8) sequence
    # knob. Without score range, entmax15 ≡ softmax ≡ uniform (the
    # observed L=4 collapse). res_attn_entropy_lambda: coefficient on
    # +mean H(α_depth) added to the loss (minimised ⇒ peaked/sparse
    # depth attention; gives the otherwise-idle router a reason to move).
    res_attn_temp: float = 1.0
    res_attn_entropy_lambda: float = 0.0


class CausalSelfAttention(nn.Module):
    def __init__(self, cfg: GPTConfig, n_head: int | None = None):
        super().__init__()
        n_head_eff = cfg.n_head if n_head is None else n_head
        if cfg.n_embd % n_head_eff != 0:
            raise ValueError(f"n_embd ({cfg.n_embd}) must be divisible by n_head ({n_head_eff})")
        self.n_head = n_head_eff
        self.head_dim = cfg.n_embd // n_head_eff
        self.c_attn = nn.Linear(cfg.n_embd, 3 * cfg.n_embd, bias=False)
        self.c_proj = nn.Linear(cfg.n_embd, cfg.n_embd, bias=False)
        self.dropout = nn.Dropout(cfg.dropout)
        # Attention norm: 'softmax' uses fast scaled_dot_product_attention;
        # other modes (entmax15, sparsemax, topk_ste, ...) use the manual
        # path via _apply_attn_norm.  Matches the Tucker attention norm
        # support so 'BP-canonical sparse attention' is usable on MHA too.
        self.attn_norm = cfg.attn_norm
        self.attn_gain = nn.Parameter(torch.tensor(0.1)) if cfg.attn_norm == "linear_gain" else None
        self.topk_k = cfg.topk_k

    def forward(self, x: torch.Tensor,
                attn_bias: torch.Tensor | None = None,
                rope_freqs: tuple[torch.Tensor, torch.Tensor] | None = None) -> torch.Tensor:
        bsz, seqlen, hidden = x.shape
        qkv = self.c_attn(x)
        q, k, v = qkv.chunk(3, dim=-1)
        q = q.view(bsz, seqlen, self.n_head, self.head_dim).transpose(1, 2)
        k = k.view(bsz, seqlen, self.n_head, self.head_dim).transpose(1, 2)
        v = v.view(bsz, seqlen, self.n_head, self.head_dim).transpose(1, 2)
        if rope_freqs is not None:
            cos, sin = rope_freqs
            q = _apply_rope(q, cos, sin)
            k = _apply_rope(k, cos, sin)
        if self.attn_norm == "softmax":
            # Fast path: scaled_dot_product_attention (built-in softmax + mask)
            if attn_bias is None:
                y = F.scaled_dot_product_attention(q, k, v, dropout_p=0.0, is_causal=True)
            else:
                T = seqlen
                causal_mask = torch.zeros(T, T, device=x.device, dtype=q.dtype)
                causal_mask.masked_fill_(
                    torch.triu(torch.ones(T, T, device=x.device, dtype=torch.bool), diagonal=1),
                    float("-inf"),
                )
                mask = attn_bias.to(q.dtype) + causal_mask
                y = F.scaled_dot_product_attention(q, k, v, attn_mask=mask, dropout_p=0.0, is_causal=False)
        else:
            # Manual path: compute scores, apply chosen norm (entmax15, etc.)
            scores = torch.einsum("bhid,bhjd->bhij", q, k) / (self.head_dim ** 0.5)
            attn = _apply_attn_norm(scores, self.attn_norm, self.attn_gain,
                                     topk_k=self.topk_k, attn_bias=attn_bias)
            y = torch.einsum("bhij,bhjd->bhid", attn, v)
        y = y.transpose(1, 2).contiguous().view(bsz, seqlen, hidden)
        return self.dropout(self.c_proj(y))


class BlockDiagAttention(nn.Module):
    """Block-diagonal-(a,b) attention with configurable output projection.

    Partition d_model into G = d_model / block_size disjoint coordinate groups.
    Each group g has its own input projections Q_g, K_g, V_g of shape (b, b)
    that operate ONLY on its own block of x — no cross-group mixing in the
    input projection (block-diagonal c_attn).

    The output projection is configurable:
      block_diag_proj=True  → block-diagonal c_proj (per-group W_O_g, no
                              cross-group mixing). Lossy on cross-feature
                              output integration.
      block_diag_proj=False → full d×d c_proj. Each group's per-position
                              attention output gets cross-mixed across all
                              groups. Tests whether the cost of pure block-
                              diagonal is on the input or output side.

    Parameter count per layer:
      block_diag_proj=True:  G·4·b² = 4·d·b           (32× compression at d=128, b=4)
      block_diag_proj=False: G·3·b² + d·d = 3·d·b + d² (≈ MHA on c_proj, much smaller c_attn)
    """

    def __init__(self, cfg: GPTConfig, block_size: int, block_diag_proj: bool = True):
        super().__init__()
        d = cfg.n_embd
        if d % block_size != 0:
            raise ValueError(f"n_embd ({d}) must be divisible by block_diag_size ({block_size})")
        self.G = d // block_size
        self.b = block_size
        self.block_diag_proj = block_diag_proj
        scale_in = self.b ** -0.5
        self.q_w = nn.Parameter(torch.randn(self.G, self.b, self.b) * scale_in)
        self.k_w = nn.Parameter(torch.randn(self.G, self.b, self.b) * scale_in)
        self.v_w = nn.Parameter(torch.randn(self.G, self.b, self.b) * scale_in)
        if block_diag_proj:
            # Per-group block-diagonal output projection
            self.o_w = nn.Parameter(torch.randn(self.G, self.b, self.b) * scale_in)
            self._c_proj_full = None
        else:
            # Full d×d output projection — restores cross-group mixing
            self.o_w = None
            self._c_proj_full = nn.Linear(d, d, bias=False)
        self.dropout = nn.Dropout(cfg.dropout)

    @property
    def c_proj(self):
        """Adapter for iter_projected_params and RankPolicy.
        Returns the actual nn.Linear when full c_proj is enabled, or a
        synthetic dense view of the block-diagonal o_w otherwise."""
        if self.block_diag_proj:
            return _BlockDiagCProjView(self)
        return self._c_proj_full

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B, T, d = x.shape
        xg = x.view(B, T, self.G, self.b)
        Q = torch.einsum("btgi,gij->btgj", xg, self.q_w)
        K = torch.einsum("btgi,gij->btgj", xg, self.k_w)
        V = torch.einsum("btgi,gij->btgj", xg, self.v_w)
        Q = Q.transpose(1, 2)  # [B, G, T, b]
        K = K.transpose(1, 2)
        V = V.transpose(1, 2)
        out = F.scaled_dot_product_attention(Q, K, V, dropout_p=0.0, is_causal=True)
        out = out.transpose(1, 2)  # [B, T, G, b]
        if self.block_diag_proj:
            out = torch.einsum("btgi,gij->btgj", out, self.o_w)
            out = out.reshape(B, T, d)
        else:
            # Concat group outputs to flat [B, T, d], then full c_proj
            out = out.reshape(B, T, d)
            out = self._c_proj_full(out)
        return self.dropout(out)


class _BlockDiagCProjView:
    """Read-only weight-shape adapter so existing iter_projected_params /
    RankPolicy / projection-related code can introspect a BlockDiagAttention
    block without crashing. The underlying training parameter is `o_w`."""

    def __init__(self, attn: "BlockDiagAttention"):
        self._attn = attn

    @property
    def weight(self) -> torch.Tensor:
        # Materialise a (d, d) block-diagonal view of o_w only when asked
        d = self._attn.G * self._attn.b
        out = self._attn.o_w.new_zeros(d, d)
        for g in range(self._attn.G):
            i = g * self._attn.b
            out[i:i + self._attn.b, i:i + self._attn.b] = self._attn.o_w[g]
        return out


def _make_rope_freqs(head_dim: int, T: int, base: float, device, dtype) -> tuple[torch.Tensor, torch.Tensor]:
    """Standard RoPE cos/sin tables of shape [T, head_dim/2]."""
    if head_dim % 2 != 0:
        raise ValueError(f"RoPE requires even head_dim, got {head_dim}")
    inv_freq = 1.0 / (base ** (torch.arange(0, head_dim, 2, device=device, dtype=torch.float32) / head_dim))
    pos = torch.arange(T, device=device, dtype=torch.float32)
    freqs = torch.outer(pos, inv_freq)  # [T, head_dim/2]
    return freqs.cos().to(dtype), freqs.sin().to(dtype)


def _apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
    """Rotate every pair of dims in x by position-dependent angle.

    x: [..., T, head_dim]; cos, sin: [T, head_dim/2].
    """
    x1 = x[..., 0::2]
    x2 = x[..., 1::2]
    # cos/sin broadcast over leading dims of x.
    rot1 = x1 * cos - x2 * sin
    rot2 = x1 * sin + x2 * cos
    out = torch.empty_like(x)
    out[..., 0::2] = rot1
    out[..., 1::2] = rot2
    return out


def _alibi_slopes(n_head: int) -> torch.Tensor:
    """Standard ALiBi per-head slopes (Press et al. 2021, eq. 2).

    For n_head a power of two: m_h = 2^(-8h/n_head) for h = 1..n_head.
    Otherwise interpolate following the reference implementation.
    """
    def _power_of_two(n: int) -> list[float]:
        start = 2.0 ** (-(2.0 ** -(math.log2(n) - 3)))
        return [start * (start ** i) for i in range(n)]

    if (n_head & (n_head - 1)) == 0:
        slopes = _power_of_two(n_head)
    else:
        closest = 1 << (n_head - 1).bit_length() >> 1
        a = _power_of_two(closest)
        b = _power_of_two(2 * closest)[0::2][: n_head - closest]
        slopes = a + b
    return torch.tensor(slopes, dtype=torch.float32)


def _alibi_bias(n_head: int, T: int, slopes: torch.Tensor, device, dtype) -> torch.Tensor:
    """Return [n_head, T, T] additive bias: m_h * (key_pos - query_pos), 0 for j > i.

    The causal mask blocks j > i separately, so this fills only the lower-triangular
    band with the linearly-decaying offset. Result is added to attention scores.
    """
    pos = torch.arange(T, device=device)
    rel = pos.unsqueeze(0) - pos.unsqueeze(1)  # [T, T]; rel[i, j] = j - i
    rel = rel.to(dtype)
    # ALiBi penalty grows in magnitude with |key - query| distance; use -|rel| so
    # closer keys have score 0 and farther keys are penalized.
    return slopes.to(device=device, dtype=dtype).view(n_head, 1, 1) * (-rel.abs()).unsqueeze(0)


def _sparsemax(z: torch.Tensor, dim: int = -1) -> torch.Tensor:
    """Exact sparsemax (Martins & Astudillo 2016) — Euclidean projection onto
    the simplex; produces exact zeros. Closed form, autograd-differentiable."""
    z_sorted, _ = torch.sort(z, descending=True, dim=dim)
    rng = torch.arange(1, z.size(dim) + 1, device=z.device, dtype=z.dtype)
    shape = [1] * z.dim()
    shape[dim] = -1
    rng = rng.view(shape)
    cssv = z_sorted.cumsum(dim) - 1.0
    support = rng * z_sorted > cssv
    k = support.sum(dim=dim, keepdim=True).clamp_min(1)
    tau = cssv.gather(dim, k - 1) / k.to(z.dtype)
    return torch.clamp(z - tau, min=0.0)


def _entmax15(z: torch.Tensor, dim: int = -1) -> torch.Tensor:
    """Exact 1.5-entmax (Peters, Niculae & Martins 2019). Fallback only —
    prefer the `entmax` package (see _depth_attn_norm). Autograd through
    `sqrt(delta)` is singular as delta→0; the +1e-8 below bounds the
    backward (≈1/2e-4) so this stays finite, but it is not the exact
    entmax Jacobian. Forward is exact."""
    z = z - z.max(dim=dim, keepdim=True).values
    z = z / 2.0
    z_sorted, _ = torch.sort(z, descending=True, dim=dim)
    rng = torch.arange(1, z.size(dim) + 1, device=z.device, dtype=z.dtype)
    shape = [1] * z.dim()
    shape[dim] = -1
    rng = rng.view(shape)
    mean = z_sorted.cumsum(dim) / rng
    mean_sq = (z_sorted ** 2).cumsum(dim) / rng
    ss = rng * (mean_sq - mean ** 2)
    delta = (1.0 - ss) / rng
    tau = mean - torch.sqrt(torch.clamp(delta, min=0.0) + 1e-8)
    support_size = (tau <= z_sorted).sum(dim=dim, keepdim=True).clamp_min(1)
    tau_star = tau.gather(dim, support_size - 1)
    return torch.clamp(z - tau_star, min=0.0) ** 2


try:  # noqa: SIM105 — prefer the package's numerically-correct custom backward
    from entmax import entmax15 as _pkg_entmax15, sparsemax as _pkg_sparsemax
    _HAVE_ENTMAX_PKG = True
except ImportError:  # pragma: no cover
    _HAVE_ENTMAX_PKG = False


def _depth_attn_norm(scores: torch.Tensor, mode: str) -> torch.Tensor:
    """Normalize depth-attention scores over the source axis (last dim).
    Used by Attention Residuals (arXiv:2603.15031) to weight preceding
    layer outputs. softmax = the paper's default (dense, full support);
    entmax15/sparsemax = exact-sparse variants (this repo's BP-structure
    extension — a sparse depth-mixing matrix is the BP factor graph).

    entmax15/sparsemax route to the `entmax` package when installed: its
    custom backward is exact, whereas autograd through the inline closed
    form differentiates `sqrt(clamp(delta,0))`, whose gradient is
    unbounded at the sparse-support boundary (delta→0) — a latent NaN
    that A2′ score-pressure drives straight into. The inline path is a
    forward-correct, eps-stabilised fallback for when the pkg is absent."""
    if mode == "softmax":
        return torch.softmax(scores, dim=-1)
    if mode == "entmax15":
        return (_pkg_entmax15(scores, dim=-1) if _HAVE_ENTMAX_PKG
                else _entmax15(scores, dim=-1))
    if mode == "sparsemax":
        return (_pkg_sparsemax(scores, dim=-1) if _HAVE_ENTMAX_PKG
                else _sparsemax(scores, dim=-1))
    raise ValueError(f"unknown res_attn_norm: {mode!r}")


def _apply_attn_norm(scores: torch.Tensor, mode: str,
                     gain: torch.Tensor | None,
                     topk_k: int = 4,
                     attn_bias: torch.Tensor | None = None) -> torch.Tensor:
    """Apply causal mask + chosen normalization to raw attention scores.

    scores: [..., T, T] where the last two dims are (query_pos, key_pos).
    gain:   optional learnable scalar Parameter for 'linear_gain' mode.
    topk_k: keys per query for 'topk_ste' mode.
    attn_bias: optional additive bias broadcast onto scores before mask/norm
               (ALiBi). Shape compatible with `scores`.
    """
    T = scores.size(-1)
    if attn_bias is not None:
        scores = scores + attn_bias
    causal = torch.triu(torch.ones(T, T, device=scores.device), diagonal=1).bool()
    if mode == "softmax":
        scores = scores.masked_fill(causal, float("-inf"))
        return F.softmax(scores, dim=-1)
    if mode == "linear_gain":
        attn = scores.masked_fill(causal, 0.0)
        return attn * (gain if gain is not None else 1.0)
    if mode == "rms_signed":
        # Row-RMS, keeping sign — BP-aligned (no partition function).
        attn = scores.masked_fill(causal, 0.0)
        rms = (attn * attn).mean(dim=-1, keepdim=True).clamp(min=1e-8).sqrt()
        return attn / rms
    if mode == "kernel_relu2":
        attn = F.relu(scores) ** 2
        attn = attn.masked_fill(causal, 0.0)
        denom = attn.sum(dim=-1, keepdim=True).clamp(min=1e-8)
        return attn / denom
    if mode == "kernel_elu_plus1":
        attn = F.elu(scores) + 1.0  # always positive
        attn = attn.masked_fill(causal, 0.0)
        denom = attn.sum(dim=-1, keepdim=True).clamp(min=1e-8)
        return attn / denom
    if mode == "sparsemax":
        # α=2 entmax — projection onto the simplex; produces exact-zero attention
        # weights for low-scoring keys. Causal mask via -inf so masked positions
        # project to 0.
        from entmax import sparsemax  # noqa: F401  (lazy import; pip-installed)
        scores_masked = scores.masked_fill(causal, float("-inf"))
        return sparsemax(scores_masked, dim=-1)
    if mode == "entmax15":
        # α=1.5 entmax — sparse like sparsemax but smoother gradient.
        # Package when present (exact custom backward); else the local
        # forward-exact _entmax15 (the eval sandbox has no `entmax` package).
        # Verified: forward matches the package to <2e-7, 0 argmax flips.
        scores_masked = scores.masked_fill(causal, float("-inf"))
        return (_pkg_entmax15(scores_masked, dim=-1) if _HAVE_ENTMAX_PKG
                else _entmax15(scores_masked, dim=-1))
    if mode == "topk_ste":
        # Forward: hard top-k softmax (only top-k entries non-zero, then renormalized).
        # Backward: gradient flows as full softmax (straight-through estimator).
        scores_masked = scores.masked_fill(causal, float("-inf"))
        soft = F.softmax(scores_masked, dim=-1)
        k = min(topk_k, T)
        # Use raw masked scores for selection (avoids ranking over softmax that may have ties).
        _, topk_idx = scores_masked.topk(k, dim=-1)
        topk_mask = torch.zeros_like(scores).scatter_(-1, topk_idx, 1.0)
        hard = soft * topk_mask
        hard = hard / hard.sum(dim=-1, keepdim=True).clamp(min=1e-9)
        # STE: forward = hard, backward via soft.
        return hard.detach() + soft - soft.detach()
    raise ValueError(f"unknown attn_norm: {mode!r}")


class SharedMixerRotorAttention(nn.Module):
    """Shared global mixer + cheap per-operator specialization + parallel output operators.

    Architecture (per the friend's Run-3 analysis):

        U     = M · x                        [T, d]    shared mixer
        Q     = U · W_Q                      [T, d_v]  shared attention pattern
        K     = U · W_K                      [T, d_v]
        A     = softmax(Q K^T / √d_v)        [T, T]    shared routing
        For each operator s = 1..m:
            V_s = R_s · U                    [T, d]    cheap per-op specialization
            Vp_s = V_s · W_V_s               [T, d_v]  down-project per op
            G_s  = A · Vp_s                  [T, d_v]  gathered
            Y_s  = G_s · W_O_s               [T, d]    lift back per op
        Y     = sum_s Y_s                    [T, d]    sum of operator contributions

    R_s can be:
      'diag':  diagonal d-vector per operator (cheapest, gain only)
      'rotor': block-diagonal SO(b) per operator (local subspace rotation)
      'full':  dense (d, d) per operator (control / upper bound)
      'none':  identity (no per-op specialization)

    The output Y is the SUM of operator contributions — providing implicit
    cross-operator mixing without an explicit c_proj projection. This is the
    key idea of the friend's architecture: avoid duplicating a full c_proj
    per operator; instead, let lifted contributions add in the residual.

    Parameter count at d=128, m=4, d_v=32, R='rotor' b=4:
        M: d² = 16384
        Q, K shared: 2 × d × d_v = 8192
        m × R rotor: m × (d/b) × b² = m × d × b = 4 × 128 × 4 = 2048
        m × W_V_s: m × d × d_v = 16384
        m × W_O_s: m × d_v × d = 16384
        Total: 59392 vs MHA's 65536 → ~10% attention compression
    """

    def __init__(self, cfg: GPTConfig, n_operators: int, d_v: int,
                 r_kind: str, rotor_block: int,
                 operator_kinds: tuple[str, ...] = (),
                 attn_norm: str = "softmax",
                 shared_wo: bool = False,
                 topk_k: int = 4):
        super().__init__()
        d = cfg.n_embd
        if d % d_v != 0:
            raise ValueError(f"n_embd ({d}) must be divisible by smr_d_v ({d_v})")
        # Resolve per-operator kinds. If operator_kinds is given, it overrides
        # n_operators and uniform r_kind.
        if operator_kinds:
            kinds_list = list(operator_kinds)
            self.m_ops = len(kinds_list)
        else:
            if r_kind not in ("diag", "rotor", "full", "none"):
                raise ValueError(f"unknown smr_r_kind: {r_kind!r}")
            kinds_list = [r_kind] * n_operators
            self.m_ops = n_operators
        valid = {"diag", "rotor", "full", "none"}
        for k in kinds_list:
            if k not in valid:
                raise ValueError(f"unknown operator kind: {k!r} (valid: {valid})")
        if "rotor" in kinds_list and d % rotor_block != 0:
            raise ValueError(f"n_embd ({d}) must be divisible by smr_rotor_block ({rotor_block})")
        self.d_v = d_v
        self.rotor_block = rotor_block
        self.G_blocks = d // rotor_block
        self.operator_kinds: tuple[str, ...] = tuple(kinds_list)
        # Shared full mixer M
        self.M = nn.Linear(d, d, bias=False)
        # Shared attention pattern: Q, K project U → d_v (low-rank QK)
        self.q_proj = nn.Linear(d, d_v, bias=False)
        self.k_proj = nn.Linear(d, d_v, bias=False)
        # Per-operator R_s as ParameterDict-like — one entry per operator with appropriate
        # shape. We use nn.ParameterList so PyTorch tracks them, even when shapes differ.
        self.R_diag = nn.ParameterList()    # entries for 'diag' kind, indexed by op idx (others omitted)
        self.R_rotor = nn.ParameterList()   # entries for 'rotor' kind
        self.R_full = nn.ParameterList()    # entries for 'full' kind
        # Maps from operator index to (storage_kind, storage_idx).
        self._r_lookup: list[tuple[str, int]] = []
        for k in kinds_list:
            if k == "diag":
                self.R_diag.append(nn.Parameter(torch.ones(d)))
                self._r_lookup.append(("diag", len(self.R_diag) - 1))
            elif k == "rotor":
                self.R_rotor.append(nn.Parameter(
                    torch.eye(rotor_block).repeat(self.G_blocks, 1, 1)
                    + torch.randn(self.G_blocks, rotor_block, rotor_block) * 0.02
                ))
                self._r_lookup.append(("rotor", len(self.R_rotor) - 1))
            elif k == "full":
                self.R_full.append(nn.Parameter(
                    torch.eye(d) + torch.randn(d, d) * 0.02
                ))
                self._r_lookup.append(("full", len(self.R_full) - 1))
            elif k == "none":
                self._r_lookup.append(("none", -1))
        # Per-operator down-project W_V_s ∈ R^{d × d_v}
        self.W_V = nn.Parameter(torch.randn(self.m_ops, d, d_v) * (d ** -0.5))
        # Output projection W_O. shared_wo=True means a single (d_v, d) shared by all
        # operators — implements Y = W_O · sum_q (A · V_q) per the 'many readers, one
        # writer' architectural rule from the Tucker tied-O empirical result.
        self.shared_wo = shared_wo
        if shared_wo:
            self.W_O = nn.Parameter(torch.randn(d_v, d) * (d_v ** -0.5))
        else:
            self.W_O = nn.Parameter(torch.randn(self.m_ops, d_v, d) * (d_v ** -0.5))
        # Attention normalization mode + optional gain for 'linear_gain' mode
        self.attn_norm = attn_norm
        self.attn_gain = nn.Parameter(torch.tensor(0.1)) if attn_norm == "linear_gain" else None
        self.topk_k = topk_k
        self.dropout = nn.Dropout(cfg.dropout)

    @property
    def c_proj(self) -> "_SmrCProjView":
        return _SmrCProjView(self)

    def _apply_R(self, U: torch.Tensor) -> torch.Tensor:
        """U: [B, T, d] → V_s: [B, m, T, d] after per-operator R_s applied.

        Iterates over operators (m_ops is small, typically 4–8) so each op
        can use a different specialization kind. If all operators have the
        same kind there's a small overhead vs the bulk-tensor path; for
        heterogeneous configs this is the cleanest implementation.
        """
        B, T, d = U.shape
        outs: list[torch.Tensor] = []
        Ug_cached = None  # only build when needed
        for kind, idx in self._r_lookup:
            if kind == "none":
                outs.append(U)
            elif kind == "diag":
                outs.append(U * self.R_diag[idx][None, None, :])
            elif kind == "rotor":
                if Ug_cached is None:
                    Ug_cached = U.view(B, T, self.G_blocks, self.rotor_block)
                rotated = torch.einsum("btgi,gij->btgj", Ug_cached, self.R_rotor[idx])
                outs.append(rotated.reshape(B, T, d))
            elif kind == "full":
                outs.append(torch.einsum("btj,jk->btk", U, self.R_full[idx]))
        return torch.stack(outs, dim=1)  # [B, m, T, d]

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B, T, d = x.shape
        U = self.M(x)                                         # [B, T, d]
        Q = self.q_proj(U)                                    # [B, T, d_v]
        K = self.k_proj(U)                                    # [B, T, d_v]
        scores = torch.einsum("bid,bjd->bij", Q, K) / (self.d_v ** 0.5)
        attn = _apply_attn_norm(scores, self.attn_norm, self.attn_gain, topk_k=self.topk_k)  # [B, T, T]
        V_per_op = self._apply_R(U)                            # [B, m, T, d]
        # Down-project per op: V_s · W_V_s → [B, m, T, d_v]
        Vp = torch.einsum("bmtd,mdk->bmtk", V_per_op, self.W_V)
        # Apply shared attention: A · Vp_s → [B, m, T, d_v]
        gathered = torch.einsum("bij,bmjk->bmik", attn, Vp)
        if self.shared_wo:
            # Sum operator outputs in d_v space, then a single shared (d_v → d) lift.
            summed = gathered.sum(dim=1)                  # [B, T, d_v]
            out = torch.einsum("btk,kd->btd", summed, self.W_O)  # [B, T, d]
        else:
            # Per-op lift, then sum.
            out = torch.einsum("bmtk,mkd->btd", gathered, self.W_O)
        return self.dropout(out)


class _SmrCProjView:
    """Read-only weight-shape adapter so iter_projected_params + RankPolicy can walk this
    attention type. Returns a (d, d) dense projection equivalent to summing per-operator
    W_O contributions (useful for shape introspection, not training)."""

    def __init__(self, attn: "SharedMixerRotorAttention"):
        self._attn = attn

    @property
    def weight(self) -> torch.Tensor:
        # Effective output projection — materialize a (d, d) view for tooling compatibility.
        if self._attn.shared_wo:
            d = self._attn.W_O.size(1)
            return self._attn.W_O.T @ self._attn.W_O  # (d, d_v) @ (d_v, d)
        d = self._attn.W_O.size(2)
        out = self._attn.W_O.new_zeros(d, d)
        for s in range(self._attn.m_ops):
            out = out + self._attn.W_O[s].T @ self._attn.W_O[s]
        return out


class TuckerAttention(nn.Module):
    """Tucker decomposition of R^ℓ[i,a,j,b]:

        R[i,a,j,b] = sum_{p,q} core[p,q] · attn_p[i,j] · W_OV,q[a,b]

    r_p attention patterns (each with its own (Q_p, K_p) → softmax(Q_p K_p^T) attention map)
    r_q OV kernels (each with its own V_q and contribution to output)
    core ∈ R^{r_p × r_q} mixes which attention pattern feeds which kernel.

    When r_p = r_q = n_head and core = I, this reduces exactly to standard multi-head
    softmax attention. With r_p ≠ r_q, the count of attention patterns and OV kernels
    decouple.

    `tied_v` / `tied_o` controls (Phase B Run 2): test whether observed wins from
    multiple OV kernels are due to operator multiplicity or to total rank/parameter
    budget by sharing parameters across operators.

      tied_v: single V projection (d → d_v) shared across all r_q operators.
              Different attention patterns applied to the *same* V, then mixed.
              Note: at r_p == 1 this is degenerate (all per-operator outputs become
              scalar multiples of one gather-and-V vector). Meaningful only at r_p ≥ 2.
      tied_o: per-operator d_v contributions all use the same (d_v → d) output
              projection, then summed. Tests whether output-side multiplicity
              matters.

    Either flag reduces parameter count without changing the per-operator structure
    visible to the model — the rank-vs-diversity discriminant the friend's Run 2
    analysis calls for.
    """

    def __init__(self, cfg: GPTConfig, r_p: int, r_q: int,
                 tied_v: bool = False, tied_o: bool = False,
                 attn_norm: str = "softmax", d_v: int = 0,
                 topk_k: int = 4):
        super().__init__()
        d = cfg.n_embd
        if d % r_p != 0:
            raise ValueError(f"n_embd ({d}) must be divisible by tucker_r_p ({r_p})")
        # d_v: explicit value if > 0, else auto = d // r_q (backward compat)
        if d_v <= 0:
            if d % r_q != 0:
                raise ValueError(
                    f"n_embd ({d}) must be divisible by tucker_r_q ({r_q}) when d_v is auto"
                )
            d_v_eff = d // r_q
        else:
            d_v_eff = d_v
        self.r_p = r_p
        self.r_q = r_q
        self.d_qk = d // r_p
        self.d_v = d_v_eff
        self.tied_v = tied_v
        self.tied_o = tied_o
        if tied_v and r_p == 1:
            # Warn but don't fail — runner can choose to test the degenerate case.
            import warnings
            warnings.warn(
                "TuckerAttention(r_p=1, tied_v=True) is degenerate: per-operator "
                "outputs are scalar multiples of the same gather-and-V vector. Use "
                "r_p >= 2 for a meaningful tied_v test.",
                stacklevel=2,
            )
        # Q, K project x → r_p · d_qk = d (reshaped to r_p attention patterns)
        self.q_proj = nn.Linear(d, r_p * self.d_qk, bias=False)
        self.k_proj = nn.Linear(d, r_p * self.d_qk, bias=False)
        # V: tied means single (d → d_v); untied means (d → r_q · d_v)
        v_out = self.d_v if tied_v else r_q * self.d_v
        self.v_proj = nn.Linear(d, v_out, bias=False)
        # Output projection:
        #   tied_o: single (d_v, d) applied per-operator and summed
        #   untied: standard (r_q · d_v = d, d) cross-operator mixing matrix
        if tied_o:
            self.c_proj = nn.Linear(self.d_v, d, bias=False)
        else:
            self.c_proj = nn.Linear(d, d, bias=False)
        # Core: initialize as identity when r_p == r_q (so default ≡ MHA), random otherwise
        if r_p == r_q:
            self.core = nn.Parameter(torch.eye(r_p))
        else:
            self.core = nn.Parameter(torch.randn(r_p, r_q) / (r_p * r_q) ** 0.5)
        # Attention normalization mode + optional gain for 'linear_gain'
        self.attn_norm = attn_norm
        self.attn_gain = nn.Parameter(torch.tensor(0.1)) if attn_norm == "linear_gain" else None
        self.topk_k = topk_k
        self.dropout = nn.Dropout(cfg.dropout)

    def forward(self, x: torch.Tensor,
                attn_bias: torch.Tensor | None = None,
                rope_freqs: tuple[torch.Tensor, torch.Tensor] | None = None) -> torch.Tensor:
        B, T, _ = x.shape
        Q = self.q_proj(x).view(B, T, self.r_p, self.d_qk).transpose(1, 2)  # [B, r_p, T, d_qk]
        K = self.k_proj(x).view(B, T, self.r_p, self.d_qk).transpose(1, 2)  # [B, r_p, T, d_qk]
        if rope_freqs is not None:
            cos, sin = rope_freqs
            Q = _apply_rope(Q, cos, sin)
            K = _apply_rope(K, cos, sin)
        if self.tied_v:
            V_shared = self.v_proj(x)  # [B, T, d_v]
            V = V_shared.unsqueeze(1).expand(B, self.r_q, T, self.d_v)  # [B, r_q, T, d_v]
        else:
            V = self.v_proj(x).view(B, T, self.r_q, self.d_v).transpose(1, 2)  # [B, r_q, T, d_v]
        # r_p attention maps with chosen normalization. attn_bias (if any) is [r_p, T, T];
        # _apply_attn_norm broadcasts it onto [B, r_p, T, T] scores.
        scores = torch.einsum("bpid,bpjd->bpij", Q, K) / (self.d_qk ** 0.5)
        attn = _apply_attn_norm(scores, self.attn_norm, self.attn_gain,
                                topk_k=self.topk_k, attn_bias=attn_bias)  # [B, r_p, T, T]
        # Core mixing: linearly combine attention maps to produce r_q kernel-specific maps
        attn_mixed = torch.einsum("pq,bpij->bqij", self.core, attn)  # [B, r_q, T, T]
        # Apply mixed attention to V → [B, r_q, T, d_v]
        out = torch.einsum("bqij,bqjd->bqid", attn_mixed, V)
        if self.tied_o:
            # Single (d_v → d) projection applied to each operator slice, then sum across operators.
            out_lifted = self.c_proj(out)  # [B, r_q, T, d]
            return self.dropout(out_lifted.sum(dim=1))  # [B, T, d]
        else:
            out = out.transpose(1, 2).contiguous().view(B, T, -1)  # [B, T, d]
            return self.dropout(self.c_proj(out))


class MLP(nn.Module):
    def __init__(self, cfg: GPTConfig, mlp_mult: int | None = None,
                 mlp_activation: str | None = None):
        super().__init__()
        mult = cfg.mlp_mult if mlp_mult is None else mlp_mult
        inter = mult * cfg.n_embd
        self._activation = cfg.mlp_activation if mlp_activation is None else mlp_activation
        if self._activation not in {"gelu", "silu", "sigmoid", "linear", "silu_entmax_gate"}:
            raise ValueError(f"unknown mlp_activation: {self._activation!r}")
        self.c_fc = nn.Linear(cfg.n_embd, inter, bias=False)
        self.c_proj = nn.Linear(inter, cfg.n_embd, bias=False)
        self.dropout = nn.Dropout(cfg.dropout)
        # Optional output-side RMSNorm calibration (BP-canonical: forces per-MLP
        # delta to a fixed log-odds scale, mirroring logit-RMSNorm on the LM head).
        # scale = 0 disables. scale > 0: y = RMSNorm(c_proj_output) * scale
        self._out_rms_scale = float(cfg.mlp_output_rmsnorm_scale)
        # Buffer for L1 activation sparsity penalty (computed in forward, read by ByteGPT).
        # Set to None when sparsity is off so no extra graph nodes are created.
        self._act_sparsity_lambda = float(cfg.mlp_act_sparsity_lambda)
        self._last_act_l1 = None
        # Hard top-K mask on inner activations (structural sparsity with STE)
        self._act_topk = int(cfg.mlp_act_topk)
        self._inter = inter

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.c_fc(x)
        if self._activation == "linear":
            pass  # exact-linear MLP (reverse-compiled adder L0/L1 form)
        elif self._activation == "gelu":
            x = F.gelu(x, approximate="tanh")
        elif self._activation == "sigmoid":
            x = torch.sigmoid(x)
        elif self._activation == "silu_entmax_gate":
            # BP-canonical: SiLU carries value, entmax15-over-channels gates.
            # Exact-zero gate entries on low-scoring channels; smooth gradients.
            # Same parameter count as plain SiLU (single c_fc).
            v = F.silu(x)
            g = (_pkg_entmax15(x, dim=-1) if _HAVE_ENTMAX_PKG else _entmax15(x, dim=-1))
            x = v * g
        else:  # silu
            x = F.silu(x)
        if self._act_topk > 0 and self._act_topk < self._inter:
            # Per-token hard top-K mask on |x|; straight-through estimator
            # in backward: gradient flows through as if mask were identity.
            k = self._act_topk
            with torch.no_grad():
                _, top_idx = x.abs().topk(k, dim=-1)
                mask = torch.zeros_like(x).scatter_(-1, top_idx, 1.0)
            # STE: forward = x * mask (hard), backward = grad flows through x
            x = x + (x * mask - x).detach()
        if self._act_sparsity_lambda > 0.0:
            # mean(|act|) over all (batch, token, channel) entries
            self._last_act_l1 = x.abs().mean()
        else:
            self._last_act_l1 = None
        x = self.c_proj(x)
        if self._out_rms_scale > 0.0:
            rms = x.pow(2).mean(dim=-1, keepdim=True).clamp(min=1e-8).sqrt()
            x = x / rms * self._out_rms_scale
        return self.dropout(x)


class Block(nn.Module):
    def __init__(self, cfg: GPTConfig, n_head: int | None = None, mlp_mult: int | None = None,
                 mlp_activation: str | None = None):
        super().__init__()
        self.ln_1 = nn.LayerNorm(cfg.n_embd)
        if cfg.attn_kind == "mha":
            self.attn = CausalSelfAttention(cfg, n_head=n_head)
        elif cfg.attn_kind == "tucker":
            self.attn = TuckerAttention(
                cfg, r_p=cfg.tucker_r_p, r_q=cfg.tucker_r_q,
                tied_v=cfg.tucker_tied_v, tied_o=cfg.tucker_tied_o,
                attn_norm=cfg.attn_norm, d_v=cfg.tucker_d_v,
                topk_k=cfg.topk_k,
            )
        elif cfg.attn_kind == "blockdiag":
            self.attn = BlockDiagAttention(
                cfg, block_size=cfg.block_diag_size,
                block_diag_proj=cfg.block_diag_proj,
            )
        elif cfg.attn_kind == "shared_mixer_rotor":
            self.attn = SharedMixerRotorAttention(
                cfg, n_operators=cfg.smr_n_operators, d_v=cfg.smr_d_v,
                r_kind=cfg.smr_r_kind, rotor_block=cfg.smr_rotor_block,
                operator_kinds=cfg.smr_operator_kinds,
                attn_norm=cfg.attn_norm,
                shared_wo=cfg.smr_shared_wo,
                topk_k=cfg.topk_k,
            )
        else:
            raise ValueError(f"unknown attn_kind: {cfg.attn_kind!r}")
        self.ln_2 = nn.LayerNorm(cfg.n_embd)
        self.mlp = MLP(cfg, mlp_mult=mlp_mult, mlp_activation=mlp_activation)

    def forward(self, x: torch.Tensor,
                attn_bias: torch.Tensor | None = None,
                rope_freqs: tuple[torch.Tensor, torch.Tensor] | None = None) -> torch.Tensor:
        attn_kwargs = {}
        if attn_bias is not None:
            attn_kwargs["attn_bias"] = attn_bias
        if rope_freqs is not None:
            attn_kwargs["rope_freqs"] = rope_freqs
        if attn_kwargs:
            x = x + self.attn(self.ln_1(x), **attn_kwargs)
        else:
            x = x + self.attn(self.ln_1(x))
        x = x + self.mlp(self.ln_2(x))
        return x


class TransformerCore(nn.Module):
    def __init__(self, cfg: GPTConfig):
        super().__init__()
        self.wte = nn.Embedding(cfg.vocab_size, cfg.n_embd)
        self.wpe = nn.Embedding(cfg.seq_len, cfg.n_embd)
        if cfg.n_heads_per_block:
            if len(cfg.n_heads_per_block) != cfg.n_layer:
                raise ValueError(
                    f"n_heads_per_block has {len(cfg.n_heads_per_block)} entries; "
                    f"expected {cfg.n_layer} (one per block)."
                )
            heads_per_block = cfg.n_heads_per_block
        else:
            heads_per_block = (cfg.n_head,) * cfg.n_layer
        if cfg.mlp_mult_per_block:
            if len(cfg.mlp_mult_per_block) != cfg.n_layer:
                raise ValueError(
                    f"mlp_mult_per_block has {len(cfg.mlp_mult_per_block)} entries; "
                    f"expected {cfg.n_layer} (one per block)."
                )
            mlp_mult_per_block = cfg.mlp_mult_per_block
        else:
            mlp_mult_per_block = (cfg.mlp_mult,) * cfg.n_layer
        if cfg.mlp_activation_per_block:
            if len(cfg.mlp_activation_per_block) != cfg.n_layer:
                raise ValueError(
                    f"mlp_activation_per_block has {len(cfg.mlp_activation_per_block)} entries; "
                    f"expected {cfg.n_layer} (one per block)."
                )
            mlp_act_per_block = cfg.mlp_activation_per_block
        else:
            mlp_act_per_block = (cfg.mlp_activation,) * cfg.n_layer
        self.h = nn.ModuleList([
            Block(cfg, n_head=heads_per_block[i], mlp_mult=mlp_mult_per_block[i],
                  mlp_activation=mlp_act_per_block[i])
            for i in range(cfg.n_layer)
        ])
        self.ln_f = nn.LayerNorm(cfg.n_embd)


class ByteGPT(nn.Module):
    def __init__(self, cfg: GPTConfig):
        super().__init__()
        if cfg.loss_kind not in {"ce", "bce"}:
            raise ValueError(f"unknown loss_kind: {cfg.loss_kind!r}")
        if cfg.position_encoding not in {"learned", "alibi", "rope", "none"}:
            raise ValueError(f"unknown position_encoding: {cfg.position_encoding!r}")
        if cfg.position_encoding in {"alibi", "rope"} and cfg.attn_kind not in {"mha", "tucker"}:
            raise ValueError(
                f"position_encoding={cfg.position_encoding!r} not supported for "
                f"attn_kind={cfg.attn_kind!r} (only mha and tucker supported)"
            )
        self.config = cfg
        self.transformer = TransformerCore(cfg)
        self.lm_head = nn.Linear(cfg.n_embd, cfg.vocab_size, bias=False)
        # ALiBi: pre-compute slopes (one per attention "channel": r_p for tucker, n_head for mha).
        if cfg.position_encoding == "alibi":
            n_slopes = cfg.tucker_r_p if cfg.attn_kind == "tucker" else cfg.n_head
            slopes = _alibi_slopes(n_slopes) * cfg.alibi_slope_scale
            self.register_buffer("_alibi_slopes", slopes, persistent=False)
            self._alibi_n_slopes = n_slopes
        else:
            self._alibi_slopes = None
            self._alibi_n_slopes = 0
        # RoPE: head_dim differs between attn kinds (d_qk = d/r_p for tucker, d/n_head for mha).
        if cfg.position_encoding == "rope":
            if cfg.attn_kind == "tucker":
                self._rope_head_dim = cfg.n_embd // cfg.tucker_r_p
            else:
                self._rope_head_dim = cfg.n_embd // cfg.n_head
            if self._rope_head_dim % 2 != 0:
                raise ValueError(f"RoPE requires even head_dim, got {self._rope_head_dim}")
        else:
            self._rope_head_dim = 0
        if cfg.block_gates:
            if cfg.block_gate_kind == "linear":
                init_val = 1.0
            elif cfg.block_gate_kind in ("sigmoid", "ste"):
                # σ(4) ≈ 0.982 ≈ "on". α=0 → γ=0.5 (mid). α=-4 → γ=0.018 ≈ "off".
                init_val = 4.0
            else:
                raise ValueError(f"unknown block_gate_kind: {cfg.block_gate_kind!r}")
            self.block_gate = nn.Parameter(torch.full((cfg.n_layer,), init_val))
        else:
            self.block_gate = None
        # Attention Residuals: one learned pseudo-query w_l ∈ R^d per
        # aggregation. n_layer queries feed block inputs h_0..h_{L-1};
        # the (L+1)-th feeds the final readout into ln_f. Zero-init →
        # uniform α at step 0 → AttnRes ≡ equal-weight depth average
        # (the paper's stability trick); the network sharpens from there.
        if cfg.res_attn != "none":
            self.res_attn_q = nn.Parameter(torch.zeros(cfg.n_layer + 1, cfg.n_embd))
        else:
            self.res_attn_q = None
        # Diagnostics: when set True (by analyze_attnres_depth_matrix.py),
        # forward stashes the token-averaged depth-mixing rows in
        # self._depth_alpha_rows (list of 1-D tensors, ragged lower-triangular).
        self._capture_depth_alpha = False
        self._depth_alpha_rows: List[torch.Tensor] | None = None

    def _gate_values(self) -> torch.Tensor | None:
        """Resolve raw block_gate parameter to per-block gate values γ_ℓ in the chosen parameterization."""
        if self.block_gate is None:
            return None
        kind = self.config.block_gate_kind
        if kind == "linear":
            return self.block_gate
        soft = torch.sigmoid(self.block_gate)
        if kind == "sigmoid":
            return soft
        # ste: forward = hard 0/1, backward = soft sigmoid gradient
        hard = (soft > 0.5).float()
        return hard + (soft - soft.detach())

    def _depth_aggregate(self, vs: List[torch.Tensor], q: torch.Tensor
                         ) -> tuple[torch.Tensor, torch.Tensor]:
        """One Attention-Residuals aggregation (Eq. 2–4 of arXiv:2603.15031).
        vs: list of S source tensors [B,T,d] (v_0=embedding, v_i=block Δ).
        q:  this aggregation's learned pseudo-query [d].
        Returns (h [B,T,d], alpha [B,T,S]). RMSNorm on keys makes the
        routing magnitude-invariant; α is normalized over the S sources."""
        V = torch.stack(vs, dim=2)                                  # [B,T,S,d]
        k = V * torch.rsqrt(V.pow(2).mean(-1, keepdim=True) + 1e-6)  # RMSNorm(v_i)
        scores = torch.einsum("btsd,d->bts", k, q)                   # w_lᵀ RMSNorm(v_i)
        scores = scores * self.config.res_attn_temp                   # A2′ temperature τ
        alpha = _depth_attn_norm(scores, self.config.res_attn_norm)  # over S
        h = torch.einsum("bts,btsd->btd", alpha, V)
        return h, alpha

    def _forward_attnres(self, x: torch.Tensor,
                         attn_bias: torch.Tensor | None,
                         rope_freqs: tuple[torch.Tensor, torch.Tensor] | None
                         ) -> torch.Tensor:
        """Block-granular Full Attention Residuals. The depth-unit is one
        transformer block; v_0 is the (post-PE) embedding and v_{i+1} is
        block i's residual-free contribution f_i(h_i) = Block_i(h_i)−h_i.
        Each block input and the final readout is a learned aggregation
        over all preceding v's. See docs/attnres-bp-research-program.md."""
        blk_kwargs = {}
        if attn_bias is not None:
            blk_kwargs["attn_bias"] = attn_bias
        if rope_freqs is not None:
            blk_kwargs["rope_freqs"] = rope_freqs
        vs: List[torch.Tensor] = [x]                       # v_0 = embedding
        rows: List[torch.Tensor] | None = [] if self._capture_depth_alpha else None
        ent_terms: List[torch.Tensor] = []

        def _entropy(a: torch.Tensor) -> torch.Tensor:
            # H(α) over the source axis; exact-zero-safe (0·log0→0).
            return -(a * a.clamp_min(1e-12).log()).sum(dim=-1).mean()

        for b, block in enumerate(self.transformer.h):
            h_in, alpha = self._depth_aggregate(vs, self.res_attn_q[b])
            if rows is not None:
                rows.append(alpha.detach().float().mean(dim=(0, 1)).cpu())
            if alpha.size(-1) > 1:
                ent_terms.append(_entropy(alpha))
            out = block(h_in, **blk_kwargs) if blk_kwargs else block(h_in)
            vs.append(out - h_in)                          # v_{b+1} = f_b(h_in)
        x_out, alpha = self._depth_aggregate(vs, self.res_attn_q[len(self.transformer.h)])
        if rows is not None:
            rows.append(alpha.detach().float().mean(dim=(0, 1)).cpu())
            self._depth_alpha_rows = rows
        if alpha.size(-1) > 1:
            ent_terms.append(_entropy(alpha))
        self._depth_entropy = (torch.stack(ent_terms).mean() if ent_terms
                               else x_out.new_zeros(()))
        return x_out

    def forward(self, idx: torch.Tensor, targets: torch.Tensor | None = None,
                loop_block_k_override: int | None = None):
        bsz, seqlen = idx.shape
        if seqlen > self.config.seq_len:
            raise ValueError(f"sequence length {seqlen} exceeds model limit {self.config.seq_len}")
        tok = self.transformer.wte(idx)
        attn_bias: torch.Tensor | None = None
        rope_freqs: tuple[torch.Tensor, torch.Tensor] | None = None
        pe = self.config.position_encoding
        if pe == "learned":
            pos = torch.arange(seqlen, device=idx.device)
            pos_emb = self.transformer.wpe(pos)[None, :, :]
            x = tok + pos_emb
        elif pe == "alibi":
            x = tok
            attn_bias = _alibi_bias(
                self._alibi_n_slopes, seqlen, self._alibi_slopes,
                device=idx.device, dtype=x.dtype,
            )
        elif pe == "rope":
            x = tok
            rope_freqs = _make_rope_freqs(
                self._rope_head_dim, seqlen, base=self.config.rope_base,
                device=idx.device, dtype=x.dtype,
            )
        else:  # none
            x = tok
        gates = self._gate_values()  # [n_layer] or None
        attractor_terms: List[torch.Tensor] = []
        attractor_set = set(self.config.block_attractor_indices)
        loop_set = set(self.config.loop_block_indices)
        unit_idx = tuple(self.config.loop_unit_indices)
        if unit_idx:
            # Must be contiguous and disjoint from loop_block_indices.
            assert all(unit_idx[k + 1] - unit_idx[k] == 1 for k in range(len(unit_idx) - 1)), \
                f"loop_unit_indices must be contiguous, got {unit_idx}"
            assert not (set(unit_idx) & loop_set), \
                "loop_unit_indices and loop_block_indices must not overlap"
        unit_start = unit_idx[0] if unit_idx else None
        unit_end = unit_idx[-1] if unit_idx else None
        K_loop = max(1, int(loop_block_k_override if loop_block_k_override is not None else self.config.loop_block_k))

        if self.config.res_attn != "none":
            # Attention Residuals replace the additive residual stream with a
            # learned depth-mixing aggregation; the v1 path is deliberately
            # not composed with the gate/loop/attractor regularizers (that is
            # stage A4 in docs/attnres-bp-research-program.md).
            if (self.block_gate is not None or self.config.loop_block_indices
                    or self.config.loop_unit_indices
                    or self.config.block_attractor_indices):
                raise ValueError(
                    "res_attn is not yet composable with block gates / loop-block / "
                    "block-attractor (v1); see docs/attnres-bp-research-program.md."
                )
            x = self._forward_attnres(x, attn_bias, rope_freqs)
        else:
            def _apply_block(block, x_in, i):
                if attn_bias is not None or rope_freqs is not None:
                    y = block(x_in, attn_bias=attn_bias, rope_freqs=rope_freqs)
                else:
                    y = block(x_in)
                contribution = (y - x_in) if gates is None else gates[i] * (y - x_in)
                return x_in + contribution

            i = 0
            n_blocks = len(self.transformer.h)
            while i < n_blocks:
                if unit_start is not None and i == unit_start:
                    # Apply blocks [unit_start..unit_end] together K_loop times.
                    x_pre_unit = x
                    for _ in range(K_loop):
                        for j in range(unit_start, unit_end + 1):
                            x = _apply_block(self.transformer.h[j], x, j)
                    if targets is not None:
                        for j in range(unit_start, unit_end + 1):
                            if j in attractor_set:
                                attractor_terms.append((x - x_pre_unit).pow(2).mean())
                                break
                    i = unit_end + 1
                    continue
                block = self.transformer.h[i]
                x_in_pos = x
                iters = K_loop if i in loop_set else 1
                for _ in range(iters):
                    x = _apply_block(block, x, i)
                if targets is not None and i in attractor_set:
                    attractor_terms.append((x - x_in_pos).pow(2).mean())
                i += 1
        x = self.transformer.ln_f(x)
        logits = self.lm_head(x)
        loss = None
        if targets is not None:
            loss = self._compute_loss(logits, targets)
            if (self.config.res_attn != "none"
                    and self.config.res_attn_entropy_lambda > 0):
                loss = loss + self.config.res_attn_entropy_lambda * self._depth_entropy
            if attractor_terms and self.config.block_attractor_lambda > 0:
                loss = loss + self.config.block_attractor_lambda * sum(attractor_terms)
            if gates is not None and self.config.block_gate_l1 > 0:
                loss = loss + self.config.block_gate_l1 * gates.abs().sum()
            if self.config.mlp_act_sparsity_lambda > 0.0:
                act_terms = [blk.mlp._last_act_l1 for blk in self.transformer.h
                             if blk.mlp._last_act_l1 is not None]
                if act_terms:
                    loss = loss + self.config.mlp_act_sparsity_lambda * sum(act_terms)
            if (self.config.mlp_byte_concentration_lambda > 0.0
                    and self.config.mlp_byte_concentration_layers):
                conc_loss = self._mlp_byte_concentration_loss()
                # Negative: maximizing concentration ⇒ subtract from loss
                loss = loss - self.config.mlp_byte_concentration_lambda * conc_loss
        return logits, loss

    def _mlp_byte_concentration_loss(self) -> torch.Tensor:
        """Mean top-K row-mass fraction of M = W_lm · W_proj · W_fc · W_wte^T
        across selected layers. Higher = MLP is more BP-readable as a sparse
        transition factor (each output byte explained by few input bytes)."""
        wte = self.transformer.wte.weight                            # [V, d]
        lm_head = self.lm_head.weight                                 # [V, d]
        K = self.config.mlp_byte_concentration_topk
        layers = self.config.mlp_byte_concentration_layers
        terms: List[torch.Tensor] = []
        for i in layers:
            blk = self.transformer.h[i]
            W_fc = blk.mlp.c_fc.weight                                # [inter, d]
            W_proj = blk.mlp.c_proj.weight                            # [d, inter]
            W_lin = W_proj @ W_fc                                      # [d, d]
            M = lm_head @ W_lin @ wte.T                                # [V, V]
            M_sq = M.pow(2)
            row_sums = M_sq.sum(dim=-1, keepdim=True).clamp(min=1e-8)
            topk_vals, _ = M_sq.topk(K, dim=-1)
            topk_mass = topk_vals.sum(dim=-1, keepdim=True) / row_sums
            terms.append(topk_mass.mean())
        return torch.stack(terms).mean() if terms else torch.tensor(0.0,
            device=self.lm_head.weight.device)

    def _compute_loss(self, logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
        flat_logits = logits.view(-1, logits.size(-1))
        flat_targets = targets.reshape(-1)
        if self.config.loss_kind == "ce":
            return F.cross_entropy(flat_logits, flat_targets, ignore_index=-100)
        # BCE: each vocab logit is an independent Bernoulli. Gradient on logit_i
        # is sigma(z_i) - y_i, so it does not couple across classes the way
        # softmax CE does through the partition function.
        # Mask out positions with target == -100 (parallels F.cross_entropy ignore_index).
        valid_mask = flat_targets != -100
        if not valid_mask.all():
            flat_logits = flat_logits[valid_mask]
            flat_targets = flat_targets[valid_mask]
        if flat_targets.numel() == 0:
            return flat_logits.new_zeros(())
        if self.config.logit_rmsnorm_scale > 0:
            flat_logits = F.rms_norm(flat_logits, (flat_logits.size(-1),)) * self.config.logit_rmsnorm_scale
        one_hot = F.one_hot(flat_targets, num_classes=flat_logits.size(-1)).to(flat_logits.dtype)
        per_elem = F.binary_cross_entropy_with_logits(flat_logits, one_hot, reduction="none")
        loss = per_elem.sum(dim=-1).mean()
        if self.config.sparsity_lambda > 0:
            neg_mask = 1.0 - one_hot
            sparsity = (torch.sigmoid(flat_logits) * neg_mask).sum(dim=-1).mean()
            loss = loss + self.config.sparsity_lambda * sparsity
        return loss

    @torch.no_grad()
    def generate(self, idx: torch.Tensor, max_new_tokens: int) -> torch.Tensor:
        for _ in range(max_new_tokens):
            idx_cond = idx[:, -self.config.seq_len :]
            logits, _ = self(idx_cond)
            next_id = torch.argmax(logits[:, -1, :], dim=-1, keepdim=True)
            idx = torch.cat([idx, next_id], dim=1)
        return idx


class RandomByteCorpus:
    def sample_batch(self, batch_size: int, seq_len: int, device: str) -> tuple[torch.Tensor, torch.Tensor]:
        raise NotImplementedError

    def stats(self) -> Dict[str, int | str]:
        raise NotImplementedError


class MixtureByteCorpus(RandomByteCorpus):
    def __init__(self, corpora: Sequence[RandomByteCorpus], weights: Sequence[float]):
        if len(corpora) != len(weights):
            raise ValueError("corpora and weights must have the same length")
        if not corpora:
            raise ValueError("mixture corpus needs at least one member")
        total = float(sum(weights))
        if total <= 0.0:
            raise ValueError("mixture weights must sum to a positive value")
        self.corpora = list(corpora)
        self.weights = [float(weight) / total for weight in weights]

    def sample_batch(self, batch_size: int, seq_len: int, device: str) -> tuple[torch.Tensor, torch.Tensor]:
        choices = random.choices(range(len(self.corpora)), weights=self.weights, k=batch_size)
        counts: Dict[int, int] = {}
        for idx in choices:
            counts[idx] = counts.get(idx, 0) + 1
        xs = []
        ys = []
        for idx, count in counts.items():
            x, y = self.corpora[idx].sample_batch(batch_size=count, seq_len=seq_len, device=device)
            xs.append(x)
            ys.append(y)
        x_full = torch.cat(xs, dim=0)
        y_full = torch.cat(ys, dim=0)
        perm = torch.randperm(x_full.size(0), device=x_full.device)
        return x_full[perm], y_full[perm]

    def stats(self) -> Dict[str, int | str]:
        payload = {"kind": "mixture"}
        for idx, (corpus, weight) in enumerate(zip(self.corpora, self.weights)):
            payload[f"member_{idx}_weight"] = round(weight, 6)
            payload[f"member_{idx}_kind"] = corpus.stats().get("kind", "unknown")
        return payload


def iter_projected_params(model: ByteGPT):
    for layer_idx, block in enumerate(model.transformer.h):
        yield (
            f"transformer.h.{layer_idx}.attn.c_proj.weight",
            block.attn.c_proj.weight,
            layer_idx,
            "attn_out",
        )
        yield (
            f"transformer.h.{layer_idx}.mlp.c_proj.weight",
            block.mlp.c_proj.weight,
            layer_idx,
            "mlp_down",
        )


class InMemoryByteCorpus(RandomByteCorpus):
    def __init__(self, data: bytes):
        self.data = torch.tensor(list(data), dtype=torch.long)

    def sample_batch(self, batch_size: int, seq_len: int, device: str) -> tuple[torch.Tensor, torch.Tensor]:
        max_start = self.data.size(0) - seq_len - 1
        if max_start <= 0:
            raise ValueError("corpus too small for requested sequence length")
        ix = torch.randint(0, max_start, (batch_size,))
        x = torch.stack([self.data[i : i + seq_len] for i in ix]).to(device)
        y = torch.stack([self.data[i + 1 : i + 1 + seq_len] for i in ix]).to(device)
        return x, y

    def stats(self) -> Dict[str, int | str]:
        return {"kind": "in_memory_bytes", "bytes": int(self.data.numel())}


def _make_addition_example(max_digits: int, reverse: bool = True) -> str:
    n_digits_a = random.randint(1, max_digits)
    n_digits_b = random.randint(1, max_digits)
    a = random.randint(0, 10 ** n_digits_a - 1)
    b = random.randint(0, 10 ** n_digits_b - 1)
    answer = str(a + b)
    if reverse:
        answer = answer[::-1]
    return f"{a}+{b}={answer}\n"


class ArithmeticByteCorpus(RandomByteCorpus):
    """Synthetic addition examples encoded as ASCII bytes (vocab=256).

    Reverse-LSB answer aligns autoregressive generation with the carry chain
    direction, which is the BP-canonical schedule for addition.
    """

    def __init__(self, max_digits: int = 3, reverse: bool = True):
        if max_digits < 1:
            raise ValueError("max_digits must be >= 1")
        self.max_digits = max_digits
        self.reverse = reverse

    def sample_batch(self, batch_size: int, seq_len: int, device: str) -> tuple[torch.Tensor, torch.Tensor]:
        need = seq_len + 1
        xs = []
        ys = []
        for _ in range(batch_size):
            buf = bytearray()
            while len(buf) < need:
                ex = _make_addition_example(self.max_digits, reverse=self.reverse)
                buf.extend(ex.encode("ascii"))
            ids = torch.tensor(list(buf[:need]), dtype=torch.long)
            xs.append(ids[:-1])
            ys.append(ids[1:])
        return torch.stack(xs).to(device), torch.stack(ys).to(device)

    def stats(self) -> Dict[str, int | str]:
        return {
            "kind": "arithmetic",
            "max_digits": self.max_digits,
            "reverse": "true" if self.reverse else "false",
        }


class ParquetTextCorpus(RandomByteCorpus):
    def __init__(self, parquet_paths: Sequence[Path], text_column: str = "text",
                 cache_target: int = 1024):
        if pq is None:
            raise RuntimeError("pyarrow is required for parquet corpora")
        if not parquet_paths:
            raise ValueError("no parquet shards provided")
        self.text_column = text_column
        self.paths = [Path(p) for p in parquet_paths]
        self._files = [pq.ParquetFile(path) for path in self.paths]
        self._row_groups: List[tuple[int, int, int]] = []
        self._total_docs = 0
        for file_idx, pf in enumerate(self._files):
            for rg_idx in range(pf.num_row_groups):
                meta = pf.metadata.row_group(rg_idx)
                rows = meta.num_rows
                self._row_groups.append((file_idx, rg_idx, rows))
                self._total_docs += rows
        # Document cache: each refill reads several row groups and pre-encodes
        # to bytes, so sample_batch's per-element call is O(1) RAM lookup + tensor
        # build. Reading a row group is hundreds of ms — without caching, that
        # cost is paid once per batch element and starves the GPU.
        self._cache_target = int(cache_target)
        self._cache: List[bytes] = []

    def _refill_cache(self) -> None:
        while len(self._cache) < self._cache_target:
            file_idx, rg_idx, _ = random.choice(self._row_groups)
            rg = self._files[file_idx].read_row_group(rg_idx, columns=[self.text_column])
            texts = rg.column(self.text_column).to_pylist()
            for t in texts:
                if t:
                    self._cache.append(t.encode("utf-8"))
        random.shuffle(self._cache)

    def _sample_document_bytes(self, min_bytes: int) -> bytes:
        for _ in range(64):
            if not self._cache:
                self._refill_cache()
            raw = self._cache.pop()
            if len(raw) >= min_bytes:
                return raw
        # Last-resort: rebuild and try once more.
        self._refill_cache()
        for _ in range(64):
            if not self._cache:
                self._refill_cache()
            raw = self._cache.pop()
            if len(raw) >= min_bytes:
                return raw
        raise ValueError(f"could not sample a document with at least {min_bytes} bytes")

    def sample_batch(self, batch_size: int, seq_len: int, device: str) -> tuple[torch.Tensor, torch.Tensor]:
        xs = []
        ys = []
        need = seq_len + 1
        for _ in range(batch_size):
            raw = self._sample_document_bytes(need)
            start_max = len(raw) - need
            start = 0 if start_max <= 0 else random.randint(0, start_max)
            chunk = raw[start : start + need]
            ids = torch.tensor(list(chunk), dtype=torch.long)
            xs.append(ids[:-1])
            ys.append(ids[1:])
        return torch.stack(xs).to(device), torch.stack(ys).to(device)

    def stats(self) -> Dict[str, int | str]:
        return {
            "kind": "parquet_text",
            "shards": len(self.paths),
            "row_groups": len(self._row_groups),
            "documents": int(self._total_docs),
        }


def load_tokenized_corpus(paths: Sequence[Path], text_key: str, tokenizer) -> List[int]:
    """Encode JSONL/text files into a flat list of token IDs using a
    HuggingFace `tokenizers.Tokenizer` instance. A single newline-id (encoded
    from "\\n") separates documents so that downstream chunking does not
    splice two unrelated stories without any boundary signal."""
    sep_ids: List[int] = tokenizer.encode("\n").ids or []
    ids: List[int] = []
    for path in paths:
        if path.suffix == ".jsonl":
            with path.open() as f:
                for line in f:
                    line = line.strip()
                    if not line:
                        continue
                    obj = json.loads(line)
                    text = obj.get(text_key, "")
                    if text:
                        ids.extend(tokenizer.encode(text).ids)
                        ids.extend(sep_ids)
        else:
            ids.extend(tokenizer.encode(path.read_text()).ids)
            ids.extend(sep_ids)
    if not ids:
        raise ValueError("no corpus text found")
    return ids


def split_tokenized(ids: List[int], val_fraction: float) -> tuple[List[int], List[int]]:
    cut = max(1, min(len(ids) - 1, int(len(ids) * (1.0 - val_fraction))))
    return ids[:cut], ids[cut:]


def load_text_corpus(paths: Sequence[Path], text_key: str) -> bytes:
    chunks: List[bytes] = []
    for path in paths:
        if path.suffix == ".jsonl":
            with path.open() as f:
                for line in f:
                    line = line.strip()
                    if not line:
                        continue
                    obj = json.loads(line)
                    text = obj.get(text_key, "")
                    if text:
                        chunks.append(text.encode("utf-8") + b"\n")
        else:
            chunks.append(path.read_bytes() + b"\n")
    if not chunks:
        raise ValueError("no corpus text found")
    return b"".join(chunks)


def split_corpus(data: bytes, val_fraction: float) -> tuple[bytes, bytes]:
    cut = max(1, min(len(data) - 1, int(len(data) * (1.0 - val_fraction))))
    train = data[:cut]
    val = data[cut:]
    return train, val


def build_corpora(
    paths: Sequence[Path],
    text_key: str,
    val_fraction: float,
    parquet_mix_weight: float,
    text_mix_weight: float,
    tokenizer=None,
) -> tuple[RandomByteCorpus, RandomByteCorpus]:
    parquet_paths = [p for p in paths if p.suffix == ".parquet"]
    non_parquet = [p for p in paths if p.suffix != ".parquet"]
    if parquet_paths:
        if tokenizer is not None:
            raise ValueError("tokenizer mode is not supported with parquet shards yet")
        if len(parquet_paths) < 2:
            raise ValueError("need at least 2 parquet shards so the last shard can serve as validation")
        parquet_paths = sorted(parquet_paths)
        parquet_train = ParquetTextCorpus(parquet_paths[:-1], text_column=text_key)
        parquet_val = ParquetTextCorpus(parquet_paths[-1:], text_column=text_key)
        if not non_parquet:
            return (
                parquet_train,
                parquet_val,
            )
        raw = load_text_corpus(non_parquet, text_key=text_key)
        train_raw, val_raw = split_corpus(raw, val_fraction=val_fraction)
        text_train = InMemoryByteCorpus(train_raw)
        text_val = InMemoryByteCorpus(val_raw)
        return (
            MixtureByteCorpus([parquet_train, text_train], [parquet_mix_weight, text_mix_weight]),
            MixtureByteCorpus([parquet_val, text_val], [parquet_mix_weight, text_mix_weight]),
        )
    if tokenizer is not None:
        ids = load_tokenized_corpus(non_parquet, text_key=text_key, tokenizer=tokenizer)
        train_ids, val_ids = split_tokenized(ids, val_fraction=val_fraction)
        return InMemoryByteCorpus(train_ids), InMemoryByteCorpus(val_ids)
    raw = load_text_corpus(non_parquet, text_key=text_key)
    train_raw, val_raw = split_corpus(raw, val_fraction=val_fraction)
    return InMemoryByteCorpus(train_raw), InMemoryByteCorpus(val_raw)


def bits_per_byte(loss: float) -> float:
    return loss / math.log(2.0)


def _json_safe(value):
    if isinstance(value, Path):
        return str(value)
    if isinstance(value, list):
        return [_json_safe(v) for v in value]
    if isinstance(value, dict):
        return {k: _json_safe(v) for k, v in value.items()}
    return value


def load_probe_rows(paths: Sequence[Path]) -> Dict[str, List[Dict]]:
    groups: Dict[str, List[Dict]] = {}
    for path in paths:
        rows: List[Dict] = []
        with path.open() as f:
            for line in f:
                line = line.strip()
                if not line:
                    continue
                rows.append(json.loads(line))
        groups[path.stem] = rows
    return groups


def complete_bytes(model: ByteGPT, device: str, prompt: str, max_new_tokens: int) -> str:
    ids = torch.tensor([list(prompt.encode("utf-8"))], dtype=torch.long, device=device)
    out = model.generate(ids, max_new_tokens=max_new_tokens)[0].tolist()
    return bytes(out).decode("utf-8", errors="ignore")


def format_sample(model: ByteGPT, device: str, prompt: str, max_new_tokens: int = 80) -> str:
    try:
        return complete_bytes(model, device=device, prompt=prompt, max_new_tokens=max_new_tokens)
    except Exception:
        return repr(prompt)


@torch.no_grad()
def evaluate_probe_groups(model: ByteGPT, probe_groups: Dict[str, List[Dict]], device: str) -> Dict[str, Dict]:
    if not probe_groups:
        return {}
    model.eval()
    payload: Dict[str, Dict] = {}
    for name, rows in probe_groups.items():
        results: List[Dict] = []
        for row in rows:
            prompt = row["prompt"]
            target = row["target"]
            match_mode = row.get("match", "exact")
            max_new = int(row.get("max_new_tokens", max(len(target.encode("utf-8")), 1) + 8))
            generated = complete_bytes(model, device=device, prompt=prompt, max_new_tokens=max_new)
            completion = generated[len(prompt) :] if generated.startswith(prompt) else generated
            got = completion.strip()
            want = target.strip()
            if match_mode == "contains":
                correct = want in got
            else:
                correct = got.startswith(want)
            results.append(
                {
                    "prompt": prompt,
                    "target": target,
                    "generated": got,
                    "correct": bool(correct),
                    "match": match_mode,
                }
            )
        payload[name] = {
            "n": len(results),
            "correct": sum(1 for row in results if row["correct"]),
            "examples": results[:20],
        }
    return payload


def save_checkpoint(
    path: Path,
    *,
    model: ByteGPT,
    optimizer: torch.optim.Optimizer,
    step: int,
    history: List[Dict],
    mode: str,
    args: argparse.Namespace,
    model_config: GPTConfig,
) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    payload = {
        "step": step,
        "mode": mode,
        "args": _json_safe(vars(args)),
        "model_config": asdict(model_config),
        "model_state": model.state_dict(),
        "optimizer_state": optimizer.state_dict(),
        "history": history,
        "rng_state": torch.random.get_rng_state(),
        "cuda_rng_state": torch.cuda.get_rng_state_all() if torch.cuda.is_available() else None,
    }
    torch.save(payload, path)


def load_checkpoint(path: Path, model: ByteGPT, optimizer: torch.optim.Optimizer, device: str) -> Dict:
    payload = torch.load(path, map_location=device, weights_only=False)
    model.load_state_dict(payload["model_state"])
    optimizer.load_state_dict(payload["optimizer_state"])
    torch.random.set_rng_state(payload["rng_state"].cpu())
    if torch.cuda.is_available() and payload.get("cuda_rng_state") is not None:
        torch.cuda.set_rng_state_all([state.cpu() for state in payload["cuda_rng_state"]])
    return payload


def infer_mlp_activation(args: argparse.Namespace, device: str) -> str:
    if args.mlp_activation is not None:
        return args.mlp_activation
    if args.resume is None:
        return "silu"
    payload = torch.load(args.resume, map_location=device, weights_only=False)
    model_cfg = payload.get("model_config", {})
    return str(model_cfg.get("mlp_activation", "gelu"))


def infer_loss_kind(args: argparse.Namespace, device: str) -> str:
    if args.loss is not None:
        return args.loss
    if args.resume is None:
        return "ce"
    payload = torch.load(args.resume, map_location=device, weights_only=False)
    model_cfg = payload.get("model_config", {})
    return str(model_cfg.get("loss_kind", "ce"))


def infer_bce_reg(args: argparse.Namespace, device: str) -> tuple[float, float]:
    """Resolve --logit-rmsnorm-scale / --sparsity-lambda, falling back to checkpoint if resuming."""
    rms = args.logit_rmsnorm_scale
    spars = args.sparsity_lambda
    if args.resume is not None and (rms is None or spars is None):
        payload = torch.load(args.resume, map_location=device, weights_only=False)
        model_cfg = payload.get("model_config", {})
        if rms is None:
            # Backward compat: old bool field "logit_rmsnorm" → 1.0 if True else 0.0
            old_bool = bool(model_cfg.get("logit_rmsnorm", False))
            rms = float(model_cfg.get("logit_rmsnorm_scale", 1.0 if old_bool else 0.0))
        if spars is None:
            spars = float(model_cfg.get("sparsity_lambda", 0.0))
    return float(rms) if rms is not None else 0.0, float(spars) if spars is not None else 0.0


def apply_coppola_projection(
    model: ByteGPT,
    controller: CoppolaPretrainingController,
    mode: str,
) -> Dict[str, Dict[str, float]]:
    stats: Dict[str, Dict[str, float]] = {"attn_out": {}, "mlp_down": {}}
    for layer_idx, block in enumerate(model.transformer.h):
        attn_param = block.attn.c_proj.weight
        if attn_param.grad is not None:
            norms = controller.gradient_component_norms(layer_idx, attn_param.grad, param_kind="attn_out")
            attn_param.grad = controller.project_gradient(layer_idx, attn_param.grad, param_kind="attn_out", mode=mode)
            stats["attn_out"][str(layer_idx)] = {
                "routing": norms.routing,
                "remainder": norms.remainder,
            }
        mlp_param = block.mlp.c_proj.weight
        if mlp_param.grad is not None:
            norms = controller.gradient_component_norms(layer_idx, mlp_param.grad, param_kind="mlp_down")
            mlp_param.grad = controller.project_gradient(layer_idx, mlp_param.grad, param_kind="mlp_down", mode=mode)
            stats["mlp_down"][str(layer_idx)] = asdict(norms)
    return stats


def projected_param_names(model: ByteGPT) -> set[str]:
    return {name for name, _param, _layer_idx, _kind in iter_projected_params(model)}


@torch.no_grad()
def snapshot_projected_params(model: ByteGPT):
    return [
        (param, param.detach().clone(), layer_idx, kind)
        for _name, param, layer_idx, kind in iter_projected_params(model)
    ]


@torch.no_grad()
def project_parameter_updates(
    snapshots,
    controller: CoppolaPretrainingController,
) -> None:
    for param, before, layer_idx, kind in snapshots:
        delta = param - before
        delta = controller.project_update_support(layer_idx, delta, param_kind=kind)
        param.copy_(before + delta)


def build_optimizer(
    model: ByteGPT,
    lr: float,
    weight_decay: float,
    projected_weight_decay: float,
) -> torch.optim.Optimizer:
    proj_names = projected_param_names(model)
    projected: List[torch.nn.Parameter] = []
    other: List[torch.nn.Parameter] = []
    for name, param in model.named_parameters():
        if not param.requires_grad:
            continue
        if name in proj_names:
            projected.append(param)
        else:
            other.append(param)
    groups = []
    if projected:
        groups.append({"params": projected, "weight_decay": projected_weight_decay})
    if other:
        groups.append({"params": other, "weight_decay": weight_decay})
    return torch.optim.AdamW(groups, lr=lr)


@torch.no_grad()
def evaluate_arithmetic_accuracy(
    model: ByteGPT,
    device: str,
    n_problems: int,
    max_digits: int,
    reverse: bool,
) -> Dict[str, float]:
    """Per-problem addition accuracy: greedy generation, compare to ground truth."""
    model.eval()
    correct = 0
    by_len: Dict[int, List[int]] = {}  # max(n_a, n_b) -> [correct, total]
    nl_id = ord("\n")
    for _ in range(n_problems):
        n_a = random.randint(1, max_digits)
        n_b = random.randint(1, max_digits)
        a = random.randint(0, 10 ** n_a - 1)
        b = random.randint(0, 10 ** n_b - 1)
        true_answer = str(a + b)
        if reverse:
            true_answer = true_answer[::-1]
        prompt = f"{a}+{b}="
        out_ids = torch.tensor([list(prompt.encode("ascii"))], dtype=torch.long, device=device)
        max_new = len(true_answer) + 2
        gen = bytearray()
        for _ in range(max_new):
            logits, _ = model(out_ids)
            next_id = int(logits[0, -1].argmax())
            if next_id == nl_id:
                break
            gen.append(next_id)
            out_ids = torch.cat([out_ids, torch.tensor([[next_id]], device=device)], dim=1)
        try:
            generated = gen.decode("ascii", errors="ignore")
        except Exception:
            generated = ""
        is_correct = generated == true_answer
        n_digits = max(n_a, n_b)
        bucket = by_len.setdefault(n_digits, [0, 0])
        bucket[1] += 1
        if is_correct:
            bucket[0] += 1
            correct += 1
    out: Dict[str, float] = {
        "accuracy": correct / max(1, n_problems),
        "n": float(n_problems),
    }
    for k, (c, t) in sorted(by_len.items()):
        out[f"acc_{k}d"] = c / max(1, t)
    return out


@torch.no_grad()
def evaluate(
    model: ByteGPT,
    corpus: RandomByteCorpus,
    batch_size: int,
    seq_len: int,
    device: str,
    eval_batches: int,
) -> Dict[str, float]:
    model.eval()
    train_losses = []
    ce_losses = []
    for _ in range(eval_batches):
        x, y = corpus.sample_batch(batch_size=batch_size, seq_len=seq_len, device=device)
        logits, loss = model(x, y)
        train_losses.append(float(loss.item()))
        # bpb is reported from CE so it stays comparable across loss_kind modes.
        ce = F.cross_entropy(logits.view(-1, logits.size(-1)), y.reshape(-1))
        ce_losses.append(float(ce.item()))
    avg_loss = sum(train_losses) / max(1, len(train_losses))
    avg_ce = sum(ce_losses) / max(1, len(ce_losses))
    return {"loss": avg_loss, "ce_loss": avg_ce, "bpb": bits_per_byte(avg_ce)}


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--input", nargs="+", type=Path, default=[ROOT / "content" / "generic_gutenberg" / "raw"])
    ap.add_argument("--text-key", default="text")
    ap.add_argument("--tokenizer-path", type=Path, default=None,
                    help="Path to a HuggingFace tokenizers .json file. If set, "
                         "text is tokenized with it and vocab_size is taken from the "
                         "tokenizer. If unset (default), training is byte-level.")
    ap.add_argument(
        "--corpus",
        choices=["text", "arithmetic"],
        default="text",
        help="'text' = byte-level text from --input; "
             "'arithmetic' = synthetic a+b=c examples (no --input needed).",
    )
    ap.add_argument("--arith-max-digits", type=int, default=3,
                    help="Max digits per operand for --corpus arithmetic.")
    ap.add_argument("--arith-reverse", action=argparse.BooleanOptionalAction, default=True,
                    help="Reverse-LSB answer (BP-canonical carry direction).")
    ap.add_argument("--arith-eval-problems", type=int, default=200,
                    help="Per-eval probe count for --corpus arithmetic.")
    ap.add_argument("--mode", choices=["baseline", "uniform", "zoned"], default="zoned")
    ap.add_argument("--seq-len", type=int, default=256)
    ap.add_argument("--batch-size", type=int, default=16)
    ap.add_argument("--steps", type=int, default=2000)
    ap.add_argument("--eval-every", type=int, default=100)
    ap.add_argument("--eval-batches", type=int, default=20)
    ap.add_argument("--lr", type=float, default=3e-4)
    ap.add_argument("--lr-schedule", choices=["constant", "cosine", "linear"], default="constant",
                    help="LR schedule: constant (default), cosine to lr-min over --steps, "
                         "or linear to lr-min over --steps. Warmup applies first if --warmup-steps > 0.")
    ap.add_argument("--lr-min", type=float, default=0.0,
                    help="Floor LR for cosine/linear schedules.")
    ap.add_argument("--warmup-steps", type=int, default=0,
                    help="Linear warmup from 0 to --lr over this many steps; 0 disables.")
    ap.add_argument("--weight-decay", type=float, default=0.01)
    ap.add_argument("--projected-weight-decay", type=float, default=0.0)
    ap.add_argument("--n-layer", type=int, default=8)
    ap.add_argument("--n-head", type=int, default=8)
    ap.add_argument("--n-embd", type=int, default=256)
    ap.add_argument("--mlp-mult", type=int, default=4)
    ap.add_argument("--mlp-activation",
                    choices=["gelu", "silu", "sigmoid", "linear", "silu_entmax_gate"],
                    default=None,
                    help="MLP activation. silu_entmax_gate = silu(x) * entmax15(x, dim=-1) — "
                         "BP-canonical sparse gating; entmax15 produces exact zeros on low-scoring "
                         "channels per token while silu carries the unbounded value on the rest.")
    ap.add_argument(
        "--loss",
        choices=["ce", "bce"],
        default=None,
        help="Per-token loss: 'ce' = softmax cross-entropy (default); "
             "'bce' = per-vocab Bernoulli (decoupled gradient, BP-aligned).",
    )
    ap.add_argument(
        "--logit-rmsnorm-scale",
        type=float,
        default=None,
        help="BCE only. Post-RMSNorm temperature scale on logits before BCE; 0 disables. "
             "Pure RMSNorm (scale=1) hits a structural BCE floor (~80 for V=256); use 8.0 to "
             "preserve sphere structure while giving σ enough dynamic range. Default 0 (off).",
    )
    ap.add_argument(
        "--sparsity-lambda",
        type=float,
        default=None,
        help="BCE only. Coefficient on penalty mean_{j!=t} σ(z_j); drives idle vocab logits "
             "to zero ('clean nop signal'). Try 1e-3.",
    )
    ap.add_argument(
        "--block-attractor-indices",
        type=str,
        default="",
        help="Comma-separated block indices to drive toward identity (e.g. '3' or '2,3'). "
             "Empty = no attractor.",
    )
    ap.add_argument(
        "--block-attractor-lambda",
        type=float,
        default=0.0,
        help="Coefficient on Σ_ℓ mean((block_ℓ(x) - x)²) for ℓ in --block-attractor-indices. "
             "Try 0.1 as a starting point.",
    )
    ap.add_argument(
        "--block-gates",
        action=argparse.BooleanOptionalAction,
        default=False,
        help="Add learnable per-block gates γ_ℓ initialized to 1, applied as "
             "x_{ℓ+1} = x_ℓ + γ_ℓ * (block_ℓ(x_ℓ) - x_ℓ). γ → 0 makes block identity.",
    )
    ap.add_argument(
        "--block-gate-kind",
        choices=["linear", "sigmoid", "ste"],
        default="linear",
        help="Gate parameterization. 'linear' = γ = α (init 1.0); 'sigmoid' = γ = σ(α) "
             "init α=4 (γ≈0.982), bistable under L1; 'ste' = forward γ ∈ {0,1} via threshold, "
             "backward via sigmoid (interpretable per-block on/off).",
    )
    ap.add_argument(
        "--block-gate-l1",
        type=float,
        default=0.0,
        help="L1 penalty on block gates: λ * Σ_ℓ |γ_ℓ|. Drives unneeded blocks toward identity. "
             "For sigmoid/ste γ ∈ (0,1) so penalty is bounded — try λ=0.05 (linear: 0.01).",
    )
    ap.add_argument(
        "--loop-block-indices",
        type=str,
        default="",
        help="Comma-separated block indices to apply K times during forward (e.g. '1' or '1,2'). "
             "Each designated block iterates --loop-block-k times. Tests recurrent BP in the "
             "designated block(s); other blocks remain feedforward. Empty = no looping.",
    )
    ap.add_argument(
        "--loop-block-k",
        type=int,
        default=1,
        help="K iterations for each block in --loop-block-indices. K=1 (default) = no looping. "
             "Backprop through K applications during training. Also the eval/inference K.",
    )
    ap.add_argument(
        "--loop-unit-indices",
        type=str,
        default="",
        help="Comma-separated CONTIGUOUS block indices applied K times as a UNIT (e.g. '1,2' "
             "applies block1→block2 K times together, vs --loop-block-indices '1,2' which "
             "applies block1 K times then block2 K times). Deeper loop body for BP iterations "
             "that need >1 block per step. K = --loop-block-k. Empty = no unit loop.",
    )
    ap.add_argument(
        "--loop-block-k-min",
        type=int,
        default=0,
        help="Minimum K for variable-K training: if > 0, each training batch samples K ~ "
             "Uniform({k_min..k_max}). Encourages the loop block to be stable across "
             "iteration counts (toward a BP fixed-point operator). 0 = disabled (use fixed "
             "--loop-block-k). Eval/inference still uses --loop-block-k.",
    )
    ap.add_argument(
        "--loop-block-k-max",
        type=int,
        default=0,
        help="Maximum K for variable-K training. Active only if --loop-block-k-min > 0.",
    )
    ap.add_argument(
        "--n-heads-per-block",
        type=str,
        default="",
        help="Comma-separated per-block n_head values (one per block, length must equal --n-layer). "
             "Each must divide --n-embd. Empty = uniform --n-head across all blocks. "
             "Example: '4,1,1,1' for OR-dominant block 0 with 4 heads, focused single-head elsewhere.",
    )
    ap.add_argument(
        "--mlp-mult-per-block",
        type=str,
        default="",
        help="Comma-separated per-block mlp_mult values (one per block, length must equal --n-layer). "
             "Empty = uniform --mlp-mult across all blocks. "
             "Example for 12L: '2,1,1,1,1,1,1,1,1,1,1,1' to keep block 0 full-width and shrink the rest.",
    )
    ap.add_argument(
        "--mlp-activation-per-block",
        type=str,
        default="",
        help="Comma-separated per-block MLP activation, choices {linear,sigmoid,silu,gelu}. "
             "Empty = uniform --mlp-activation. "
             "Per reverse-compiled adder (docs/coppola-bp-weight-compilation.md): early blocks want "
             "linear (assembly), middle blocks want sigmoid (conjunction-threshold) or silu, "
             "late blocks want sigmoid for commitment. "
             "Example for 12L: 'linear,linear,sigmoid,silu,silu,silu,silu,silu,silu,silu,sigmoid,silu'.",
    )
    ap.add_argument(
        "--mlp-output-rmsnorm-scale",
        type=float,
        default=0.0,
        help="If > 0, apply RMSNorm to each MLP's output before residual add, with this scale. "
             "Calibrates per-MLP delta to a fixed log-odds magnitude, mirroring "
             "--logit-rmsnorm-scale on the LM head side. Hypothesis: enables BP-canonical "
             "byte-identity-emergence (signature 5) on the MLP side. 0 = off.",
    )
    ap.add_argument(
        "--mlp-act-sparsity-lambda",
        type=float,
        default=0.0,
        help="L1 penalty on the post-activation inner MLP tensor (analogue of --sparsity-lambda "
             "but on the MLP intermediate, not vocab probabilities). Pushes the model toward "
             "per-token typed-channel firing — only a small fraction of inner neurons should "
             "activate per token. Typical values: 1e-3 to 1e-1. 0 = off.",
    )
    ap.add_argument(
        "--mlp-act-topk",
        type=int,
        default=0,
        help="Hard top-K mask on inner MLP activations per (batch, token) with "
             "straight-through estimator. Each token keeps only K of "
             "(mlp_mult * n_embd) inner neurons active; gradients flow through "
             "as if the mask weren't there. Tests whether MLPs implement typed "
             "BP channels when forced to structural per-token sparse firing. "
             "Typical: 8, 16, 32 out of 512. 0 = off.",
    )
    ap.add_argument(
        "--mlp-byte-concentration-lambda",
        type=float,
        default=0.0,
        help="Direct gradient pressure on byte-basis row-top-K mass of "
             "M = W_lm · W_proj · W_fc · W_wte^T. Loss subtracts λ · mean_top_k_fraction. "
             "Pushes MLP toward sparse-transition BP-readability. Use with "
             "--mlp-byte-concentration-layers. Typical: 0.1 to 1.0. 0 = off.",
    )
    ap.add_argument(
        "--mlp-byte-concentration-topk",
        type=int,
        default=4,
        help="K for top-K mass fraction in --mlp-byte-concentration-lambda. Default 4.",
    )
    ap.add_argument(
        "--mlp-byte-concentration-layers",
        type=str,
        default="",
        help="Comma-separated layer indices to apply MLP byte-concentration loss to. "
             "Empty = off. Example: '9,10,11' for late blocks only.",
    )
    ap.add_argument(
        "--attn-kind",
        choices=["mha", "tucker", "blockdiag", "shared_mixer_rotor"],
        default="mha",
        help="Attention factorization. 'mha' = standard multi-head softmax. 'tucker' = "
             "Tucker decomposition of R: r_p attention patterns × r_q OV kernels mixed by a "
             "[r_p,r_q] core. r_p=r_q=n_head with core=I matches MHA exactly. 'blockdiag' = "
             "block-diagonal in (a,b) feature axis: G = d_model/block_diag_size independent "
             "per-group attentions, no cross-group mixing in c_attn or c_proj. "
             "'shared_mixer_rotor' = Y = sum_s O_s · A · R_s · M · x with shared mixer M, "
             "shared attention A, cheap per-operator R_s (diag/rotor/full/none) and per-op "
             "down-project + lift. The friend's Run-3 architecture.",
    )
    ap.add_argument(
        "--tucker-r-p",
        type=int,
        default=4,
        help="Tucker only: number of attention patterns r_p. Must divide --n-embd.",
    )
    ap.add_argument(
        "--tucker-r-q",
        type=int,
        default=4,
        help="Tucker only: number of OV kernels r_q. Must divide --n-embd.",
    )
    ap.add_argument(
        "--tucker-tied-v",
        action=argparse.BooleanOptionalAction,
        default=False,
        help="Tucker only: tie V projection across all r_q operators (rank-vs-diversity "
             "Run-2 control). Degenerate at r_p=1; meaningful at r_p≥2.",
    )
    ap.add_argument(
        "--tucker-tied-o",
        action=argparse.BooleanOptionalAction,
        default=False,
        help="Tucker only: tie output projection (single d_v→d shared across r_q operators, "
             "sum operator outputs). Tests output-side multiplicity.",
    )
    ap.add_argument(
        "--tucker-d-v",
        type=int,
        default=0,
        help="Tucker only: explicit per-operator value dim. 0 = auto (d / r_q). "
             "Setting > 0 decouples d_v from r_q so reader count can scale without "
             "shrinking the writer (tied-O W_O size = d_v · d).",
    )
    ap.add_argument(
        "--block-diag-size",
        type=int,
        default=4,
        help="blockdiag only: per-group block size b. Must divide --n-embd. "
             "Default 4 matches RotorQuant findings.",
    )
    ap.add_argument(
        "--block-diag-proj",
        action=argparse.BooleanOptionalAction,
        default=True,
        help="blockdiag only: if True (default), output projection is also block-diagonal "
             "(no cross-group mixing). If False, output projection is full d×d, restoring "
             "cross-group mixing while keeping input-side c_attn block-diagonal.",
    )
    ap.add_argument(
        "--smr-n-operators",
        type=int,
        default=4,
        help="shared_mixer_rotor only: number of parallel operators (parallel feature "
             "transforms). Each gets its own R_s, W_V_s, W_O_s.",
    )
    ap.add_argument(
        "--smr-d-v",
        type=int,
        default=32,
        help="shared_mixer_rotor only: per-operator value dim. Must divide --n-embd.",
    )
    ap.add_argument(
        "--smr-r-kind",
        choices=["diag", "rotor", "full", "none"],
        default="rotor",
        help="shared_mixer_rotor only: per-operator R_s specialization. 'diag' = diagonal "
             "scaling, 'rotor' = block-diagonal SO(b) rotation, 'full' = dense d×d (control), "
             "'none' = identity (only W_V_s and W_O_s differ across operators).",
    )
    ap.add_argument(
        "--smr-rotor-block",
        type=int,
        default=4,
        help="shared_mixer_rotor only: when --smr-r-kind=rotor, block size for SO(b) "
             "rotor. Must divide --n-embd.",
    )
    ap.add_argument(
        "--smr-operator-kinds",
        type=str,
        default="",
        help="shared_mixer_rotor only: heterogeneous per-operator R kinds (Run 4). "
             "Comma-separated list, one entry per operator (overrides --smr-n-operators "
             "and --smr-r-kind). Each ∈ {diag, rotor, full, none}. Examples: "
             "'full,diag,diag,diag' for one rich + cheap tail; "
             "'full,rotor,rotor,rotor' for one rich + rotor tail.",
    )
    ap.add_argument(
        "--attn-norm",
        choices=["softmax", "linear_gain", "rms_signed", "kernel_relu2", "kernel_elu_plus1",
                 "sparsemax", "entmax15", "topk_ste"],
        default="softmax",
        help="Attention score normalization (applies to tucker / shared_mixer_rotor). "
             "softmax (default, BP-foreign convex combination), "
             "linear_gain (signed unbounded, most BP-canonical), "
             "rms_signed (signed but bounded), "
             "kernel_relu2/kernel_elu_plus1 (positive but unbounded), "
             "sparsemax/entmax15 (sparse exact-zero attention via simplex projection), "
             "topk_ste (forward = hard top-k softmax, backward = soft via STE).",
    )
    ap.add_argument(
        "--topk-k",
        type=int,
        default=4,
        help="topk_ste only: keys per query at inference. Default 4. "
             "Smaller = sparser attention pattern.",
    )
    ap.add_argument(
        "--res-attn-temp", type=float, default=1.0,
        help="A2′: fixed scalar τ multiplying the depth scores w_lᵀk before "
             "Φ (res-attn only). τ>1 gives entmax15 the score range it needs "
             "to sparsify; the depth analogue of --logit-rmsnorm-scale.")
    ap.add_argument(
        "--res-attn-entropy-lambda", type=float, default=0.0,
        help="A2′: coefficient on +mean H(α_depth) added to the loss "
             "(res-attn only). Minimised ⇒ peaked/sparse depth attention; "
             "gives the otherwise-idle L=4 router a reason to specialise.")
    ap.add_argument(
        "--position-encoding",
        choices=["learned", "alibi", "rope", "none"],
        default="learned",
        help="'learned' (default) = absolute learned wpe (BP-incompatible). "
             "'alibi' = no wpe; per-head linear-bias position encoding (Press et al. "
             "2021), aligning with BP framing where position enters as additive log-odds. "
             "'rope' = rotary position encoding (Su et al. 2021); relative positions via "
             "rotation of Q/K, no learned position parameters.",
    )
    ap.add_argument("--rope-base", type=float, default=10000.0,
                    help="RoPE base (default 10000). For short byte context, try 1000 or 100 "
                         "so slow-frequency dims actually rotate noticeably across the window.")
    ap.add_argument("--alibi-slope-scale", type=float, default=1.0,
                    help="Multiplier on standard ALiBi slopes (default 1.0). Lower values "
                         "(e.g. 0.5) make positional bias gentler — useful when r_p is small.")
    ap.add_argument(
        "--res-attn",
        choices=["none", "full"],
        default="none",
        help="Attention Residuals (Kimi, arXiv:2603.15031). 'none' (default) = "
             "standard additive residual. 'full' = block-granular Full AttnRes: "
             "each block input is a learned softmax/entmax aggregation over ALL "
             "preceding block outputs. See docs/attnres-bp-research-program.md.")
    ap.add_argument(
        "--res-attn-norm",
        choices=["softmax", "entmax15", "sparsemax"],
        default="softmax",
        help="Normalization Φ for the depth-attention weights (res-attn only). "
             "'softmax' = paper default (dense). 'entmax15'/'sparsemax' = "
             "exact-sparse: the learned depth-mixing matrix becomes a sparse "
             "BP factor graph (this repo's BP-structure extension).")
    ap.add_argument(
        "--smr-shared-wo",
        action=argparse.BooleanOptionalAction,
        default=False,
        help="shared_mixer_rotor only: 'many readers, one writer' compressed form. "
             "Single (d_v, d) W_O shared across operators (sum operator outputs in d_v "
             "space then one shared lift). Empirically motivated by Tucker tied-O win.",
    )
    ap.add_argument("--dropout", type=float, default=0.0)
    ap.add_argument("--val-fraction", type=float, default=0.02)
    ap.add_argument("--parquet-mix-weight", type=float, default=0.5)
    ap.add_argument("--text-mix-weight", type=float, default=0.5)
    ap.add_argument("--basis-update-interval", type=int, default=250)
    ap.add_argument("--uniform-routing-scale", type=float, default=1.0)
    ap.add_argument("--uniform-and-scale", type=float, default=1.0)
    ap.add_argument("--uniform-or-scale", type=float, default=1.0)
    ap.add_argument("--uniform-remainder-scale", type=float, default=0.0)
    ap.add_argument("--uniform-coupling-scale", type=float, default=1.0)
    ap.add_argument("--rank-energy", type=float, default=0.99)
    ap.add_argument("--max-rank", type=int, default=0)
    ap.add_argument("--seed", type=int, default=0)
    ap.add_argument("--sample-prompt", default="The ")
    ap.add_argument("--sample-tokens", type=int, default=80)
    ap.add_argument("--probe-jsonl", nargs="*", type=Path, default=[])
    ap.add_argument("--skip-update-support-projection", action="store_true")
    ap.add_argument("--save-every", type=int, default=0)
    ap.add_argument("--checkpoint-dir", type=Path, default=ROOT / "results" / "coppola_pretrain_tiny_ckpts")
    ap.add_argument("--resume", type=Path, default=None)
    ap.add_argument("-o", "--output", type=Path, default=ROOT / "results" / "coppola_pretrain_tiny.json")
    args = ap.parse_args()

    set_seed(args.seed)
    if args.corpus == "arithmetic":
        train_corpus = ArithmeticByteCorpus(max_digits=args.arith_max_digits, reverse=args.arith_reverse)
        val_corpus = ArithmeticByteCorpus(max_digits=args.arith_max_digits, reverse=args.arith_reverse)
    else:
        input_paths: List[Path] = []
        for path in args.input:
            if path.is_dir():
                input_paths.extend(sorted(p for p in path.iterdir() if p.is_file()))
            else:
                input_paths.append(path)
        tokenizer = None
        if args.tokenizer_path is not None:
            from tokenizers import Tokenizer
            tokenizer = Tokenizer.from_file(str(args.tokenizer_path))
            print(f"loaded tokenizer ({tokenizer.get_vocab_size()} vocab) from {args.tokenizer_path}")
        train_corpus, val_corpus = build_corpora(
            input_paths,
            text_key=args.text_key,
            val_fraction=args.val_fraction,
            parquet_mix_weight=args.parquet_mix_weight,
            text_mix_weight=args.text_mix_weight,
            tokenizer=tokenizer,
        )
        if tokenizer is not None:
            args.vocab_size = tokenizer.get_vocab_size()
        else:
            args.vocab_size = 256
    probe_groups = load_probe_rows(args.probe_jsonl)

    device = "cuda" if torch.cuda.is_available() else "cpu"
    mlp_activation = infer_mlp_activation(args, device=device)
    args.mlp_activation = mlp_activation
    loss_kind = infer_loss_kind(args, device=device)
    args.loss = loss_kind
    logit_rmsnorm_scale, sparsity_lambda = infer_bce_reg(args, device=device)
    args.logit_rmsnorm_scale = logit_rmsnorm_scale
    args.sparsity_lambda = sparsity_lambda
    if loss_kind != "bce" and (logit_rmsnorm_scale > 0 or sparsity_lambda > 0):
        print("warning: --logit-rmsnorm-scale / --sparsity-lambda are BCE-only; ignoring under loss=ce")

    block_attractor_indices: tuple[int, ...] = (
        tuple(int(x.strip()) for x in args.block_attractor_indices.split(",") if x.strip())
        if args.block_attractor_indices
        else ()
    )
    loop_block_indices: tuple[int, ...] = (
        tuple(int(x.strip()) for x in args.loop_block_indices.split(",") if x.strip())
        if args.loop_block_indices
        else ()
    )
    loop_unit_indices: tuple[int, ...] = (
        tuple(int(x.strip()) for x in args.loop_unit_indices.split(",") if x.strip())
        if args.loop_unit_indices
        else ()
    )
    n_heads_per_block: tuple[int, ...] = (
        tuple(int(x.strip()) for x in args.n_heads_per_block.split(",") if x.strip())
        if args.n_heads_per_block
        else ()
    )
    mlp_mult_per_block: tuple[int, ...] = (
        tuple(int(x.strip()) for x in args.mlp_mult_per_block.split(",") if x.strip())
        if args.mlp_mult_per_block
        else ()
    )
    mlp_activation_per_block: tuple[str, ...] = (
        tuple(s.strip() for s in args.mlp_activation_per_block.split(",") if s.strip())
        if args.mlp_activation_per_block
        else ()
    )
    smr_operator_kinds: tuple[str, ...] = (
        tuple(s.strip() for s in args.smr_operator_kinds.split(",") if s.strip())
        if args.smr_operator_kinds
        else ()
    )

    cfg = GPTConfig(
        vocab_size=getattr(args, "vocab_size", 256),
        seq_len=args.seq_len,
        n_layer=args.n_layer,
        n_head=args.n_head,
        n_embd=args.n_embd,
        mlp_mult=args.mlp_mult,
        mlp_activation=mlp_activation,
        dropout=args.dropout,
        loss_kind=loss_kind,
        logit_rmsnorm_scale=logit_rmsnorm_scale,
        sparsity_lambda=sparsity_lambda,
        block_attractor_indices=block_attractor_indices,
        block_attractor_lambda=float(args.block_attractor_lambda),
        block_gates=bool(args.block_gates),
        block_gate_kind=str(args.block_gate_kind),
        block_gate_l1=float(args.block_gate_l1),
        loop_block_indices=loop_block_indices,
        loop_block_k=int(args.loop_block_k),
        loop_unit_indices=loop_unit_indices,
        n_heads_per_block=n_heads_per_block,
        mlp_mult_per_block=mlp_mult_per_block,
        mlp_activation_per_block=mlp_activation_per_block,
        mlp_output_rmsnorm_scale=float(args.mlp_output_rmsnorm_scale),
        mlp_act_sparsity_lambda=float(args.mlp_act_sparsity_lambda),
        mlp_act_topk=int(args.mlp_act_topk),
        mlp_byte_concentration_lambda=float(args.mlp_byte_concentration_lambda),
        mlp_byte_concentration_topk=int(args.mlp_byte_concentration_topk),
        mlp_byte_concentration_layers=(
            tuple(int(x.strip()) for x in args.mlp_byte_concentration_layers.split(",") if x.strip())
            if args.mlp_byte_concentration_layers else ()
        ),
        attn_kind=str(args.attn_kind),
        tucker_r_p=int(args.tucker_r_p),
        tucker_r_q=int(args.tucker_r_q),
        tucker_tied_v=bool(args.tucker_tied_v),
        tucker_tied_o=bool(args.tucker_tied_o),
        tucker_d_v=int(args.tucker_d_v),
        block_diag_size=int(args.block_diag_size),
        block_diag_proj=bool(args.block_diag_proj),
        smr_n_operators=int(args.smr_n_operators),
        smr_d_v=int(args.smr_d_v),
        smr_r_kind=str(args.smr_r_kind),
        smr_rotor_block=int(args.smr_rotor_block),
        smr_operator_kinds=smr_operator_kinds,
        attn_norm=str(args.attn_norm),
        smr_shared_wo=bool(args.smr_shared_wo),
        topk_k=int(args.topk_k),
        position_encoding=str(args.position_encoding),
        rope_base=float(args.rope_base),
        alibi_slope_scale=float(args.alibi_slope_scale),
        res_attn=str(args.res_attn),
        res_attn_norm=str(args.res_attn_norm),
        res_attn_temp=float(args.res_attn_temp),
        res_attn_entropy_lambda=float(args.res_attn_entropy_lambda),
    )
    model = ByteGPT(cfg).to(device)

    controller = None
    if args.mode != "baseline":
        rank_policy = RankPolicy(energy=args.rank_energy, max_rank=(args.max_rank or None))
        controller = CoppolaPretrainingController(
            CoppolaPretrainingConfig(
                n_head=args.n_head,
                attn_output_policy=rank_policy,
                mlp_output_policy=rank_policy,
                mlp_input_policy=rank_policy,
                basis_update_interval=args.basis_update_interval,
                uniform_scales=FamilyScales(
                    routing=args.uniform_routing_scale,
                    and_transport=args.uniform_and_scale,
                    or_update=args.uniform_or_scale,
                    remainder=args.uniform_remainder_scale,
                    coupling=args.uniform_coupling_scale,
                ),
            )
        )
        controller.refresh_from_model(model)

    optimizer = build_optimizer(
        model,
        lr=args.lr,
        weight_decay=args.weight_decay,
        projected_weight_decay=args.projected_weight_decay if controller is not None else args.weight_decay,
    )

    history: List[Dict] = []
    start_step = 1
    if args.resume is not None:
        payload = load_checkpoint(args.resume, model, optimizer, device=device)
        history = payload.get("history", [])
        start_step = int(payload["step"]) + 1
        if controller is not None:
            controller.refresh_from_model(model)

    last_projection_stats: Dict[str, Dict[str, float]] | None = None

    def _current_lr(step: int) -> float:
        lr_max = args.lr
        lr_min = args.lr_min
        if args.warmup_steps > 0 and step <= args.warmup_steps:
            return lr_max * step / max(1, args.warmup_steps)
        if args.lr_schedule == "constant":
            return lr_max
        # progress over the post-warmup phase, in [0, 1]
        decay_start = args.warmup_steps
        decay_total = max(1, args.steps - decay_start)
        progress = min(1.0, max(0.0, (step - decay_start) / decay_total))
        if args.lr_schedule == "cosine":
            return lr_min + 0.5 * (lr_max - lr_min) * (1.0 + math.cos(math.pi * progress))
        if args.lr_schedule == "linear":
            return lr_max - (lr_max - lr_min) * progress
        return lr_max

    for step in range(start_step, args.steps + 1):
        model.train()
        if controller is not None and step % args.basis_update_interval == 1 and step > 1:
            controller.refresh_from_model(model)
        cur_lr = _current_lr(step)
        for pg in optimizer.param_groups:
            pg["lr"] = cur_lr
        x, y = train_corpus.sample_batch(batch_size=args.batch_size, seq_len=args.seq_len, device=device)
        optimizer.zero_grad(set_to_none=True)
        # Variable-K training: sample K per batch if --loop-block-k-min > 0.
        k_override = None
        if args.loop_block_k_min > 0 and args.loop_block_k_max >= args.loop_block_k_min:
            k_override = random.randint(args.loop_block_k_min, args.loop_block_k_max)
        _, loss = model(x, y, loop_block_k_override=k_override)
        loss.backward()
        param_snapshots = None
        if controller is not None:
            last_projection_stats = apply_coppola_projection(model, controller, mode=args.mode)
            if not args.skip_update_support_projection:
                param_snapshots = snapshot_projected_params(model)
        optimizer.step()
        if controller is not None and param_snapshots is not None:
            project_parameter_updates(param_snapshots, controller)

        if step == 1 or step % args.eval_every == 0 or step == args.steps:
            metrics = evaluate(
                model,
                val_corpus,
                batch_size=args.batch_size,
                seq_len=args.seq_len,
                device=device,
                eval_batches=args.eval_batches,
            )
            sample = format_sample(model, device=device, prompt=args.sample_prompt, max_new_tokens=args.sample_tokens)
            payload = {
                "step": step,
                "train_loss": float(loss.item()),
                "val_loss": metrics["loss"],
                "val_ce": metrics["ce_loss"],
                "val_bpb": metrics["bpb"],
                "sample": sample,
            }
            if last_projection_stats is not None:
                payload["projection_stats"] = last_projection_stats
            probe_payload = evaluate_probe_groups(model, probe_groups, device=device)
            if probe_payload:
                payload["probe_eval"] = probe_payload
            if args.corpus == "arithmetic":
                payload["arith"] = evaluate_arithmetic_accuracy(
                    model,
                    device=device,
                    n_problems=args.arith_eval_problems,
                    max_digits=args.arith_max_digits,
                    reverse=args.arith_reverse,
                )
            history.append(payload)
            line = (
                f"step={step:5d} train_loss={loss.item():.4f} "
                f"val_loss={metrics['loss']:.4f} val_ce={metrics['ce_loss']:.4f} "
                f"val_bpb={metrics['bpb']:.4f}"
            )
            if "arith" in payload:
                line += f" arith_acc={payload['arith']['accuracy']:.3f}"
            if model.block_gate is not None:
                with torch.no_grad():
                    gate_vals = model._gate_values()
                gates = gate_vals.detach().cpu().tolist() if gate_vals is not None else []
                line += " γ=[" + " ".join(f"{g:+.3f}" for g in gates) + "]"
                payload["block_gates"] = gates
            print(line, flush=True)
        if args.save_every > 0 and (step % args.save_every == 0 or step == args.steps):
            ckpt_path = args.checkpoint_dir / f"{args.mode}_step_{step}.pt"
            save_checkpoint(
                ckpt_path,
                model=model,
                optimizer=optimizer,
                step=step,
                history=history,
                mode=args.mode,
                args=args,
                model_config=cfg,
            )
            save_checkpoint(
                args.checkpoint_dir / f"{args.mode}_latest.pt",
                model=model,
                optimizer=optimizer,
                step=step,
                history=history,
                mode=args.mode,
                args=args,
                model_config=cfg,
            )

    args.output.parent.mkdir(parents=True, exist_ok=True)
    output = {
        "mode": args.mode,
        "steps": args.steps,
        "seed": args.seed,
        "config": _json_safe(vars(args)),
        "model_config": asdict(cfg),
        "train_corpus": _json_safe(train_corpus.stats()),
        "val_corpus": _json_safe(val_corpus.stats()),
        "history": history,
    }
    if history:
        output["final_val_bpb"] = history[-1].get("val_bpb")
        output["best_val_bpb"] = min(
            row["val_bpb"] for row in history if row.get("val_bpb") is not None
        )
    args.output.write_text(json.dumps(output, indent=2))
    print(f"saved {args.output}")


if __name__ == "__main__":
    main()