modeling_suprabrain.py

"""
© SupraLabs 2026 - Official pretraining code for SupraBrain 50M v0.1
Hybrid: Gated DeltaNet (3:1) + Sliding-Window Attention + Surprise Gating
Optimizer: Muon (2D) + AdamW (rest) | Schedule: WSD | Data: FineWeb-Edu -> Anneal
"""
import os
os.environ.setdefault("PYTORCH_ALLOC_CONF", "expandable_segments:True")
os.environ.setdefault("CUDA_VISIBLE_DEVICES", "0")
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")

import json
import math
import shutil
from dataclasses import dataclass
from typing import List, Optional, Tuple

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import Dataset, SequentialSampler

from transformers import PretrainedConfig, PreTrainedModel
from transformers.modeling_outputs import CausalLMOutputWithPast
from transformers import GenerationMixin

# ---------------------------------------------------------------------------
# Optional backends
# ---------------------------------------------------------------------------
try:
    from fla.ops.gated_delta_rule import chunk_gated_delta_rule as _fla_gdn
    _HAS_FLA = True
except Exception:
    _fla_gdn = None
    _HAS_FLA = False

try:
    from torch.nn.attention.flex_attention import (
        flex_attention as _flex_raw, create_block_mask as _create_block_mask)
    _HAS_FLEX = True
except Exception:
    _flex_raw, _create_block_mask = None, None
    _HAS_FLEX = False

_flex_attention = None  # lazy compiled


def _get_flex():
    global _flex_attention
    if _flex_attention is None:
        _flex_attention = torch.compile(_flex_raw, dynamic=False)
    return _flex_attention


# ===========================================================================
#  1. CONFIG
# ===========================================================================
class SupraBrainConfig(PretrainedConfig):
    model_type = "suprabrain"

    def __init__(
        self,
        vocab_size: int = 23_808,
        hidden_size: int = 384,
        num_hidden_layers: int = 28,
        intermediate_size: int = 1152,          # 3 * d  (2-Matrix-MLP)
        mlp_act: str = "rational",              # "rational" | "relu2" | "swiglu"
        # --- Gated DeltaNet ---
        gdn_num_heads: int = 3,                 # 3 x 128 statt 6 x 64 -> 2x Recall
        gdn_head_dim: int = 128,
        gdn_conv_size: int = 4,
        gdn_gate_rank: int = 32,                # Low-Rank Output-Gate
        # --- Surprise Gating ---
        surprise: bool = True,
        surprise_win: int = 16,
        # --- Attention ---
        attn_num_heads: int = 6,
        attn_num_kv_heads: int = 2,
        attn_head_dim: int = 64,
        attn_window: int = 512,
        attn_every: int = 4,                    # 3 GDN : 1 Attn
        full_attn_layers: Tuple[int, ...] = (19,),
        rope_theta: float = 10_000.0,
        # --- Head / Loss ---
        unembed_rank: int = 32,                 # half-untied adapter
        z_loss: float = 1e-4,
        logit_softcap: float = 0.0,
        ce_chunk: int = 1024,                   # checkpointed chunked CE
        # --- Misc ---
        max_position_embeddings: int = 1024,
        tie_word_embeddings: bool = True,
        emb_shortcut: bool = True,
        value_residual: bool = True,
        norm_eps: float = 1e-5,
        initializer_range: float = 0.02,
        grad_ckpt: bool = False,
        compile_blocks: bool = True,
        pad_token_id: int = 1,
        bos_token_id: int = 0,
        eos_token_id: int = 2,
        **kw,
    ):
        self.vocab_size = vocab_size
        self.hidden_size = hidden_size
        self.num_hidden_layers = num_hidden_layers
        self.intermediate_size = intermediate_size
        self.mlp_act = mlp_act
        self.gdn_num_heads = gdn_num_heads
        self.gdn_head_dim = gdn_head_dim
        self.gdn_conv_size = gdn_conv_size
        self.gdn_gate_rank = gdn_gate_rank
        self.surprise = surprise
        self.surprise_win = surprise_win
        self.attn_num_heads = attn_num_heads
        self.attn_num_kv_heads = attn_num_kv_heads
        self.attn_head_dim = attn_head_dim
        self.attn_window = attn_window
        self.attn_every = attn_every
        self.full_attn_layers = tuple(full_attn_layers)
        self.rope_theta = rope_theta
        self.unembed_rank = unembed_rank
        self.z_loss = z_loss
        self.logit_softcap = logit_softcap
        self.ce_chunk = ce_chunk
        self.max_position_embeddings = max_position_embeddings
        self.emb_shortcut = emb_shortcut
        self.value_residual = value_residual
        self.norm_eps = norm_eps
        self.initializer_range = initializer_range
        self.grad_ckpt = grad_ckpt
        self.compile_blocks = compile_blocks
        super().__init__(pad_token_id=pad_token_id, bos_token_id=bos_token_id,
                         eos_token_id=eos_token_id,
                         tie_word_embeddings=tie_word_embeddings, **kw)

    @property
    def attn_layers(self) -> List[int]:
        """Layer-Indices with attention. Never in Layer 0/1 (GDN suffices there)."""
        base = [i for i in range(self.num_hidden_layers)
                if i % self.attn_every == (self.attn_every - 1) and i >= 2]
        return sorted(set(base) | set(self.full_attn_layers))


# ===========================================================================
#  2. BUILDING BLOCKS
# ===========================================================================
class RMSNorm(nn.Module):
    def __init__(self, d, eps=1e-5):
        super().__init__()
        self.w = nn.Parameter(torch.ones(d))
        self.eps = eps

    def forward(self, x):
        dt = x.dtype
        x = x.float()
        x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
        return (x * self.w.float()).to(dt)


