File size: 6,775 Bytes
73919ed | 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 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 | """Llama-style decoder-only transformer.
RMSNorm, rotary position embeddings, SwiGLU MLP, untied input/output
embeddings. Deliberately small and dependency-light (just torch) so it runs
unchanged on the GH200 nodes and on a login-node CPU smoke test.
Configs live in configs/*.yaml; ModelConfig mirrors the yaml `model:` block.
"""
from dataclasses import dataclass
import torch
import torch.nn as nn
import torch.nn.functional as F
@dataclass
class ModelConfig:
vocab_size: int = 65536
dim: int = 2048
n_layers: int = 16
n_heads: int = 16
n_kv_heads: int | None = None # None -> = n_heads (no GQA)
ffn_dim: int = 5632
max_seq_len: int = 2048
rope_theta: float = 10000.0
norm_eps: float = 1e-5
@property
def kv_heads(self) -> int:
return self.n_kv_heads or self.n_heads
@property
def head_dim(self) -> int:
return self.dim // self.n_heads
class RMSNorm(nn.Module):
def __init__(self, dim: int, eps: float):
super().__init__()
self.eps = eps
self.weight = nn.Parameter(torch.ones(dim))
def forward(self, x):
dt = x.dtype
x = x.float()
x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
return (x * self.weight.float()).to(dt)
def _rope_cache(seq_len: int, head_dim: int, theta: float, device, dtype):
inv = 1.0 / (theta ** (torch.arange(0, head_dim, 2, device=device).float() / head_dim))
t = torch.arange(seq_len, device=device).float()
freqs = torch.outer(t, inv) # (T, head_dim/2)
return torch.cos(freqs).to(dtype), torch.sin(freqs).to(dtype)
def _apply_rope(x, cos, sin):
# x: (B, H, T, D). split even/odd halves (rotate-half convention)
x1, x2 = x[..., ::2], x[..., 1::2]
cos = cos[None, None, :, :]
sin = sin[None, None, :, :]
o1 = x1 * cos - x2 * sin
o2 = x1 * sin + x2 * cos
out = torch.empty_like(x)
out[..., ::2] = o1
out[..., 1::2] = o2
return out
class Attention(nn.Module):
def __init__(self, cfg: ModelConfig):
super().__init__()
self.n_heads = cfg.n_heads
self.kv_heads = cfg.kv_heads
self.head_dim = cfg.head_dim
self.wq = nn.Linear(cfg.dim, cfg.n_heads * cfg.head_dim, bias=False)
self.wk = nn.Linear(cfg.dim, cfg.kv_heads * cfg.head_dim, bias=False)
self.wv = nn.Linear(cfg.dim, cfg.kv_heads * cfg.head_dim, bias=False)
self.wo = nn.Linear(cfg.n_heads * cfg.head_dim, cfg.dim, bias=False)
def forward(self, x, cos, sin):
B, T, _ = x.shape
q = self.wq(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
k = self.wk(x).view(B, T, self.kv_heads, self.head_dim).transpose(1, 2)
v = self.wv(x).view(B, T, self.kv_heads, self.head_dim).transpose(1, 2)
q = _apply_rope(q, cos, sin)
k = _apply_rope(k, cos, sin)
if self.kv_heads != self.n_heads:
rep = self.n_heads // self.kv_heads
k = k.repeat_interleave(rep, dim=1)
v = v.repeat_interleave(rep, dim=1)
out = F.scaled_dot_product_attention(q, k, v, is_causal=True)
out = out.transpose(1, 2).contiguous().view(B, T, -1)
return self.wo(out)
class SwiGLU(nn.Module):
def __init__(self, cfg: ModelConfig):
super().__init__()
self.w1 = nn.Linear(cfg.dim, cfg.ffn_dim, bias=False) # gate
self.w3 = nn.Linear(cfg.dim, cfg.ffn_dim, bias=False) # up
self.w2 = nn.Linear(cfg.ffn_dim, cfg.dim, bias=False) # down
def forward(self, x):
return self.w2(F.silu(self.w1(x)) * self.w3(x))
class Block(nn.Module):
def __init__(self, cfg: ModelConfig):
super().__init__()
self.attn_norm = RMSNorm(cfg.dim, cfg.norm_eps)
self.attn = Attention(cfg)
self.ffn_norm = RMSNorm(cfg.dim, cfg.norm_eps)
self.ffn = SwiGLU(cfg)
def forward(self, x, cos, sin):
x = x + self.attn(self.attn_norm(x), cos, sin)
x = x + self.ffn(self.ffn_norm(x))
return x
class Transformer(nn.Module):
def __init__(self, cfg: ModelConfig):
super().__init__()
self.cfg = cfg
self.tok_emb = nn.Embedding(cfg.vocab_size, cfg.dim)
self.layers = nn.ModuleList(Block(cfg) for _ in range(cfg.n_layers))
self.norm = RMSNorm(cfg.dim, cfg.norm_eps)
self.lm_head = nn.Linear(cfg.dim, cfg.vocab_size, bias=False) # untied
self._rope = None
self.apply(self._init)
# scale residual-projection inits by depth (GPT-2/Llama convention)
for name, p in self.named_parameters():
if name.endswith("wo.weight") or name.endswith("w2.weight"):
nn.init.normal_(p, mean=0.0, std=0.02 / (2 * cfg.n_layers) ** 0.5)
def _init(self, m):
if isinstance(m, nn.Linear):
nn.init.normal_(m.weight, mean=0.0, std=0.02)
elif isinstance(m, nn.Embedding):
nn.init.normal_(m.weight, mean=0.0, std=0.02)
def _rope_for(self, T, device, dtype):
if self._rope is None or self._rope[0].shape[0] < T or self._rope[0].device != device:
self._rope = _rope_cache(self.cfg.max_seq_len, self.cfg.head_dim,
self.cfg.rope_theta, device, dtype)
cos, sin = self._rope
return cos[:T], sin[:T]
def forward(self, idx, targets=None):
B, T = idx.shape
x = self.tok_emb(idx)
cos, sin = self._rope_for(T, idx.device, x.dtype)
for layer in self.layers:
x = layer(x, cos, sin)
x = self.norm(x)
if targets is None:
return self.lm_head(x[:, -1:, :])
logits = self.lm_head(x)
loss = F.cross_entropy(logits.view(-1, logits.size(-1)),
targets.reshape(-1), ignore_index=-100)
return logits, loss
@torch.no_grad()
def layer_reps(self, idx):
"""Per-layer hidden states for representation analysis (MEXA).
Returns a tensor (n_layers+1, B, T, dim): index 0 is the embedding
output, index i>=1 is the output of block i. Causal attention means
right-padding never contaminates real positions, so callers can pool
over a length mask safely.
"""
B, T = idx.shape
x = self.tok_emb(idx)
cos, sin = self._rope_for(T, idx.device, x.dtype)
reps = [x]
for layer in self.layers:
x = layer(x, cos, sin)
reps.append(x)
return torch.stack(reps, dim=0)
def num_params(self, embedding: bool = True) -> int:
n = sum(p.numel() for p in self.parameters())
if not embedding:
n -= self.tok_emb.weight.numel() + self.lm_head.weight.numel()
return n
|