modeling_alpha.py

"""alpha-er — a PyTorch reference implementation.

alpha-er was trained end to end on a from-scratch GPU stack: our own ioctl
driver, memory manager, sm_86 SASS assembler and kernel IR, with no CUDA, no
cuBLAS and no vendor runtime anywhere in the training path. This file exists so
the resulting weights are usable by people who do not have that stack — it is a
faithful re-expression of the same arithmetic in PyTorch, not the trainer.

Three things about the architecture are unusual enough to explain, because each
one is why the model cannot be loaded as a Llama and would be silently wrong if
you tried.

FACTORED PROJECTIONS. The QKV projection, the attention output projection and
the LM head are each a rank-r bottleneck rather than a dense matrix: x @ down.T,
then a LayerNorm on the r-dimensional bottleneck, then @ up.T. The norm is not
decoration — without it the factored form diverged in training (grad_norm 51
against a dense baseline's 1.25). r = 128 here for all three.

CONDITIONAL MLP. The feed-forward block is partitioned into G = 64 experts of
width ffn/G = 320. A token is routed to exactly one, so the model stores a
20,480-wide FFN but any single token pays for 320. Routing is POSITIONAL and
depends only on the token's index within its own sequence:

    expert(t) = floor(t * G / T)

That detail matters more than it looks. An earlier version routed on the index
in the flattened batch, which made a token's expert depend on how many other
sequences shared its batch: each sequence reached only 4 of the 64 experts, and
the weights were meaningful only at the exact batch shape they were trained at.
Routing on t makes a checkpoint portable — the same sequence gives identical
logits at any batch width — and lets every sequence reach every expert.

SEQUENCE LENGTH IS PART OF THE ARCHITECTURE. Because the expert boundaries fall
at multiples of T/G, the model reproduces its training behaviour only at its
trained context length. Generation must pad the prompt to block_size and read
the logits at the last real position, which is exact rather than approximate:
attention is causal, so padding after the prompt cannot influence it.

Trained on 1.97B tokens (FineWeb-Edu/DCLM + Concordance-EN + SmolTalk) at
~96,000 tokens/second on a single RTX 3070.
"""

from __future__ import annotations

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


class AlphaErConfig:
    """Mirrors the trainer's ModelConfig. Field names match the checkpoint header."""

    model_type = "alpha-er"

    def __init__(
        self,
        vocab_size: int = 12288,
        block_size: int = 512,
        n_layer: int = 2,
        n_embd: int = 1024,
        n_head: int = 8,
        ffn_dim: int = 20480,
        mlp_experts: int = 64,
        lm_head_rank: int = 128,
        attn_rank: int = 128,
        attn_soft_cap: float = 30.0,
        layer_norm_eps: float = 1e-5,
        **kwargs,
    ):
        self.vocab_size = vocab_size
        self.block_size = block_size
        self.n_layer = n_layer
        self.n_embd = n_embd
        self.n_head = n_head
        self.ffn_dim = ffn_dim
        self.mlp_experts = mlp_experts
        self.lm_head_rank = lm_head_rank
        self.attn_rank = attn_rank
        self.attn_soft_cap = attn_soft_cap
        self.layer_norm_eps = layer_norm_eps
        for k, v in kwargs.items():
            setattr(self, k, v)


def gelu(x: torch.Tensor) -> torch.Tensor:
    """The tanh approximation, which is what the native kernel computes.

    Using the exact erf form here would leave the weights unchanged and the
    outputs subtly different, which is the kind of mismatch that survives a
    smoke test and shows up as degraded generation.
    """
    return F.gelu(x, approximate="tanh")


class FactoredProjection(nn.Module):
    """x @ down.T -> LayerNorm on the rank-r bottleneck -> @ up.T."""

    def __init__(self, d_in: int, rank: int, d_out: int, eps: float):
        super().__init__()
        self.down = nn.Parameter(torch.empty(rank, d_in))
        self.mid_w = nn.Parameter(torch.ones(rank))
        self.mid_b = nn.Parameter(torch.zeros(rank))
        self.up = nn.Parameter(torch.empty(d_out, rank))
        self.eps = eps

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        z = F.linear(x, self.down)
        z = F.layer_norm(z, (z.shape[-1],), self.mid_w, self.mid_b, self.eps)
        return F.linear(z, self.up)