class GatedRMSNorm(nn.Module):
    """RMSNorm(x) * silu(gate) -- the usual GDN/Mamba output standard."""
    def __init__(self, d, eps=1e-5):
        super().__init__()
        self.w = nn.Parameter(torch.ones(d))
        self.eps = eps

    def forward(self, x, gate):
        dt = x.dtype
        x = x.float()
        x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
        return (x * self.w.float() * F.silu(gate.float())).to(dt)


class RationalAct(nn.Module):
    """Per-channel learnable nonlinearity (the cheap KAN substitute).
       f(x) = (a0 + a1 x + a2 x^2) / (1 + |b1 x|)   -> 4 params per channel."""
    def __init__(self, d):
        super().__init__()
        self.a0 = nn.Parameter(torch.zeros(d))
        self.a1 = nn.Parameter(torch.ones(d))
        self.a2 = nn.Parameter(torch.full((d,), 0.5))
        self.b1 = nn.Parameter(torch.full((d,), 0.5))

    def forward(self, x):
        num = self.a0 + self.a1 * x + self.a2 * x * x
        den = 1.0 + (self.b1 * x).abs()
        return num / den


class MLP(nn.Module):
    def __init__(self, cfg: SupraBrainConfig):
        super().__init__()
        d, h, act = cfg.hidden_size, cfg.intermediate_size, cfg.mlp_act
        self.act_type = act
        if act == "swiglu":
            self.gate = nn.Linear(d, h, bias=False)
            self.up = nn.Linear(d, h, bias=False)
        else:
            self.up = nn.Linear(d, h, bias=False)
            self.act = RationalAct(h) if act == "rational" else None
        self.down = nn.Linear(h, d, bias=False)

    def forward(self, x):
        if self.act_type == "swiglu":
            return self.down(F.silu(self.gate(x)) * self.up(x))
        h = self.up(x)
        h = self.act(h) if self.act is not None else F.relu(h).pow(2)
        return self.down(h)


class ShortConv(nn.Module):
    """Causal depthwise Conv1d + SiLU (Mamba/GDN standard)."""
    def __init__(self, d, k=4):
        super().__init__()
        self.k = k
        self.conv = nn.Conv1d(d, d, k, groups=d, bias=False)

    def forward(self, x):                       # (B,T,D)
        y = self.conv(F.pad(x.transpose(1, 2), (self.k - 1, 0)))
        return F.silu(y.transpose(1, 2))


class SurpriseBeta(nn.Module):
    """beta_t (B,T,H) in (0,2). Zero-init Surprise-Path -> Start == Vanilla-GDN.
       Fully parallel; no state dependency -> chunkwise scan remains intact."""
    def __init__(self, d, n_heads, win=16, enable=True):
        super().__init__()
        self.enable = enable
        self.win = win
        self.to_beta = nn.Linear(d, n_heads, bias=True)
        if enable:
            self.pred = nn.Conv1d(d, d, 3, groups=d, bias=False)   # causal, k=3
            self.w_s = nn.Parameter(torch.zeros(n_heads))          # ZERO-INIT

    def forward(self, x):                       # x: (B,T,D)
        logits = self.to_beta(x).float()
        if self.enable:
            xt = x.transpose(1, 2)
            xhat = self.pred(F.pad(xt, (2, 0)))
            xt32 = xt.float()
            num = (xt32 - xhat.float()).pow(2).sum(1, keepdim=True)
            den = xt32.pow(2).sum(1, keepdim=True) + 1e-6
            s = num / den                                      # scale-invariant!
            # causal box filter = "past surprise" (avg_pool1d remains fp32)
            s = F.avg_pool1d(F.pad(s, (self.win - 1, 0)), self.win, 1)
            s = torch.log1p(s).transpose(1, 2)                 # (B,T,1)
            logits = logits + self.w_s.float() * s
        return 2.0 * torch.sigmoid(logits)                     # beta in (0,2)


# ---------------------------------------------------------------------------
#  Gated DeltaNet Kernel-Adapter
# ---------------------------------------------------------------------------
def _gdn_reference(q, k, v, g, beta, use_l2=True):
    """Slow but guaranteed-correct reference (CPU / no FLA / unit test)."""
    B, T, H, Dk = q.shape
    Dv = v.shape[-1]
    q, k, v = q.float(), k.float(), v.float()
    if use_l2:
        q = F.normalize(q, dim=-1)
        k = F.normalize(k, dim=-1)
    S = q.new_zeros(B, H, Dk, Dv)
    out = q.new_zeros(B, T, H, Dv)
    scale = Dk ** -0.5
    for t in range(T):
        gt = g[:, t].float().exp().unsqueeze(-1).unsqueeze(-1)      # (B,H,1,1)
        S = S * gt
        kt, vt, bt = k[:, t], v[:, t], beta[:, t].float().unsqueeze(-1)
        pred = torch.einsum('bhk,bhkv->bhv', kt, S)
        S = S + bt.unsqueeze(-1) * kt.unsqueeze(-1) * (vt - pred).unsqueeze(-2)
        out[:, t] = torch.einsum('bhk,bhkv->bhv', q[:, t], S) * scale
    return out.to(v.dtype)


