nexus_model.py

"""
NEXUS v7.1 — Depth-Recurrent LM with Causal Attention + Memory-Efficient Training
===================================================================================
Changes from v7:
- Chunked cross-entropy loss: never materializes full [B*T, vocab] logits tensor
  → Saves ~100MB per supervision step, critical for base/medium on T4
- Skip redundant final logits computation during training
- T4-friendly base config: N_sup=4, T_inner=6, backprop_depth=3

Architecture (Huginn-style):
- Prelude: embed + RoPE
- Core R: shared transformer block, applied T_inner times per N_sup step
- Coda: output norm + LM head
- State injection: linear adapter on cat(state, embedding) per iteration
- Sandwich norm: 4 RMSNorms per layer (critical for depth stability)
- Deep supervision with k-step truncated backprop
- BitNet b1.58 ternary on MLP weights (attention weights kept fp for stability)
"""

import json, math
import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Optional, Dict


def _rms(d):
    if hasattr(nn, 'RMSNorm'):
        return nn.RMSNorm(d)
    class _R(nn.Module):
        def __init__(self, dim):
            super().__init__()
            self.w = nn.Parameter(torch.ones(dim))
            self.e = 1e-6
        def forward(self, x):
            return (x.float() * x.float().pow(2).mean(-1, keepdim=True).add(self.e).rsqrt()).to(x.dtype) * self.w
    return _R(d)


class RotaryEmbedding(nn.Module):
    """RoPE positional encoding."""
    def __init__(self, dim, max_seq_len=4096, base=10000.0):
        super().__init__()
        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
        self.register_buffer("inv_freq", inv_freq)
        self.max_seq_len = max_seq_len

    def forward(self, seq_len, device):
        t = torch.arange(seq_len, device=device, dtype=self.inv_freq.dtype)
        freqs = torch.outer(t, self.inv_freq)
        emb = torch.cat((freqs, freqs), dim=-1)
        return emb.cos(), emb.sin()


def rotate_half(x):
    x1, x2 = x.chunk(2, dim=-1)
    return torch.cat((-x2, x1), dim=-1)


def apply_rotary_pos_emb(q, k, cos, sin):
    cos = cos[None, None, :, :]
    sin = sin[None, None, :, :]
    q_embed = (q * cos) + (rotate_half(q) * sin)
    k_embed = (k * cos) + (rotate_half(k) * sin)
    return q_embed, k_embed


class BitLinear158(nn.Module):
    """Ternary linear for MLP weights."""
    def __init__(self, inf, outf):
        super().__init__()
        self.weight = nn.Parameter(torch.empty(outf, inf))
        self.norm = _rms(inf)
        nn.init.normal_(self.weight, 0, 0.02)

    def forward(self, x):
        x = self.norm(x)
        gw = self.weight.abs().mean() + 1e-7
        w_q = (self.weight / gw).round().clamp(-1, 1)
        w_ste = self.weight + (w_q * gw - self.weight).detach()
        return F.linear(x, w_ste)


class SandwichBlock(nn.Module):
    """Huginn-style transformer block with sandwich norm + RoPE."""
    def __init__(self, d_model, n_heads, expansion=4):
        super().__init__()
        self.n1 = _rms(d_model)
        self.n2 = _rms(d_model)
        self.n3 = _rms(d_model)
        self.n4 = _rms(d_model)
        self.n_heads = n_heads
        self.head_dim = d_model // n_heads
        self.qkv = nn.Linear(d_model, 3 * d_model, bias=False)
        self.o_proj = nn.Linear(d_model, d_model, bias=False)
        df = int(d_model * expansion * 2 / 3)
        self.gate = BitLinear158(d_model, df)
        self.up = BitLinear158(d_model, df)
        self.down = BitLinear158(df, d_model)

    def forward(self, x, cos=None, sin=None):
        B, T, D = x.shape
        h = self.n1(x)
        qkv = self.qkv(h).reshape(B, T, 3, self.n_heads, self.head_dim)
        q, k, v = qkv.unbind(2)
        q = q.transpose(1, 2)
        k = k.transpose(1, 2)
        v = v.transpose(1, 2)
        if cos is not None and sin is not None:
            q, k = apply_rotary_pos_emb(q, k, cos, sin)
        h = F.scaled_dot_product_attention(q, k, v, is_causal=True)
        h = h.transpose(1, 2).reshape(B, T, D)
        h = self.o_proj(h)
        x = self.n2(x + h)
        h = self.n3(x)
        h = self.down(F.silu(self.gate(h)) * self.up(h))
        x = self.n4(x + h)
        return x


