| """A decoder-only transformer whose feed-forward block is swappable. |
| |
| Attention, RoPE and RMSNorm are carried over from Assignment 1 and are held |
| fixed across all five ablations -- the feed-forward network is the only thing |
| that changes, which is what makes the comparison mean anything. |
| """ |
|
|
| import math |
|
|
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
| from .config import TransformerConfig |
| from .ffn import build_ffn, MoE |
|
|
|
|
| class RMSNorm(nn.Module): |
| def __init__(self, d_model, eps=1e-6): |
| super().__init__() |
| self.gamma = nn.Parameter(torch.ones(d_model)) |
| self.eps = eps |
|
|
| def forward(self, x): |
| dtype = x.dtype |
| x = x.float() |
| rms = torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) |
| return (self.gamma.float() * x * rms).to(dtype) |
|
|
|
|
| class RotaryEmbedding(nn.Module): |
| """Rotary position embedding, applied to q and k inside attention.""" |
|
|
| def __init__(self, head_dim, max_len, base=10000.0): |
| super().__init__() |
| inv = 1.0 / (base ** (torch.arange(0, head_dim, 2).float() / head_dim)) |
| t = torch.arange(max_len).float() |
| freqs = torch.outer(t, inv) |
| emb = torch.cat((freqs, freqs), dim=-1) |
| self.register_buffer("cos", emb.cos()[None, None], persistent=False) |
| self.register_buffer("sin", emb.sin()[None, None], persistent=False) |
|
|
| @staticmethod |
| def _rotate_half(x): |
| x1, x2 = x.chunk(2, dim=-1) |
| return torch.cat((-x2, x1), dim=-1) |
|
|
| def forward(self, q, k, offset=0): |
| n = q.size(-2) |
| cos = self.cos[..., offset:offset + n, :].to(q.dtype) |
| sin = self.sin[..., offset:offset + n, :].to(q.dtype) |
| q = q * cos + self._rotate_half(q) * sin |
| k = k * cos + self._rotate_half(k) * sin |
| return q, k |
|
|
|
|
| class CausalSelfAttention(nn.Module): |
| def __init__(self, config: TransformerConfig, rope: RotaryEmbedding): |
| super().__init__() |
| self.n_heads = config.n_heads |
| self.n_kv = config.n_kv_heads or config.n_heads |
| assert self.n_heads % self.n_kv == 0 |
| self.n_groups = self.n_heads // self.n_kv |
| self.d_head = config.d_head |
| self.rope = rope |
|
|
| d = config.d_model |
| self.wq = nn.Linear(d, self.n_heads * self.d_head, bias=False) |
| self.wk = nn.Linear(d, self.n_kv * self.d_head, bias=False) |
| self.wv = nn.Linear(d, self.n_kv * self.d_head, bias=False) |
| self.wo = nn.Linear(self.n_heads * self.d_head, d, bias=False) |
| self.dropout = config.dropout |
|
|
| def forward(self, x): |
| B, N, _ = x.shape |
| q = self.wq(x).view(B, N, self.n_heads, self.d_head).transpose(1, 2) |
| k = self.wk(x).view(B, N, self.n_kv, self.d_head).transpose(1, 2) |
| v = self.wv(x).view(B, N, self.n_kv, self.d_head).transpose(1, 2) |
|
|
| q, k = self.rope(q, k) |
| if self.n_groups > 1: |
| k = k.repeat_interleave(self.n_groups, dim=1) |
| v = v.repeat_interleave(self.n_groups, dim=1) |
|
|
| out = F.scaled_dot_product_attention( |
| q, k, v, is_causal=True, |
| dropout_p=self.dropout if self.training else 0.0) |
| out = out.transpose(1, 2).reshape(B, N, -1) |
| return self.wo(out) |
|
|
|
|
| class Block(nn.Module): |
| """Pre-norm decoder block; ``self.ffn`` is the MoE/MLP drop-in point.""" |
|
|
| def __init__(self, config: TransformerConfig, rope: RotaryEmbedding): |
| super().__init__() |
| self.norm1 = RMSNorm(config.d_model) |
| self.attn = CausalSelfAttention(config, rope) |
| self.norm2 = RMSNorm(config.d_model) |
| self.ffn = build_ffn(config) |
|
|
| def forward(self, x): |
| x = x + self.attn(self.norm1(x)) |
| x = x + self.ffn(self.norm2(x)) |
| return x |
|
|
|
|
| class Transformer(nn.Module): |
| def __init__(self, config: TransformerConfig): |
| super().__init__() |
| self.config = config |
| self.embed = nn.Embedding(config.vocab_size, config.d_model) |
| self.rope = RotaryEmbedding(config.d_head, config.n_ctx, config.rope_base) |
| self.blocks = nn.ModuleList( |
| [Block(config, self.rope) for _ in range(config.n_layers)]) |
| self.norm_f = RMSNorm(config.d_model) |
| self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False) |
| if config.tie_embeddings: |
| self.lm_head.weight = self.embed.weight |
| self.apply(self._init) |
| |
| |
| for name, p in self.named_parameters(): |
| if name.endswith("wo.weight") or name.endswith("fc2.weight"): |
| nn.init.normal_(p, std=0.02 / math.sqrt(2 * config.n_layers)) |
|
|
| @staticmethod |
| def _init(m): |
| if isinstance(m, nn.Linear): |
| nn.init.normal_(m.weight, std=0.02) |
| if m.bias is not None: |
| nn.init.zeros_(m.bias) |
| elif isinstance(m, nn.Embedding): |
| nn.init.normal_(m.weight, std=0.02) |
|
|
| def forward(self, idx, targets=None, loss_mask=None): |
| |
| x = self.embed(idx) |
| for block in self.blocks: |
| x = block(x) |
| logits = self.lm_head(self.norm_f(x)) |
|
|
| if targets is None: |
| return logits, None, self.aux_losses() |
|
|
| loss = F.cross_entropy( |
| logits.view(-1, logits.size(-1)).float(), targets.reshape(-1), |
| ignore_index=-100, reduction="none") |
| if loss_mask is not None: |
| m = loss_mask.reshape(-1).float() |
| loss = (loss * m).sum() / m.sum().clamp_min(1.0) |
| else: |
| loss = loss[targets.reshape(-1) != -100].mean() |
| return logits, loss, self.aux_losses() |
|
|
| def aux_losses(self): |
| """Summed router losses over every MoE layer (zero for the dense model).""" |
| aux = torch.zeros((), device=self.embed.weight.device) |
| z = torch.zeros((), device=self.embed.weight.device) |
| for block in self.blocks: |
| if isinstance(block.ffn, MoE): |
| aux = aux + block.ffn.aux_loss |
| z = z + block.ffn.z_loss |
| return aux, z |
|
|
| def expert_counts(self): |
| """Per-layer token counts from the last forward, or None if dense.""" |
| out = [] |
| for block in self.blocks: |
| if isinstance(block.ffn, MoE): |
| out.append(block.ffn.last_counts) |
| return out or None |
|
|
| def expert_assignments(self): |
| """Per-layer (B, N, k) expert ids from the last forward.""" |
| out = [] |
| for block in self.blocks: |
| if isinstance(block.ffn, MoE): |
| out.append(block.ffn.last_assign) |
| return out or None |
|
|
| def num_params(self, non_embedding=False): |
| n = sum(p.numel() for p in self.parameters()) |
| if non_embedding: |
| n -= self.embed.weight.numel() |
| return n |
|
|