@torch._dynamo.disable()
def gdn_op(q, k, v, g, beta):
    """q,k,v: (B,T,H,D) bf16 | g: (B,T,H) fp32 LOG-space | beta: (B,T,H) fp32."""
    if not _HAS_FLA or not q.is_cuda:
        return _gdn_reference(q, k, v, g, beta, use_l2=True)
    try:
        out = _fla_gdn(q=q, k=k, v=v, g=g, beta=beta,
                       use_qk_l2norm_in_kernel=True, output_final_state=False)
    except TypeError:                       # older FLA-Version
        q = F.normalize(q.float(), dim=-1).to(v.dtype)
        k = F.normalize(k.float(), dim=-1).to(v.dtype)
        out = _fla_gdn(q=q, k=k, v=v, g=g, beta=beta, output_final_state=False)
    return out[0] if isinstance(out, (tuple, list)) else out


class GatedDeltaNet(nn.Module):
    def __init__(self, cfg: SupraBrainConfig):
        super().__init__()
        d = cfg.hidden_size
        self.H, self.D = cfg.gdn_num_heads, cfg.gdn_head_dim
        inner = self.H * self.D
        assert inner == d, "gdn_num_heads * gdn_head_dim must be == hidden_size"

        self.q_proj = nn.Linear(d, inner, bias=False)
        self.k_proj = nn.Linear(d, inner, bias=False)
        self.v_proj = nn.Linear(d, inner, bias=False)
        self.cq = ShortConv(inner, cfg.gdn_conv_size)
        self.ck = ShortConv(inner, cfg.gdn_conv_size)
        self.cv = ShortConv(inner, cfg.gdn_conv_size)

        # Data-dependent Decay (enzymatic Clearance), per Head, LOG-space
        self.a_proj = nn.Linear(d, self.H, bias=True)
        self.A_log = nn.Parameter(torch.log(torch.empty(self.H).uniform_(1, 16)))
        self.dt_bias = nn.Parameter(torch.full((self.H,), math.log(math.expm1(0.02))))

        self.beta = SurpriseBeta(d, self.H, cfg.surprise_win, cfg.surprise)

        # Low-Rank Output-Gate: Saves ~120k/layer in contrast to d x d projection
        r = cfg.gdn_gate_rank
        self.g_down = nn.Linear(d, r, bias=False)
        self.g_up = nn.Linear(r, inner, bias=False)
        self.g_bias = nn.Parameter(torch.ones(inner))

        self.gnorm = GatedRMSNorm(self.D, cfg.norm_eps)
        self.o_proj = nn.Linear(inner, d, bias=False)

    def forward(self, x, doc_start=None):
        B, T, _ = x.shape
        q = self.cq(self.q_proj(x)).view(B, T, self.H, self.D)
        k = self.ck(self.k_proj(x)).view(B, T, self.H, self.D)
        v = self.cv(self.v_proj(x)).view(B, T, self.H, self.D)

        # --- Decay in fp32, Log-Space (Kernel-Convention) ---
        a = self.a_proj(x).float()
        g = -torch.exp(self.A_log.float()) * F.softplus(a + self.dt_bias.float())
        if doc_start is not None:
            # DOCUMENT STATE RESET: g -> ~0 deletes S_{t-1} at document begin.
            # No elementwise -> no dynamic shapes, compile-safe.
            g = torch.where(doc_start.unsqueeze(-1), g.new_full((), -25.0), g)

        beta = self.beta(x)                                   # (B,T,H) fp32

        o = gdn_op(q, k, v, g, beta)                          # (B,T,H,Dv)

        gate = (self.g_up(self.g_down(x)) + self.g_bias).view(B, T, self.H, self.D)
        o = self.gnorm(o, gate).reshape(B, T, -1)
        return self.o_proj(o)


class WindowAttention(nn.Module):
    def __init__(self, cfg: SupraBrainConfig, layer_idx: int):
        super().__init__()
        d = cfg.hidden_size
        self.Hq, self.Hk, self.Dh = (cfg.attn_num_heads, cfg.attn_num_kv_heads,
                                     cfg.attn_head_dim)
        self.rep = self.Hq // self.Hk
        self.is_full = layer_idx in cfg.full_attn_layers
        self.window = cfg.attn_window
        self.q_proj = nn.Linear(d, self.Hq * self.Dh, bias=False)
        self.k_proj = nn.Linear(d, self.Hk * self.Dh, bias=False)
        self.v_proj = nn.Linear(d, self.Hk * self.Dh, bias=False)
        self.o_proj = nn.Linear(self.Hq * self.Dh, d, bias=False)
        self.qn = RMSNorm(self.Dh, cfg.norm_eps)              # QK-Norm: Obligatory!
        self.kn = RMSNorm(self.Dh, cfg.norm_eps)
        self.use_vres = cfg.value_residual
        if self.use_vres:
            self.lam = nn.Parameter(torch.zeros(1))           # 0 -> saves = pure v

    def forward(self, x, rope, masks, v_first):
        B, T, _ = x.shape
        cos, sin = rope
        q = self.qn(self.q_proj(x).view(B, T, self.Hq, self.Dh))
        k = self.kn(self.k_proj(x).view(B, T, self.Hk, self.Dh))
        v = self.v_proj(x).view(B, T, self.Hk, self.Dh)

        # Value-Residual (modded-nanogpt): ~1 scalar, big loss-win!
        if self.use_vres:
            if v_first is None:
                v_first = v
            else:
                lam = torch.sigmoid(self.lam).to(v.dtype)
                v = (1 - lam) * v + lam * v_first

        q, k = apply_rope(q, cos, sin), apply_rope(k, cos, sin)
        q = q.transpose(1, 2)
        k = k.repeat_interleave(self.rep, dim=2).transpose(1, 2)
        vv = v.repeat_interleave(self.rep, dim=2).transpose(1, 2)

        bm = masks["full"] if self.is_full else masks["win"]
        if bm is not None:
            o = _get_flex()(q, k, vv, block_mask=bm)
        else:
            o = F.scaled_dot_product_attention(
                q, k, vv,
                attn_mask=None if self.is_full else masks["sdpa_win"],
                is_causal=self.is_full)
        o = o.transpose(1, 2).reshape(B, T, -1)
        return self.o_proj(o), v_first


