File size: 5,382 Bytes
bec4ff4 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 | """Decoder-only transformer with a pluggable FFN (dense MLP or MoE variants)."""
from __future__ import annotations
import math
import torch
import torch.nn as nn
import torch.nn.functional as F
from .config import TransformerConfig
from .ffn import MoE, build_ffn
def apply_rope(x: torch.Tensor, positions: torch.Tensor, base: float) -> torch.Tensor:
# x: (B, H, T, hd), positions: (B, T)
hd = x.shape[-1]
inv_freq = 1.0 / (base ** (torch.arange(0, hd, 2, device=x.device).float() / hd))
ang = positions[:, None, :, None].float() * inv_freq # (B, 1, T, hd/2)
cos, sin = ang.cos().to(x.dtype), ang.sin().to(x.dtype)
x1, x2 = x[..., 0::2], x[..., 1::2]
return torch.stack((x1 * cos - x2 * sin, x1 * sin + x2 * cos), dim=-1).flatten(-2)
class Attention(nn.Module):
def __init__(self, config: TransformerConfig):
super().__init__()
assert config.D_model % config.n_heads == 0
self.n_heads = config.n_heads
self.head_dim = config.D_model // config.n_heads
self.rope_base = config.rope_base
self.qkv = nn.Linear(config.D_model, 3 * config.D_model, bias=False)
self.proj = nn.Linear(config.D_model, config.D_model, bias=False)
self.dropout = config.dropout
def forward(self, x, positions, attn_mask=None):
B, T, D = x.shape
q, k, v = self.qkv(x).view(B, T, 3, self.n_heads, self.head_dim).permute(2, 0, 3, 1, 4)
q, k = apply_rope(q, positions, self.rope_base), apply_rope(k, positions, self.rope_base)
y = F.scaled_dot_product_attention(
q, k, v,
attn_mask=attn_mask,
dropout_p=self.dropout if self.training else 0.0,
is_causal=attn_mask is None,
)
return self.proj(y.transpose(1, 2).reshape(B, T, D))
class Block(nn.Module):
def __init__(self, config: TransformerConfig):
super().__init__()
self.ln1 = nn.RMSNorm(config.D_model)
self.attn = Attention(config)
self.ln2 = nn.RMSNorm(config.D_model)
self.ffn = build_ffn(config)
def forward(self, x, positions, attn_mask=None, pad_mask=None):
x = x + self.attn(self.ln1(x), positions, attn_mask)
return x + self.ffn(self.ln2(x), pad_mask)
class Transformer(nn.Module):
def __init__(self, config: TransformerConfig):
super().__init__()
self.config = config
self.tok_emb = nn.Embedding(config.N_vocab, config.D_model)
self.blocks = nn.ModuleList(Block(config) for _ in range(config.n_layers))
self.ln_f = nn.RMSNorm(config.D_model)
self.lm_head = nn.Linear(config.D_model, config.N_vocab, bias=False)
self.lm_head.weight = self.tok_emb.weight # weight tying
self.apply(self._init)
for name, p in self.named_parameters(): # GPT-2 style residual scaling
if name.endswith(("attn.proj.weight", "down.weight")): # incl. experts
nn.init.normal_(p, std=0.02 / math.sqrt(2 * config.n_layers))
@staticmethod
def _init(m):
if isinstance(m, (nn.Linear, nn.Embedding)):
nn.init.normal_(m.weight, std=0.02)
def forward(self, ids, pad_mask=None, left_padded=False):
# Input shape: (B, N_ctx) token ids
# Output shape: (B, N_ctx, N_vocab)
# pad_mask: (B, N_ctx) bool, True for real tokens. It is always handed to
# the FFN (MoE ignores padding in its aux loss / usage stats). With right
# padding (training/eval) plain causal attention is exact, since real
# tokens never attend to trailing pads. For left-padded batches
# (generation) set left_padded=True to mask pad keys and shift positions.
B, T = ids.shape
if not left_padded:
positions = torch.arange(T, device=ids.device).expand(B, T)
attn_mask = None
else:
positions = (pad_mask.long().cumsum(-1) - 1).clamp(min=0)
causal = torch.ones(T, T, dtype=torch.bool, device=ids.device).tril()
eye = torch.eye(T, dtype=torch.bool, device=ids.device)
# `eye` keeps pad queries from having an all-masked row (NaNs)
attn_mask = (causal & pad_mask[:, None, None, :]) | eye # (B, 1, T, T)
x = self.tok_emb(ids)
for blk in self.blocks:
x = blk(x, positions, attn_mask, pad_mask)
return self.lm_head(self.ln_f(x))
# ---- MoE helpers ----
def moe_layers(self) -> list[MoE]:
return [b.ffn for b in self.blocks if isinstance(b.ffn, MoE)]
def aux_loss(self) -> torch.Tensor:
layers = self.moe_layers()
if not layers:
return torch.zeros((), device=self.tok_emb.weight.device)
return sum(m.aux_loss for m in layers) / len(layers)
def set_record_routing(self, flag: bool) -> None:
for m in self.moe_layers():
m.record_routing = flag
def param_counts(self) -> dict[str, int]:
"""total / active parameters (active = per-token, excluding unused experts)."""
from src.utils import count_params
total = count_params(self)
ffn_total = sum(count_params(b.ffn) for b in self.blocks)
ffn_active = sum(b.ffn.active_params() for b in self.blocks)
return {"total": total, "active": total - ffn_total + ffn_active,
"ffn_total": ffn_total, "ffn_active": ffn_active}
|