""" © 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()