def build_rope(T, dim, theta, device, dtype):
    inv = 1.0 / (theta ** (torch.arange(0, dim, 2, device=device).float() / dim))
    t = torch.arange(T, device=device).float()
    f = torch.outer(t, inv)
    return f.cos().to(dtype), f.sin().to(dtype)


def apply_rope(x, cos, sin):                    # x: (B,T,H,D)
    x1, x2 = x.float().chunk(2, dim=-1)
    c, s = cos[None, :, None, :], sin[None, :, None, :]
    return torch.cat([x1 * c - x2 * s, x1 * s + x2 * c], -1).to(x.dtype)


class Block(nn.Module):
    def __init__(self, cfg: SupraBrainConfig, i: int):
        super().__init__()
        self.is_attn = i in cfg.attn_layers
        self.norm1 = RMSNorm(cfg.hidden_size, cfg.norm_eps)
        self.mixer = WindowAttention(cfg, i) if self.is_attn else GatedDeltaNet(cfg)
        self.norm2 = RMSNorm(cfg.hidden_size, cfg.norm_eps)
        self.mlp = MLP(cfg)
        self.use_emb_sc = cfg.emb_shortcut
        if self.use_emb_sc:
            self.lam_emb = nn.Parameter(torch.zeros(1))       # zero-init

    def forward(self, x, x_emb, doc_start, rope, masks, v_first):
        h = self.norm1(x)
        if self.is_attn:
            y, v_first = self.mixer(h, rope, masks, v_first)
        else:
            y = self.mixer(h, doc_start)
        x = x + y
        x = x + self.mlp(self.norm2(x))
        if self.use_emb_sc:
            x = x + self.lam_emb.to(x.dtype) * x_emb           # U-Net-like skip
        return x, v_first


# ===========================================================================
#  3. MODEL
# ===========================================================================
class SupraBrainPreTrainedModel(PreTrainedModel):
    config_class = SupraBrainConfig
    base_model_prefix = "model"
    supports_gradient_checkpointing = False
    _no_split_modules = ["Block"]

    def _init_weights(self, m):
        std = self.config.initializer_range
        deep = std / math.sqrt(2 * self.config.num_hidden_layers)
        if isinstance(m, nn.Linear):
            nn.init.normal_(m.weight, 0.0, std)
            if m.bias is not None:
                nn.init.zeros_(m.bias)
        elif isinstance(m, nn.Embedding):
            nn.init.normal_(m.weight, 0.0, std)
        elif isinstance(m, (GatedDeltaNet, WindowAttention)):
            nn.init.normal_(m.o_proj.weight, 0.0, deep)        # Residual-Scaling
        elif isinstance(m, MLP):
            nn.init.normal_(m.down.weight, 0.0, deep)


class SupraBrainModel(SupraBrainPreTrainedModel):
    def __init__(self, cfg: SupraBrainConfig):
        super().__init__(cfg)
        self.embed_tokens = nn.Embedding(cfg.vocab_size, cfg.hidden_size)
        self.embed_norm = RMSNorm(cfg.hidden_size, cfg.norm_eps)
        self.layers = nn.ModuleList([Block(cfg, i)
                                     for i in range(cfg.num_hidden_layers)])
        self.norm = RMSNorm(cfg.hidden_size, cfg.norm_eps)
        self._rope_cache = None
        self._sdpa_mask = None
        self.post_init()

    def _rope(self, T, device, dtype):
        if self._rope_cache is None or self._rope_cache[0].shape[0] < T:
            self._rope_cache = build_rope(max(T, self.config.max_position_embeddings),
                                          self.config.attn_head_dim,
                                          self.config.rope_theta, device, torch.float32)
        c, s = self._rope_cache
        return c[:T], s[:T]

    def _masks(self, doc_ids, T, device):
        cfg = self.config
        if not _HAS_FLEX or not doc_ids.is_cuda:
            if self._sdpa_mask is None or self._sdpa_mask.shape[-1] < T:
                i = torch.arange(T, device=device)
                d = i[:, None] - i[None, :]
                self._sdpa_mask = ((d >= 0) & (d < cfg.attn_window))[None, None]
            return {"win": None, "full": None, "sdpa_win": self._sdpa_mask[..., :T, :T]}

        W = cfg.attn_window

        def win_mod(b, h, qi, ki):
            return (qi >= ki) & ((qi - ki) < W) & (doc_ids[b, qi] == doc_ids[b, ki])

        def full_mod(b, h, qi, ki):
            return (qi >= ki) & (doc_ids[b, qi] == doc_ids[b, ki])

        B = doc_ids.shape[0]
        mk = lambda f: _create_block_mask(f, B, None, T, T, device=device,
                                          BLOCK_SIZE=128, _compile=True)
        return {"win": mk(win_mod), "full": mk(full_mod), "sdpa_win": None}

    def forward(self, input_ids):
        cfg = self.config
        B, T = input_ids.shape
        dev = input_ids.device

        # ---- Derive document boundaries from EOS (costs 0 parameters)) ----
        is_eos = (input_ids == cfg.eos_token_id)
        prev_eos = F.pad(is_eos[:, :-1], (1, 0), value=True)     # pos 0 = doc start
        doc_start = prev_eos
        doc_ids = torch.cumsum(prev_eos.int(), dim=1)

        x_emb = self.embed_norm(self.embed_tokens(input_ids))
        x = x_emb
        rope = self._rope(T, dev, torch.float32)
        masks = self._masks(doc_ids, T, dev)
        v_first = None

        for blk in self.layers:
            if cfg.grad_ckpt and self.training:
                x, v_first = torch.utils.checkpoint.checkpoint(
                    blk, x, x_emb, doc_start, rope, masks, v_first,
                    use_reentrant=False)
            else:
                x, v_first = blk(x, x_emb, doc_start, rope, masks, v_first)
        return self.norm(x)


