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