class AlphaErBlock(nn.Module):
    def __init__(self, cfg: AlphaErConfig):
        super().__init__()
        self.cfg = cfg
        d, r = cfg.n_embd, cfg.attn_rank
        self.ln1 = nn.LayerNorm(d, eps=cfg.layer_norm_eps)
        self.qkv = FactoredProjection(d, r, 3 * d, cfg.layer_norm_eps)
        self.proj = FactoredProjection(d, r, d, cfg.layer_norm_eps)
        self.ln2 = nn.LayerNorm(d, eps=cfg.layer_norm_eps)

        G, fe = cfg.mlp_experts, cfg.ffn_dim // cfg.mlp_experts
        self.fc1s = nn.Parameter(torch.empty(G, fe, d))
        self.fc2s = nn.Parameter(torch.empty(G, d, fe))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B, T, d = x.shape
        H = self.cfg.n_head
        hd = d // H

        a = self.ln1(x)
        q, k, v = self.qkv(a).split(d, dim=-1)
        q = q.view(B, T, H, hd).transpose(1, 2)
        k = k.view(B, T, H, hd).transpose(1, 2)
        v = v.view(B, T, H, hd).transpose(1, 2)
        # Attention with LOGIT SOFT-CAPPING, which is not optional here: the
        # trainer applies softCap = 30 by default for non-RoPE models, and
        # F.scaled_dot_product_attention has no equivalent. Omitting it leaves
        # the weights untouched and the logits completely different — the port
        # disagreed with the real model by 82% of logit magnitude until this
        # was added, while still producing entirely plausible-looking numbers.
        cap = self.cfg.attn_soft_cap
        scores = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(hd))
        if cap and cap > 0:
            scores = torch.tanh(scores / cap) * cap
        mask = torch.ones(T, T, dtype=torch.bool, device=x.device).tril()
        scores = scores.masked_fill(~mask, float("-inf"))
        att = torch.softmax(scores, dim=-1) @ v
        x = x + self.proj(att.transpose(1, 2).reshape(B, T, d))

        # Conditional MLP. einsum expresses the routing directly: split the
        # sequence into G contiguous windows and give window g to expert g.
        G = self.cfg.mlp_experts
        S = T // G
        b = self.ln2(x).view(B, G, S, d)
        h = gelu(torch.einsum("bgsd,gfd->bgsf", b, self.fc1s))
        y = torch.einsum("bgsf,gdf->bgsd", h, self.fc2s).reshape(B, T, d)
        return x + y


class AlphaErForCausalLM(nn.Module):
    def __init__(self, cfg: AlphaErConfig):
        super().__init__()
        self.cfg = cfg
        self.wte = nn.Embedding(cfg.vocab_size, cfg.n_embd)
        self.wpe = nn.Embedding(cfg.block_size, cfg.n_embd)
        self.blocks = nn.ModuleList([AlphaErBlock(cfg) for _ in range(cfg.n_layer)])
        self.ln_f = nn.LayerNorm(cfg.n_embd, eps=cfg.layer_norm_eps)
        self.lm_head = FactoredProjection(
            cfg.n_embd, cfg.lm_head_rank, cfg.vocab_size, cfg.layer_norm_eps
        )

    def forward(self, input_ids: torch.Tensor) -> torch.Tensor:
        B, T = input_ids.shape
        if T % self.cfg.mlp_experts != 0:
            raise ValueError(
                f"sequence length {T} must be a multiple of mlp_experts="
                f"{self.cfg.mlp_experts}: the expert boundaries fall at multiples "
                f"of T/G, so pad the input to block_size ({self.cfg.block_size}) "
                f"and read the logits at the last real position."
            )
        pos = torch.arange(T, device=input_ids.device)
        x = self.wte(input_ids) + self.wpe(pos)
        for blk in self.blocks:
            x = blk(x)
        return self.lm_head(self.ln_f(x))

    @torch.no_grad()
    def generate(self, input_ids, max_new_tokens=64, temperature=0.8, top_k=40):
        """Pad to block_size and read the last real position — see the module docstring."""
        cfg = self.cfg
        ids = input_ids[0].tolist()
        for _ in range(max_new_tokens):
            window = ids[-cfg.block_size:]
            padded = window + [0] * (cfg.block_size - len(window))
            t = torch.tensor([padded], device=input_ids.device)
            logits = self(t)[0, len(window) - 1]
            if temperature <= 0:
                ids.append(int(logits.argmax()))
                continue
            v, i = torch.topk(logits, min(top_k, logits.numel()))
            p = torch.softmax(v / temperature, dim=-1)
            ids.append(int(i[torch.multinomial(p, 1)]))
        return torch.tensor([ids], device=input_ids.device)