def _ce_chunk(h, W, labels, z_coef, softcap):
    logits = (h @ W.t()).float()
    if softcap > 0:
        logits = softcap * torch.tanh(logits / softcap)
    loss = F.cross_entropy(logits, labels, reduction="sum")
    if z_coef > 0:
        loss = loss + z_coef * torch.logsumexp(logits, -1).pow(2).sum()
    return loss


class SupraBrainForCausalLM(SupraBrainPreTrainedModel, GenerationMixin):
    _tied_weights_keys = {}

    def __init__(self, cfg: SupraBrainConfig):
        super().__init__(cfg)
        self.model = SupraBrainModel(cfg)
        r = cfg.unembed_rank
        if r > 0:   # half-untied head: W_eff = W (I + A B^T), only 2*d*r params
            self.ub_a = nn.Parameter(torch.zeros(cfg.hidden_size, r))
            self.ub_b = nn.Parameter(torch.zeros(cfg.hidden_size, r))
            nn.init.normal_(self.ub_a, 0.0, cfg.initializer_range)
        
        # GenerationMixin requires main_input_name to be specified
        self.main_input_name = "input_ids"
        self.post_init()

    def get_input_embeddings(self):
        return self.model.embed_tokens

    def set_input_embeddings(self, v):
        self.model.embed_tokens = v

    def get_output_embeddings(self):
        return None

    def tie_weights(self, *args, **kwargs):
        return  # The head uses the embedding matrix directly.

    def _unembed_weight(self):
        W = self.model.embed_tokens.weight
        if self.config.unembed_rank > 0:
            W = W + (W @ self.ub_a) @ self.ub_b.t()
        return W

    def forward(self, input_ids=None, labels=None, **kw):
        h = self.model(input_ids)
        W = self._unembed_weight()
        cfg = self.config

        if labels is None:
            logits = h @ W.t()
            if cfg.logit_softcap > 0:
                logits = cfg.logit_softcap * torch.tanh(logits / cfg.logit_softcap)
            return CausalLMOutputWithPast(logits=logits)

        # Shift
        hs = h[:, :-1].reshape(-1, cfg.hidden_size)
        ls = labels[:, 1:].reshape(-1)
        # Chunked + checkpointed CE: (T*V) fp32 logits would consume 1.5 GB
        tot = hs.new_zeros((), dtype=torch.float32)
        for hc, lc in zip(hs.split(cfg.ce_chunk, 0), ls.split(cfg.ce_chunk, 0)):
            tot = tot + torch.utils.checkpoint.checkpoint(
                _ce_chunk, hc, W, lc, cfg.z_loss, cfg.logit_softcap,
                use_reentrant=False)
        return CausalLMOutputWithPast(loss=tot / ls.numel(), logits=None)

    def prepare_inputs_for_generation(
        self, input_ids, past_key_values=None, attention_mask=None, **kwargs
    ):
        return {
            "input_ids": input_ids,
            "past_key_values": past_key_values,
            "use_cache": kwargs.get("use_cache", False),
            "attention_mask": attention_mask,
        }


# ---- HF Auto-Registration -------------------------------------------------
def register_hf():
    from transformers import AutoConfig, AutoModelForCausalLM
    try:
        AutoConfig.register("suprabrain", SupraBrainConfig)
        AutoModelForCausalLM.register(SupraBrainConfig, SupraBrainForCausalLM)
    except Exception:
        pass


register_hf()


# ===========================================================================
#  4. MUON + ADAMW HYBRID (one optimizer object -> trainer-compatible)
# ===========================================================================
def newtonschulz5(G, steps=5, a=3.4445, b=-4.7750, c=2.0315):
    X = G.bfloat16()
    transposed = X.size(-2) > X.size(-1)
    if transposed:
        X = X.mT
    X = X / (X.norm(dim=(-2, -1), keepdim=True) + 1e-7)
    for _ in range(steps):
        A = X @ X.mT
        B = b * A + c * (A @ A)
        X = a * X + B @ X
    return (X.mT if transposed else X).to(G.dtype)


