"""Causal transformer; RoPE conventions adapted from our A1, other layers rebuilt.""" from types import SimpleNamespace import torch from torch import nn import torch.nn.functional as F from src.part1.ffn import FeedForward # Multi-head causal attention with rotary positional embeddings (RoPE). class Attention(nn.Module): def __init__(self, width, heads): """Create projections and RoPE frequencies for equal-width attention heads.""" super().__init__() # Split the model width evenly across attention heads. self.heads = heads self.head_dim = width // heads # q, k, v and o are query, key, value and output projections. # Keep these attribute names: saved Turing checkpoints use them. self.q = nn.Linear(width, width, bias=False) self.k = nn.Linear(width, width, bias=False) self.v = nn.Linear(width, width, bias=False) self.o = nn.Linear(width, width, bias=False) # A buffer follows the model device but is not optimized. # persistent=False omits these reconstructible frequencies from checkpoints. self.register_buffer( "inverse_frequency", 10000 ** (-torch.arange(0, self.head_dim, 2).float() / self.head_dim), persistent=False, ) def rotate(self, x, positions): """Rotate adjacent feature pairs using each token's absolute position.""" # x has shape [batch, heads, time, head_dim]. # Broadcast positions against frequencies to obtain one angle per pair. angles = ( positions.float()[:, None, :, None] * self.inverse_frequency[None, None, None, :] ) # RoPE rotates adjacent even/odd feature pairs without changing their norm. even, odd = x[..., 0::2], x[..., 1::2] cos, sin = angles.cos().to(x.dtype), angles.sin().to(x.dtype) # stack interleaves the rotated pair; flatten restores head_dim features. return torch.stack( (even * cos - odd * sin, even * sin + odd * cos), -1 ).flatten(-2) def forward(self, x, mask, past=None, use_cache=False): """Return attention output and optionally the extended key/value cache.""" # Short shape names: b=batch size, t=new token count, d=model width. b, t, d = x.shape # Cached keys already cover previous tokens during incremental decoding. previous = past[0].shape[-2] if past is not None else 0 # Count visible tokens, ignoring left padding, to get zero-based positions. positions = mask.long().cumsum(-1).sub(1).clamp_min(0)[:, -t:] # Project [b,t,d], separate heads, then move heads before the time axis. def split(module): """Reshape one projected tensor into [batch,heads,time,head_dim].""" return module(x).reshape(b, t, self.heads, self.head_dim).transpose(1, 2) # Apply position rotations to queries and keys; values stay unrotated. q, k, v = ( self.rotate(split(self.q), positions), self.rotate(split(self.k), positions), split(self.v), ) # Append only the new keys/values to the cached sequence (axis -2). if past is not None: k, v = torch.cat((past[0], k), -2), torch.cat((past[1], v), -2) s = k.shape[-2] # Unpadded training/prefill can use PyTorch SDPA # (scaled dot-product attention) with its built-in causal mask. if previous == 0 and bool(mask.all()): out = F.scaled_dot_product_attention(q, k, v, is_causal=True) else: # With padding or a cache, build the full [new_time,total_time] mask. # A query can see earlier keys, but never a future key or a PAD position. causal = ( torch.arange(s, device=x.device)[None, :] <= (torch.arange(t, device=x.device) + previous)[:, None] ) visible = causal[None, None, :, :] & mask[:, None, None, :].bool() out = F.scaled_dot_product_attention(q, k, v, attn_mask=visible) # Merge the heads back into [batch,time,width] before the output projection. merged = out.transpose(1, 2).reshape(b, t, d) return self.o(merged), (k, v) if use_cache else None # One pre-norm transformer block: attention residual, then FFN residual. class Block(nn.Module): def __init__(self, config): """Construct attention, two pre-normalizations and the chosen FFN variant.""" super().__init__() d = config["d_model"] # RMSNorm normalizes root mean square magnitude before each sublayer. self.norm1, self.norm2 = nn.RMSNorm(d), nn.RMSNorm(d) self.attention = Attention(d, config["heads"]) self.ffn = FeedForward( d, config["hidden"], config["ffn"], config["aux_coefficient"] ) def forward(self, x, mask, past=None, use_cache=False, collect=False): """Return updated hidden states, router loss, attention cache and usage map.""" # Pre-norm means normalize first, then add the sublayer output to x. attended, cache = self.attention(self.norm1(x), mask, past, use_cache) x = x + attended # The FFN receives only the mask of current tokens, not cached tokens. fed, aux, usage = self.ffn(self.norm2(x), mask[:, -x.shape[1] :], collect) return x + fed, aux, cache, usage # Shared decoder backbone for all Part 1 variants and Part 2 optimizers. class Transformer(nn.Module): def __init__(self, config): """Build the shared decoder and tie output weights to input embeddings.""" super().__init__() self.config = dict(config) self.embedding = nn.Embedding(config["vocab_size"], config["d_model"]) self.blocks = nn.ModuleList([Block(config) for _ in range(config["layers"])]) self.norm = nn.RMSNorm(config["d_model"]) self.head = nn.Linear(config["d_model"], config["vocab_size"], bias=False) # Tie input embeddings and output classifier weights to save parameters. self.head.weight = self.embedding.weight # Recursively initialize modules in the original Turing order. self.apply(self.initialize) @staticmethod def initialize(module): """Use the original normal initialization for each linear/embedding module.""" if isinstance(module, (nn.Linear, nn.Embedding)): nn.init.normal_(module.weight, std=0.02) def forward( self, input_ids, attention_mask=None, past_key_values=None, use_cache=False, collect_routing=False, **kwargs, ): """Return logits, mean router loss, optional caches and expert usage.""" b, t = input_ids.shape past = past_key_values prior = past[0][0].shape[-2] if past is not None else 0 # A missing mask means every input/cached position is visible. if attention_mask is None: attention_mask = torch.ones( b, prior + t, device=input_ids.device, dtype=torch.bool ) x = self.embedding(input_ids) # new_zeros creates a scalar on the same device; router losses use FP32. aux = x.new_zeros((), dtype=torch.float32) caches, routing = [], [] # Run layers in order and keep one key/value cache and usage map per layer. for i, block in enumerate(self.blocks): x, loss, cache, usage = block( x, attention_mask, past[i] if past is not None else None, use_cache, collect_routing, ) aux = aux + loss if use_cache: caches.append(cache) if usage is not None: routing.append(usage) # Expose named outputs like a Hugging Face forward result. # The auxiliary loss is averaged over layers; logits are [b,t,vocab]. return SimpleNamespace( logits=self.head(self.norm(x)), aux_loss=aux / max(len(self.blocks), 1), past_key_values=tuple(caches) if use_cache else None, routing=routing, ) def parameter_counts(self): """Count all weights and the minimum/maximum active weights per token.""" # parameters() counts tied embedding/head weights only once. total = sum(p.numel() for p in self.parameters()) active_min = active_max = total if self.config["ffn"] != "dense": # Always-active parameters exclude expert/router FFNs; add back each # router plus the minimum/maximum selected expert sizes per token. shared = total - sum( p.numel() for b in self.blocks for p in b.ffn.parameters() ) minima, maxima = [], [] d = self.config["d_model"] for b in self.blocks: widths = b.ffn.widths if b.ffn.kind == "shared": lo, hi = min(widths[:3]) + widths[3], max(widths[:3]) + widths[3] elif b.ffn.kind == "top1": lo, hi = min(widths), max(widths) else: lo, hi = sum(sorted(widths)[:2]), sum(sorted(widths)[-2:]) # Each hidden expert unit has width input and width output weights. # The 4*width router weights are active regardless of selected experts. minima.append(2 * d * lo + 4 * d) maxima.append(2 * d * hi + 4 * d) active_min, active_max = shared + sum(minima), shared + sum(maxima) return {"total": total, "active_min": active_min, "active_max": active_max}