def chunked_cross_entropy(hidden, weight, labels, z_loss_weight=1e-4, chunk_size=128,
                          ignore_index=-100):
    """
    Memory-efficient cross-entropy: projects hidden→logits in chunks along the
    sequence dimension, never materializing the full [B*T, vocab] tensor.
    
    Saves ~100MB per call for base config (B=4, T=256, V=49152).
    """
    B, T, d = hidden.shape
    V = weight.shape[0]
    
    # Shift: predict token t+1 from position t
    hidden = hidden[:, :-1].contiguous()
    labels = labels[:, 1:].contiguous()
    T_eff = T - 1
    
    total_ce = torch.tensor(0.0, device=hidden.device, dtype=torch.float32)
    total_zl = torch.tensor(0.0, device=hidden.device, dtype=torch.float32)
    n_tokens = 0
    
    for start in range(0, T_eff, chunk_size):
        end = min(start + chunk_size, T_eff)
        h_chunk = hidden[:, start:end]
        l_chunk = labels[:, start:end]
        
        logits_chunk = F.linear(h_chunk, weight)
        
        ce = F.cross_entropy(
            logits_chunk.reshape(-1, V), l_chunk.reshape(-1),
            ignore_index=ignore_index, reduction='sum'
        )
        
        lse = torch.logsumexp(logits_chunk.reshape(-1, V).float(), dim=-1)
        zl = (lse ** 2).sum()
        
        mask = (l_chunk != ignore_index)
        n_valid = mask.sum().item()
        n_tokens += n_valid
        
        total_ce = total_ce + ce
        total_zl = total_zl + zl
        
        del logits_chunk, lse
    
    n_tokens = max(n_tokens, 1)
    return total_ce / n_tokens, z_loss_weight * total_zl / n_tokens


class NexusConfig:
    def __init__(self, vocab_size=49152, d_model=512, n_heads=8, n_layers=2,
                 expansion=4, max_seq_len=256, N_sup=4, T_inner=6,
                 tie_weights=True, z_loss_weight=1e-4, backprop_depth=4,
                 use_rope=True, **kw):
        self.vocab_size = vocab_size
        self.d_model = d_model
        self.n_heads = n_heads
        self.n_layers = n_layers
        self.expansion = expansion
        self.max_seq_len = max_seq_len
        self.N_sup = N_sup
        self.T_inner = T_inner
        self.tie_weights = tie_weights
        self.z_loss_weight = z_loss_weight
        self.backprop_depth = backprop_depth
        self.use_rope = use_rope
        for k, v in kw.items():
            if not hasattr(self, k):
                setattr(self, k, v)

    def to_dict(self):
        return {k: v for k, v in self.__dict__.items() if not k.startswith('_')}

    @classmethod
    def from_dict(cls, d):
        return cls(**d)

    def __repr__(self):
        return f"NexusConfig({json.dumps(self.to_dict(), indent=2)})"