class MuonAdamW(torch.optim.Optimizer):
    """Parameter groups with use_muon=True/False. Muon for the 2D body, AdamW for the rest."""
    def __init__(self, groups):
        defaults = dict(lr=1e-3, wd=0.0, use_muon=False, momentum=0.95,
                        nesterov=True, ns_steps=5, betas=(0.9, 0.95), eps=1e-10,
                        lr_adjust="original")
        super().__init__(groups, defaults)

    @torch.no_grad()
    def step(self, closure=None):
        loss = closure() if closure is not None else None
        for g in self.param_groups:
            lr, wd = g["lr"], g["wd"]
            if g["use_muon"]:
                for p in g["params"]:
                    if p.grad is None:
                        continue
                    st = self.state[p]
                    if "m" not in st:
                        st["m"] = torch.zeros_like(p)
                    buf = st["m"]
                    buf.lerp_(p.grad, 1 - g["momentum"])
                    upd = p.grad.lerp(buf, g["momentum"]) if g["nesterov"] else buf
                    upd = newtonschulz5(upd.reshape(len(upd), -1), g["ns_steps"])
                    if g["lr_adjust"] == "match_rms_adamw":
                        scale = 0.2 * math.sqrt(max(p.size(-2), p.size(-1)))
                    else:                                  # Keller Jordan original
                        scale = max(1.0, p.size(-2) / p.size(-1)) ** 0.5
                    if wd:
                        p.mul_(1 - lr * wd)
                    p.add_(upd.view_as(p), alpha=-lr * scale)
            else:
                b1, b2 = g["betas"]
                for p in g["params"]:
                    if p.grad is None:
                        continue
                    st = self.state[p]
                    if "step" not in st:
                        st["step"] = 0
                        st["m"] = torch.zeros_like(p)
                        st["v"] = torch.zeros_like(p)
                    st["step"] += 1
                    t = st["step"]
                    st["m"].lerp_(p.grad, 1 - b1)
                    st["v"].mul_(b2).addcmul_(p.grad, p.grad, value=1 - b2)
                    mh = st["m"] / (1 - b1 ** t)
                    vh = st["v"] / (1 - b2 ** t)
                    if wd:
                        p.mul_(1 - lr * wd)
                    p.addcdiv_(mh, vh.sqrt().add_(g["eps"]), value=-lr)
        return loss


def build_optimizer(model, muon_lr=0.015, embed_lr=3e-3, other_lr=1e-3,
                    muon_wd=0.05, momentum=0.95):
    muon, embed, other = [], [], []
    for n, p in model.named_parameters():
        if not p.requires_grad:
            continue
        if "embed_tokens" in n or n.startswith("ub_"):
            embed.append(p)                       # Head/Embedding -> AdamW
        elif p.ndim == 2 and min(p.shape) >= 32:
            muon.append(p)                        # Body-Matrices -> Muon
        else:
            other.append(p)                       # Norms, Gates, Conv1d, Skalare
    groups = [
        dict(params=muon, use_muon=True, lr=muon_lr, wd=muon_wd,
             momentum=momentum, nesterov=True, ns_steps=5, lr_adjust="original"),
        dict(params=embed, use_muon=False, lr=embed_lr, wd=0.0,
             betas=(0.9, 0.95), eps=1e-10),
        dict(params=other, use_muon=False, lr=other_lr, wd=0.0,
             betas=(0.9, 0.95), eps=1e-10),
    ]
    n = lambda ps: sum(p.numel() for p in ps)
    print(f"[*] Optimizer: Muon {n(muon)/1e6:.2f}M | Embed {n(embed)/1e6:.2f}M "
          f"| Other {n(other)/1e3:.1f}k")
    return MuonAdamW(groups)


def wsd_lambda(total, warmup=0.02, stable_end=0.72, floor=0.0):
    w, s = int(total * warmup), int(total * stable_end)

    def f(step):
        if step < w:
            return (step + 1) / w
        if step < s:
            return 1.0
        p = (step - s) / max(1, total - s)
        return max(floor, 1.0 - math.sqrt(p))       # 1-sqrt Decay > linear
    return f


# ===========================================================================
#  5. DATA: streaming -> .bin (uint16, memmap, RAM-gentle)
# ===========================================================================
def build_bin(path, target_tokens, text_iter, tokenizer, eos_id,
              batch_texts=1000, flush_every=1_000_000, meta_path=None):
    from tqdm import tqdm
    if os.path.exists(path) and os.path.getsize(path) >= target_tokens * 2:
        print(f"[=] Reusing {path}")
        return
    print(f"[*] Streaming + tokenizing {target_tokens:,} tokens -> {path}")
    mm = np.memmap(path, dtype=np.uint16, mode="w+", shape=(target_tokens,))
    written, n_bytes, buf, texts = 0, 0, [], []
    pbar = tqdm(total=target_tokens, desc=f"[*] {os.path.basename(path)}", unit="tok")

    def flush():
        nonlocal written, buf
        if not buf:
            return False
        k = min(len(buf), target_tokens - written)
        mm[written:written + k] = np.asarray(buf[:k], dtype=np.uint16)
        written += k
        pbar.update(k)
        del buf[:k]
        return written >= target_tokens

    done = False
    for txt in text_iter:
        texts.append(txt)
        n_bytes += len(txt.encode("utf-8")) + 1
        if len(texts) >= batch_texts:
            for e in tokenizer.encode_batch(texts):
                buf.extend(e.ids)
                buf.append(eos_id)
            texts.clear()
            if len(buf) >= flush_every and flush():
                done = True
                break
    if not done and texts:
        for e in tokenizer.encode_batch(texts):
            buf.extend(e.ids)
            buf.append(eos_id)
        flush()
    pbar.close()
    mm.flush()
    del mm
    print(f"[+] {written:,} tokens ({os.path.getsize(path)/1e9:.2f} GB)")
    if meta_path:
        with open(meta_path, "w") as f:
            json.dump({"tokens": int(written), "bytes": int(n_bytes),
                       "bytes_per_token": n_bytes / max(1, written)}, f, indent=2)


