File size: 6,939 Bytes
a2ec932 | 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 | """
MiniTransformer — Decoder-only Causal Language Model
──────────────────────────────────────────────────────
Architecture at a glance:
Input tokens (B, T)
│
Token Embedding [vocab_size → d_model]
│
┌────┴──────────────────────────────────────────────────────┐
│ RoPE (cos, sin) — computed once, shared across all layers │
└────┬──────────────────────────────────────────────────────┘
│
┌────┴──────────────────────────┐
│ TransformerBlock × n_layers │ (Pre-LN, RMSNorm, GQA + SDPA, SwiGLU)
└────┬──────────────────────────┘
│
Final RMSNorm
│
LM Head [d_model → vocab_size] (weight-tied to token embedding)
│
Logits (B, T, vocab_size)
WEIGHT TYING:
The token embedding matrix (vocab_size × d_model) and the LM head weight
are shared — the same parameter serves both roles. This:
1. Saves vocab_size × d_model parameters (~4M for 8192 vocab, 512 dim)
2. Enforces consistency: the model learns the same representation for
"what token am I?" and "what token should come next?"
Standard in GPT-2, Llama, and most open-weight models.
WEIGHT INITIALISATION:
All Linear and Embedding weights are initialised with N(0, 0.02) following
GPT-2. This keeps activations in a stable range at the start of training.
Biases (where they exist) are zeroed.
"""
import torch
import torch.nn as nn
from model.config import ModelConfig
from model.block import TransformerBlock, RMSNorm
from model.positional_encoding import RotaryEmbedding
class MiniTransformer(nn.Module):
"""
Decoder-only causal language model.
Args:
config: A ModelConfig dataclass instance. All hyperparameters flow
through this object — never pass raw ints directly.
Usage:
config = ModelConfig(vocab_size=8192, d_model=512, n_layers=6)
model = MiniTransformer(config)
logits, kvs = model(token_ids) # logits: [B, T, vocab_size]
"""
def __init__(self, config: ModelConfig):
super().__init__()
self.config = config
# ── Embeddings ────────────────────────────────────────────────────────
self.token_emb = nn.Embedding(config.vocab_size, config.d_model)
# RoPE: produces (cos, sin) for each position; applied inside attention
self.rope = RotaryEmbedding(
dim = config.head_dim, # operates per attention-head
max_seq_len = config.context_len,
base = config.rope_base,
)
# ── Transformer Layers ───────────────────────────────────────────────
self.blocks = nn.ModuleList([
TransformerBlock(config) for _ in range(config.n_layers)
])
# ── Output ────────────────────────────────────────────────────────────
self.ln_final = RMSNorm(config.d_model, eps=config.norm_eps)
self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)
# Weight tying: lm_head reads the same matrix as token_emb
if config.tie_weights:
self.lm_head.weight = self.token_emb.weight
# Initialise all parameters
self.apply(self._init_weights)
# ── Initialisation ────────────────────────────────────────────────────────
def _init_weights(self, module: nn.Module) -> None:
"""GPT-2 style weight initialisation: N(0, 0.02) for all linear/embedding."""
if isinstance(module, (nn.Linear, nn.Embedding)):
nn.init.normal_(module.weight, mean=0.0, std=0.02)
if isinstance(module, nn.Linear) and module.bias is not None:
nn.init.zeros_(module.bias)
# ── Forward Pass ──────────────────────────────────────────────────────────
def forward(
self,
x: torch.Tensor,
past_kvs: list | None = None,
) -> tuple[torch.Tensor, list]:
"""
Args:
x: Token indices, shape [B, T].
past_kvs: List of (k, v) tuples — one per layer — for KV-cache
inference. Pass None during training or the first
generation step.
Returns:
logits: Shape [B, T, vocab_size].
present_kvs: Updated KV-cache list (one tuple per layer).
Feed this back as `past_kvs` for the next token.
"""
B, T = x.shape
# Token embeddings — no positional embedding added here (RoPE handles it)
x = self.token_emb(x) # [B, T, d_model]
# Build RoPE tables with correct sequence offset for KV-cache decoding
start_pos = 0 if past_kvs is None else past_kvs[0][0].shape[2]
cos, sin = self.rope(T, start_pos=start_pos) # both [1, 1, T, head_dim]
# Pass through each transformer block, collecting KV-cache outputs
present_kvs: list = []
for i, block in enumerate(self.blocks):
past_kv = past_kvs[i] if past_kvs is not None else None
x, present_kv = block(x, cos, sin, past_kv)
present_kvs.append(present_kv)
# Final normalisation and language model head
x = self.ln_final(x) # [B, T, d_model]
logits = self.lm_head(x) # [B, T, vocab_size]
return logits, present_kvs
# ── Helpers ───────────────────────────────────────────────────────────────
@property
def num_parameters(self) -> int:
"""Count of trainable parameters."""
return sum(p.numel() for p in self.parameters() if p.requires_grad)
def __repr__(self) -> str:
return (
f"MiniTransformer | {self.num_parameters:,} params\n"
f" {self.config}"
)
|