class NexusModel(nn.Module):
    """
    NEXUS v7.1: Depth-Recurrent LM with memory-efficient training.
    
    Uses chunked cross-entropy during deep supervision to avoid
    materializing full [B*T, vocab] logits tensors.
    """
    def __init__(self, config):
        super().__init__()
        self.config = config
        d = config.d_model
        self.emb = nn.Embedding(config.vocab_size, d)
        self.emb_norm = _rms(d)
        if config.use_rope:
            self.rope = RotaryEmbedding(d // config.n_heads, max_seq_len=config.max_seq_len)
        else:
            self.rope = None
        self.adapter = nn.Linear(2 * d, d, bias=False)
        self.adapter_norm = _rms(d)
        self.core = nn.ModuleList([
            SandwichBlock(d, config.n_heads, config.expansion)
            for _ in range(config.n_layers)
        ])
        self.iter_norm = _rms(d)
        self.out_norm = _rms(d)
        self.head = nn.Linear(d, config.vocab_size, bias=False)
        if config.tie_weights:
            self.head.weight = self.emb.weight
        self._init()

    def _init(self):
        eff = max(1, self.config.n_layers * self.config.T_inner * self.config.N_sup)
        scale = (2 / (5 * self.config.d_model)) ** 0.5
        for m in self.modules():
            if isinstance(m, nn.Linear):
                nn.init.normal_(m.weight, 0, scale)
                if m.bias is not None:
                    nn.init.zeros_(m.bias)
            elif isinstance(m, nn.Embedding):
                nn.init.normal_(m.weight, 0, scale)
        for name, p in self.named_parameters():
            if 'o_proj' in name or 'down' in name:
                p.data.mul_(1.0 / math.sqrt(eff))

    def _get_rope(self, seq_len, device):
        if self.rope is not None:
            return self.rope(seq_len, device)
        return None, None

    def _core_step(self, state, embedding, cos=None, sin=None):
        h = self.adapter_norm(self.adapter(torch.cat([state, embedding], dim=-1)))
        for block in self.core:
            h = block(h, cos=cos, sin=sin)
        return self.iter_norm(h)

    def forward(self, input_ids, labels=None):
        B, T = input_ids.shape
        n_sup = self.config.N_sup
        T_inner = self.config.T_inner
        k = min(self.config.backprop_depth, T_inner)

        embedding = self.emb_norm(self.emb(input_ids))
        cos, sin = self._get_rope(T, input_ids.device)
        state = torch.zeros_like(embedding)

        total_loss = torch.tensor(0.0, device=input_ids.device)
        n_losses = 0

        for sup in range(n_sup):
            n_nograd = T_inner - k
            if n_nograd > 0:
                with torch.no_grad():
                    for _ in range(n_nograd):
                        state = self._core_step(state, embedding, cos, sin)

            for _ in range(k):
                state = self._core_step(state, embedding, cos, sin)

            if labels is not None:
                h = self.out_norm(state)
                ce, zl = chunked_cross_entropy(
                    h, self.head.weight, labels,
                    z_loss_weight=self.config.z_loss_weight,
                    chunk_size=64,
                )
                total_loss = total_loss + ce + zl
                n_losses += 1
                del h
            state = state.detach()

        if labels is not None:
            return {
                'loss': total_loss / max(n_losses, 1),
                'logits': None,
            }
        else:
            logits = self.head(self.out_norm(state))
            return {
                'loss': None,
                'logits': logits,
            }

    @torch.no_grad()
    def generate(self, input_ids, max_new_tokens=256, temperature=0.8, top_p=0.9, top_k=50):
        self.eval()
        gen = input_ids.clone()
        for _ in range(max_new_tokens):
            ctx = gen[:, -self.config.max_seq_len:]
            T = ctx.shape[1]
            emb = self.emb_norm(self.emb(ctx))
            cos, sin = self._get_rope(T, ctx.device)
            state = torch.zeros_like(emb)
            for _ in range(self.config.N_sup * self.config.T_inner):
                state = self._core_step(state, emb, cos, sin)
            logits = self.head(self.out_norm(state))[:, -1] / temperature
            if top_k > 0:
                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
                logits[logits < v[..., -1, None]] = float('-inf')
            if top_p < 1.0:
                sl, si = torch.sort(logits, descending=True)
                cp = torch.cumsum(F.softmax(sl, -1), -1)
                rm = cp > top_p
                rm[..., 1:] = rm[..., :-1].clone()
                rm[..., 0] = 0
                logits[rm.scatter(1, si, rm)] = float('-inf')
            gen = torch.cat([gen, torch.multinomial(F.softmax(logits, -1), 1)], 1)
        return gen

    def count_params(self):
        c = {}
        t = 0
        for n, p in self.named_parameters():
            k = n.split('.')[0]
            c[k] = c.get(k, 0) + p.numel()
            t += p.numel()
        c['total'] = t
        return c

    def memory_estimate(self):
        p = sum(x.numel() for x in self.parameters())
        return {
            'model_fp16_mb': p * 2 / 1e6,
            'model_ternary_mb': p * 1.58 / 8 / 1e6,
            'total_ternary_mb': p * 1.58 / 8 / 1e6 + 1,
        }


def nexus_tiny():
    return NexusConfig(d_model=256, n_heads=4, n_layers=2, expansion=4,
                       max_seq_len=256, N_sup=4, T_inner=4, backprop_depth=4)

def nexus_small():
    return NexusConfig(d_model=512, n_heads=8, n_layers=2, expansion=4,
                       max_seq_len=256, N_sup=4, T_inner=6, backprop_depth=4)

def nexus_base():
    """T4-friendly base: reduced N_sup/T_inner vs original, backprop_depth=3."""
    return NexusConfig(d_model=768, n_heads=12, n_layers=3, expansion=4,
                       max_seq_len=256, N_sup=4, T_inner=6, backprop_depth=3)

def nexus_medium():
    return NexusConfig(d_model=1024, n_heads=16, n_layers=3, expansion=4,
                       max_seq_len=256, N_sup=4, T_inner=6, backprop_depth=2)