class PhasedMemmapDataset(Dataset):
    """Sequentially across phases (WSD data change), permuted intra-phase."""
    def __init__(self, specs, seq_len, seed=1234):
        self.seq_len = seq_len
        self.parts = []
        for i, (path, ntok) in enumerate(specs):
            nc = ntok // seq_len
            perm = np.random.default_rng(seed + i).permutation(nc)
            self.parts.append({"path": path, "nc": nc, "perm": perm, "mm": None})
        self.cum = np.cumsum([p["nc"] for p in self.parts])

    def __len__(self):
        return int(self.cum[-1])

    def _mm(self, p):
        if p["mm"] is None:                       # lazy -> multiprocessing-safe
            p["mm"] = np.memmap(p["path"], dtype=np.uint16, mode="r",
                                shape=(p["nc"] * self.seq_len,))
        return p["mm"]

    def __getitem__(self, i):
        pi = int(np.searchsorted(self.cum, i, side="right"))
        local = i - (0 if pi == 0 else int(self.cum[pi - 1]))
        p = self.parts[pi]
        s = int(p["perm"][local]) * self.seq_len
        ids = torch.from_numpy(np.asarray(self._mm(p)[s:s + self.seq_len],
                                          dtype=np.int64))
        return {"input_ids": ids, "labels": ids.clone()}


def collate_fn(batch):
    ids = torch.stack([b["input_ids"] for b in batch])
    return {"input_ids": ids, "labels": torch.stack([b["labels"] for b in batch])}


# ===========================================================================
#  6. TRAINING
# ===========================================================================
SEQ_LEN        = 1024
TOTAL_TOKENS   = 5_000_000_000
STABLE_TOKENS  = 3_600_000_000        # Phase 1: FineWeb-Edu (broad)
ANNEAL_TOKENS  = 1_400_000_000        # Phase 2: Top-Tier + Cosmopedia
VAL_TOKENS     = 4_194_304
MICRO_BS       = 16                    # 16 GB VRAM: 8 x 1024
GRAD_ACC       = 8                   # -> 262144 Tokens/Step, ~19.1k Steps
TOK_DIR        = "./suprabrain-tokenizer"
OUT_DIR        = "./SupraBrain-50M-v0.1"
SMOKE_TEST     = False                # True: Mini-Run for Validation


def make_config():
    return SupraBrainConfig(
        vocab_size=23_808, hidden_size=384, num_hidden_layers=28,
        intermediate_size=1152, mlp_act="rational",
        gdn_num_heads=3, gdn_head_dim=128,
        attn_num_heads=6, attn_num_kv_heads=2, attn_head_dim=64,
        attn_window=256, attn_every=4, full_attn_layers=(19,),
        max_position_embeddings=SEQ_LEN, surprise=True,
        unembed_rank=32, z_loss=1e-4, compile_blocks=True,
    )


def main():
    from datasets import load_dataset, interleave_datasets
    from tokenizers import Tokenizer
    from transformers import (PreTrainedTokenizerFast, Trainer, TrainingArguments,
                              TrainerCallback)

    torch.set_float32_matmul_precision("high")
    print(f"[*] FLA={_HAS_FLA} | FlexAttention={_HAS_FLEX} | "
          f"CUDA={torch.cuda.is_available()}")
    if not _HAS_FLA:
        print("[!] WARNING: flash-linear-attention missing -> slow Reference-GDN!")
        print("[!]          pip install flash-linear-attention")

    # ---------------- Tokenizer ----------------
    raw_tok = Tokenizer.from_file(os.path.join(TOK_DIR, "tokenizer.json"))
    tokenizer = PreTrainedTokenizerFast.from_pretrained(TOK_DIR)
    EOS = tokenizer.eos_token_id

    tot, stab, ann, nval = (TOTAL_TOKENS, STABLE_TOKENS, ANNEAL_TOKENS, VAL_TOKENS)
    if SMOKE_TEST:
        tot, stab, ann, nval = 8_000_000, 6_000_000, 2_000_000, 1_000_000

    # ---------------- Build bins ----------------
    os.makedirs("data", exist_ok=True)
    fw = load_dataset("HuggingFaceFW/fineweb-edu", "sample-100BT",
                      split="train", streaming=True)
    stream = iter(fw)

    # Val first from the head of the stream -> guarantee disjunked from training data!
    build_bin("data/val.bin", nval, (next(stream)["text"] for _ in iter(int, 1)),
              raw_tok, EOS, meta_path="data/val_meta.json")
    # Stable-Phase: same stream after val
    build_bin("data/stable.bin", stab,
              (next(stream)["text"] for _ in iter(int, 1)), raw_tok, EOS)

    # Anneal-Phase: 65% FineWeb-Edu score>=4.2 + 35% Cosmopedia-v2
    hi = load_dataset("HuggingFaceFW/fineweb-edu", "default",
                       split="train", streaming=True)
    cosmo = load_dataset("HuggingFaceTB/smollm-corpus", "cosmopedia-v2",
                         split="train", streaming=True)
    mix = interleave_datasets([hi, cosmo], probabilities=[0.65, 0.35], seed=42,
                              stopping_strategy="all_exhausted")
    build_bin("data/anneal.bin", ann, (e["text"] for e in mix), raw_tok, EOS)

    train_ds = PhasedMemmapDataset([("data/stable.bin", stab),
                                    ("data/anneal.bin", ann)], SEQ_LEN)
    val_ds = PhasedMemmapDataset([("data/val.bin", nval)], SEQ_LEN, seed=0)
    print(f"[+] Train: {len(train_ds):,} chunks | Val: {len(val_ds):,} chunks")

    # ---------------- Model ----------------
    cfg = make_config()
    cfg.pad_token_id, cfg.bos_token_id, cfg.eos_token_id = (
        tokenizer.pad_token_id, tokenizer.bos_token_id, tokenizer.eos_token_id)
    if SMOKE_TEST:
        cfg.num_hidden_layers, cfg.compile_blocks = 8, False
        cfg.full_attn_layers = ()

    model = SupraBrainForCausalLM(cfg)
    n_all = model.num_parameters()
    n_emb = cfg.vocab_size * cfg.hidden_size
    print(f"[*] Parameters      : {n_all:,}  ({n_all/1e6:.2f} M)")
    print(f"[*] Non-Embedding   : {(n_all-n_emb)/1e6:.2f} M")
    print(f"[*] Attention layers: {cfg.attn_layers}")
    assert n_all < 50_000_000, f"Budget busted: {n_all:,} > 50M"

    # ---------------- Optimizer + Schedule ----------------
    steps = len(train_ds) // (MICRO_BS * GRAD_ACC)
    opt = build_optimizer(model, muon_lr=0.015, embed_lr=3e-3, other_lr=1e-3)
    sched = torch.optim.lr_scheduler.LambdaLR(opt, wsd_lambda(steps))
    print(f"[*] Steps: {steps:,} | Tokens/Step: {MICRO_BS*GRAD_ACC*SEQ_LEN:,}")

    class MomentumWarmup(TrainerCallback):
        """0.85 -> 0.95 over 200 steps: prevents early Muon-Spikes."""
        def on_step_begin(self, args, state, control, optimizer=None, **kw):
            if optimizer is None or state.global_step > 200:
                return
            m = 0.85 + 0.10 * min(1.0, state.global_step / 200)
            for g in optimizer.param_groups:
                if g.get("use_muon"):
                    g["momentum"] = m

    class BPBCallback(TrainerCallback):
        """The only fair comparison to the old baseline: Bits-per-Byte."""
        def __init__(self, meta="data/val_meta.json"):
            self.bpt = json.load(open(meta))["bytes_per_token"] if \
                os.path.exists(meta) else None

        def on_evaluate(self, args, state, control, metrics=None, **kw):
            if metrics and "eval_loss" in metrics and self.bpt:
                bpb = metrics["eval_loss"] / (math.log(2) * self.bpt)
                metrics["eval_bpb"] = bpb
                print(f"[BPB] step {state.global_step}: loss="
                      f"{metrics['eval_loss']:.4f} bpb={bpb:.4f} "
                      f"(Baseline Llama-50M @20B: 1.068)")

    class SeqTrainer(Trainer):
        def __init__(self, *a, **kw):
            super().__init__(*a, **kw)
            self.model_accepts_loss_kwargs = False   # wir normalisieren nicht per num_items_in_batch

        def _get_train_sampler(self, *a, **kw):
            return SequentialSampler(self.train_dataset)   # WSD-Phasing plan!

    ta_kw = dict(
        output_dir=OUT_DIR, max_steps=steps if not SMOKE_TEST else 30,
        per_device_train_batch_size=MICRO_BS,
        gradient_accumulation_steps=GRAD_ACC,
        per_device_eval_batch_size=MICRO_BS,
        logging_steps=1, save_steps=1000, save_total_limit=3,
        eval_steps=1000, prediction_loss_only=True,
        bf16=True, fp16=False, torch_compile=False,      # we compile by ourselves
        max_grad_norm=1.0,                               # Muon is scale-invariant
        dataloader_num_workers=min(8, (os.cpu_count() or 4) // 2),
        dataloader_pin_memory=True, dataloader_drop_last=True,
        report_to="none", remove_unused_columns=False, seed=1234,
    )
    try:
        args = TrainingArguments(eval_strategy="steps", **ta_kw)
    except TypeError:
        args = TrainingArguments(evaluation_strategy="steps", **ta_kw)

    trainer = SeqTrainer(model=model, args=args, train_dataset=train_ds,
                         eval_dataset=val_ds, data_collator=collate_fn,
                         optimizers=(opt, sched),
                         callbacks=[MomentumWarmup(), BPBCallback()])

    if cfg.compile_blocks and torch.cuda.is_available():
        for b in trainer.model.model.layers:
            b.compile(mode="max-autotune-no-cudagraphs", dynamic=False)
        print("[*] Blocks compiled (regional, in-place).")

    print("[*] Starting training...")
    trainer.train()

    # ---------------- HF-compatible export ----------------
    final = OUT_DIR + "-FINAL"
    os.makedirs(final, exist_ok=True)
    inner = model
    for i, b in enumerate(inner.model.layers):        # remove compile-Wrapper
        if type(b).__name__ == "OptimizedModule":
            inner.model.layers[i] = b._orig_mod
    inner.config.auto_map = {
        "AutoConfig": "modeling_suprabrain.SupraBrainConfig",
        "AutoModelForCausalLM": "modeling_suprabrain.SupraBrainForCausalLM",
    }
    inner.save_pretrained(final, safe_serialization=True)
    tokenizer.save_pretrained(final)
    shutil.copyfile(os.path.abspath(__file__),
                    os.path.join(final, "modeling_suprabrain.py"))
    print(f"[+] Saved to {final}")
    print("[i] Load via:  AutoModelForCausalLM.from_pretrained("
          f"'{final}', trust_remote_code=True)")
    print("[*] Training finished.")


if __name__ == "__main__":
    main()