""" Stage 4 — model.py A full encoder-decoder transformer, in the shape of "Attention Is All You Need". Design choice: the encoder and decoder *layers* are hand-written so you can see exactly which tensor is the query and which are the keys/values in each of the three attention calls — that is the whole point of this project. The attention operation itself uses nn.MultiheadAttention rather than a hand-rolled softmax(QK^T/sqrt(d))V, because that part you have already studied and it is where subtle numerical bugs hide. Swap it out later if you want to write the scaled dot-product by hand; the surrounding code will not change. THE THREE ATTENTION CALLS ------------------------- 1. Encoder self-attention Q,K,V = source NOT causal (bidirectional) 2. Decoder self-attention Q,K,V = target CAUSAL 3. Decoder cross-attention Q = target K,V = encoder output NOT causal Call 3 is the new piece. The decoder asks, at each target position, "which source tokens are relevant to the word I am about to produce?" — so the query comes from the decoder and the keys/values come from the encoder's final output (`memory`). It is NOT causally masked in either direction: target position 0 is allowed to look at the whole source sentence. Only *target* positions are restricted, and only in call 2. MASK CONVENTIONS (kept identical everywhere on purpose) ------------------------------------------------------- key_padding_mask : (B, S) bool. True = that key position is , ignore it. attn_mask : (T, T) bool. True = that (query, key) pair is FORBIDDEN. Both are "True means blocked". PyTorch uses this polarity for nn.MultiheadAttention; flipping it silently trains a model that attends only to padding, which does not crash and does not produce NaNs — it just produces garbage. Hence: one convention, stated once, used everywhere. Padding must be masked on BOTH sides: - source padding -> encoder self-attention keys (call 1) - source padding -> cross-attention keys (call 3) <- easy to forget - target padding -> decoder self-attention keys (call 2) Pre-LN vs post-LN: the 2017 paper puts LayerNorm *after* each residual addition. Pre-LN (norm on the branch input, as written below) is what most modern implementations use because it trains stably without a carefully tuned warmup. We still use warmup, but this makes the model far more forgiving if you change the learning rate. Each stack gets a final LayerNorm, which pre-LN requires. ============================================================================ THIS FILE'S PLACE IN THE CHAIN ============================================================================ model.py has ONE project-internal dependency: config.py (for ModelConfig). It does not know about tokenizers, datasets, or files on disk — it is pure architecture, operating on already-tokenized integer id tensors. Who calls into this file: - train.py's main(): cfg = ModelConfig(...); model = build_model(cfg, pad_id, device) then, every step, run_epoch() calls model(src, tgt_in, ...) which routes to TransformerTranslator.forward() below. - decoding.py's greedy_decode()/beam_search_decode(): call model.encode(...) ONCE, then model.decode(...) repeatedly in a loop — this is why encode() and decode() are exposed as separate methods instead of only forward(). - evaluate.py's main(): build_model(cfg, ...) again from a saved ModelConfig, then model.load_state_dict(ckpt["model"]) to restore trained weights. - sanity_checks.py: imports build_model and causal_mask directly to probe the model's masking behavior with hand-crafted inputs. Internal call chain within this file (top to bottom of what actually runs during a training step): TransformerTranslator.forward(src, tgt_in, ...) -> self.encode(src, ...) -> self.embed(src) * sqrt(d_model) [token ids -> vectors] -> self.pos_enc(x) [add position information] -> for each layer in self.encoder_layers: EncoderLayer.forward(x, ...) -> ATTENTION CALL 1 (self.self_attn, not masked causally) -> self.ff(x) (FeedForward) -> self.enc_norm(x) [final pre-LN norm] => returns `memory`, shape (B, S, d) -> self.decode(tgt_in, memory, ...) -> self.embed(tgt_in) * sqrt(d_model) [SAME embedding table as above] -> self.pos_enc(x) -> cm = causal_mask(T, device) [built once per decode() call] -> for each layer in self.decoder_layers: DecoderLayer.forward(x, memory, cm, ...) -> ATTENTION CALL 2 (self.self_attn, WITH causal_mask) -> ATTENTION CALL 3 (self.cross_attn, query=decoder, key/value=memory) -> self.ff(x) -> self.dec_norm(x) -> self.output_proj(x) [-> logits over the vocab] """ import math import torch import torch.nn as nn from config import ModelConfig # --------------------------------------------------------------------------- # Positional encodings # --------------------------------------------------------------------------- class SinusoidalPositionalEncoding(nn.Module): """Fixed sin/cos encodings from the original paper. PE[pos, 2i] = sin(pos / 10000^(2i/d)) PE[pos, 2i+1] = cos(pos / 10000^(2i/d)) No parameters, and it extrapolates (sort of) to lengths never seen in training — which matters here because translation output length is not bounded by anything the model saw. Instantiated once inside TransformerTranslator.__init__ (via build_positional_encoding below) and called from BOTH encode() and decode() — the SAME module instance handles source and target positions, since "position 3" means the same thing (a fixed sin/cos vector) in either sequence. """ def __init__(self, d_model, max_len=512, dropout=0.1): super().__init__() self.dropout = nn.Dropout(dropout) # Precompute the ENTIRE table up to max_len positions, once, at # construction time — forward() below just slices into it. pe = torch.zeros(max_len, d_model) position = torch.arange(max_len, dtype=torch.float).unsqueeze(1) # (max_len, 1) # Computed in log space for numerical stability. # div[i] = 10000^(-2i/d_model), one value per even dimension index. div = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)) # Broadcasting: position (max_len,1) * div (d_model/2,) -> (max_len, d_model/2) pe[:, 0::2] = torch.sin(position * div) # even dims: sin pe[:, 1::2] = torch.cos(position * div) # odd dims: cos # register_buffer: saved with the model, moved by .to(device), not a parameter. # (i.e. it's part of state_dict()/checkpoint, and follows model.to(device), # but torch.optim never updates it — it's fixed math, not learned.) self.register_buffer("pe", pe.unsqueeze(0)) # (1, max_len, d) def forward(self, x): # x: (B, L, d) # Slice the precomputed table down to this sequence's actual length L, # broadcast-add across the batch dimension, then apply dropout. # Called from TransformerTranslator.encode()/decode() right after the # scaled embedding lookup — see the module docstring's call chain. return self.dropout(x + self.pe[:, : x.size(1)]) class LearnedPositionalEncoding(nn.Module): """An ordinary embedding table indexed by position. Ablation 6. Selected via ModelConfig.pos_encoding="learned" instead of "sinusoidal" — see build_positional_encoding() below, which is the only place that chooses between this class and SinusoidalPositionalEncoding above. Unlike the sinusoidal version, this one has trainable parameters (self.emb.weight) and CANNOT extrapolate past max_len (see the assert). """ def __init__(self, d_model, max_len=512, dropout=0.1): super().__init__() self.dropout = nn.Dropout(dropout) self.emb = nn.Embedding(max_len, d_model) # one learned vector per position 0..max_len-1 nn.init.normal_(self.emb.weight, mean=0.0, std=0.02) self.max_len = max_len def forward(self, x): L = x.size(1) assert L <= self.max_len, f"sequence of {L} exceeds learned max_len {self.max_len}" # arange(L) -> position indices [0, 1, ..., L-1], broadcast to batch # dim 1 so nn.Embedding can look each one up. pos = torch.arange(L, device=x.device).unsqueeze(0) # (1, L) return self.dropout(x + self.emb(pos)) def build_positional_encoding(kind, d_model, max_len, dropout): """Factory function: picks which positional-encoding class to instantiate based on ModelConfig.pos_encoding. Called exactly once, from TransformerTranslator.__init__, to build self.pos_enc.""" if kind == "sinusoidal": return SinusoidalPositionalEncoding(d_model, max_len, dropout) if kind == "learned": return LearnedPositionalEncoding(d_model, max_len, dropout) raise ValueError(f"unknown pos_encoding: {kind}") # --------------------------------------------------------------------------- # Layers # --------------------------------------------------------------------------- class FeedForward(nn.Module): """Position-wise FFN: d_model -> d_ff -> d_model, applied identically at every position (it is where most of the parameters live). "Position-wise" means: the SAME two Linear layers are applied independently to every one of the L positions in (B, L, d_model) — there is no mixing ACROSS positions here (that's what attention is for). This is why nn.Linear works directly on a 3D tensor: PyTorch applies it to the last dimension only, batching over everything else automatically. Instantiated once per EncoderLayer and once per DecoderLayer below (self.ff), so with 4+4 default layers there are 8 independent FeedForward modules, each with its own weights. """ def __init__(self, d_model, d_ff, dropout): super().__init__() self.net = nn.Sequential( nn.Linear(d_model, d_ff), # expand nn.ReLU(), # nonlinearity — this is what makes the FFN # do more than a linear projection nn.Dropout(dropout), nn.Linear(d_ff, d_model), # project back down to d_model so residual add works ) def forward(self, x): # Called from EncoderLayer.forward() and DecoderLayer.forward(), each # time wrapped in a residual connection: x = x + self.ff(norm(x)). return self.net(x) class EncoderLayer(nn.Module): """self-attention (bidirectional) -> feed-forward, both residual + pre-LN. Stacked cfg.n_encoder_layers times inside TransformerTranslator.__init__ (self.encoder_layers). TransformerTranslator.encode() runs input `x` through each of these in sequence, output of one feeding the next. """ def __init__(self, d_model, n_heads, d_ff, dropout): super().__init__() # nn.MultiheadAttention does the actual softmax(QK^T/sqrt(d))V math # (scaled dot-product attention across n_heads parallel heads) — # see the module docstring for why this project doesn't hand-roll it. # batch_first=True: tensors are (batch, seq, feature), matching every # other shape convention in this file (the PyTorch default is # (seq, batch, feature), which would be inconsistent here). self.self_attn = nn.MultiheadAttention(d_model, n_heads, dropout=dropout, batch_first=True) self.ff = FeedForward(d_model, d_ff, dropout) self.norm1 = nn.LayerNorm(d_model) # pre-LN before the attention sublayer self.norm2 = nn.LayerNorm(d_model) # pre-LN before the feed-forward sublayer self.drop = nn.Dropout(dropout) # applied to each sublayer's OUTPUT before the residual add def forward(self, x, src_key_padding_mask=None): # x: (B, S, d) — called from TransformerTranslator.encode(), once per # layer in self.encoder_layers, x is this layer's input AND (after # this function returns) the next layer's input. # ATTENTION CALL 1: query = key = value = the source itself. # No attn_mask -> every source token sees every other source token in # both directions. That is the point of the encoder. h = self.norm1(x) # pre-LN: normalize BEFORE feeding into attention attn, _ = self.self_attn(h, h, h, # (query, key, value) — all three are `h` key_padding_mask=src_key_padding_mask, # ^ (B, S) bool, True = ignore this source # position (it's ) — this is # dataset.py's "src_key_padding_mask" batch # key, passed down unchanged from # TransformerTranslator.forward()/encode(). need_weights=False) # we never inspect attention # weights, so skip computing them (speed) x = x + self.drop(attn) # residual connection: original x + (dropped) attention output h = self.norm2(x) # pre-LN before the FFN sublayer x = x + self.drop(self.ff(h)) # second residual connection return x class DecoderLayer(nn.Module): """causal self-attention -> cross-attention -> feed-forward. Stacked cfg.n_decoder_layers times inside TransformerTranslator.__init__ (self.decoder_layers). TransformerTranslator.decode() runs decoder state `x` AND the fixed `memory` (encoder output) through each of these in sequence — `memory` itself is NOT modified by the decoder stack, only `x` is. """ def __init__(self, d_model, n_heads, d_ff, dropout): super().__init__() self.self_attn = nn.MultiheadAttention(d_model, n_heads, dropout=dropout, batch_first=True) # ^ ATTENTION CALL 2's engine — target attends to target, causally masked below. self.cross_attn = nn.MultiheadAttention(d_model, n_heads, dropout=dropout, batch_first=True) # ^ ATTENTION CALL 3's engine — target attends to encoder memory. # A SEPARATE nn.MultiheadAttention instance from self_attn above — # different learned weights, even though the shapes look similar. self.ff = FeedForward(d_model, d_ff, dropout) self.norm1 = nn.LayerNorm(d_model) # before self-attention self.norm2 = nn.LayerNorm(d_model) # before cross-attention self.norm3 = nn.LayerNorm(d_model) # before feed-forward self.drop = nn.Dropout(dropout) def forward(self, x, memory, causal_mask=None, tgt_key_padding_mask=None, memory_key_padding_mask=None): # x : (B, T, d) decoder states (the target so far) # memory : (B, S, d) encoder output (the whole source, always) # Called from TransformerTranslator.decode(), once per layer in # self.decoder_layers, with the SAME `memory` tensor passed to every # layer (memory is computed once by encode() and reused). # ATTENTION CALL 2 — decoder self-attention, CAUSAL. # attn_mask is the (T, T) upper-triangular block: position i may not # look at any j > i. Without it, the model can read the answer during # teacher forcing and will look perfect in training and useless at # inference, where the future genuinely does not exist yet. h = self.norm1(x) sa, _ = self.self_attn(h, h, h, # query=key=value=decoder states attn_mask=causal_mask, # ^ (T, T) bool — built by the standalone # causal_mask() function (below in this # file) and passed down from # TransformerTranslator.decode(). Same # causal_mask tensor is reused for every # decoder layer within one decode() call. key_padding_mask=tgt_key_padding_mask, # ^ (B, T) bool — masks positions WITHIN # the target itself (relevant during # training with a padded batch; None during # greedy/beam generation, where there are no # pads yet in the growing prefix). need_weights=False) x = x + self.drop(sa) # ATTENTION CALL 3 — CROSS-ATTENTION. The new piece. # query <- decoder states (T positions: "what am I writing?") # key <- encoder memory (S positions: "what's available?") # value <- encoder memory (S positions: "what do I copy across?") # Deliberately NO attn_mask: target position 0 may attend to the whole # source, including its last word. Word order differs between Russian # and English, so restricting this would make translation impossible. # key_padding_mask IS passed: in the source must not be attended # to. Forgetting it does not crash — it quietly costs a few BLEU. h = self.norm2(x) ca, _ = self.cross_attn(query=h, key=memory, value=memory, key_padding_mask=memory_key_padding_mask, # ^ (B, S) bool — this is the SAME # src_key_padding_mask used in encoder call # 1, just renamed "memory_key_padding_mask" # here because from the decoder's point of # view it's masking the KEYS (memory), not # its own sequence. need_weights=False) x = x + self.drop(ca) h = self.norm3(x) x = x + self.drop(self.ff(h)) return x # --------------------------------------------------------------------------- # Full model # --------------------------------------------------------------------------- def causal_mask(size, device): """(size, size) bool, True = forbidden. Row i has True for all columns j>i. [[F, T, T], [F, F, T], [F, F, F]] Called from TransformerTranslator.decode() every time it runs (built fresh from the CURRENT target length T each call — cheap, so no need to cache it). Also called directly by sanity_checks.py's test_causal_mask_shape() to verify this exact shape/polarity in isolation. torch.triu(..., diagonal=1): keeps only the STRICTLY upper triangle (diagonal=1 excludes the main diagonal itself), which is exactly "j > i". That's why position i CAN attend to itself (j == i is False, i.e. allowed) but not to anything after it. """ return torch.triu(torch.ones(size, size, dtype=torch.bool, device=device), diagonal=1) class TransformerTranslator(nn.Module): """The whole model. Constructed via build_model() below (never instantiated directly outside this file/sanity_checks.py).""" def __init__(self, cfg: ModelConfig, pad_id=0): super().__init__() self.cfg = cfg self.pad_id = pad_id self.d_model = cfg.d_model # ONE embedding table, shared by: # - the encoder input (Russian tokens) # - the decoder input (English tokens) # - the output projection (via weight tying, below) # This is only possible because the BPE vocabulary is shared. # padding_idx=pad_id: tells nn.Embedding that row `pad_id` should # never receive a gradient update (kept pinned at whatever it's # initialized to — see the explicit zeroing right below). self.embed = nn.Embedding(cfg.vocab_size, cfg.d_model, padding_idx=pad_id) nn.init.normal_(self.embed.weight, mean=0.0, std=cfg.d_model ** -0.5) # ^ std=d_model^-0.5 matches the scale expected by the # "* sqrt(d_model)" multiplication done in encode()/decode() below # — see the comment there for why. with torch.no_grad(): self.embed.weight[pad_id].zero_() # belt-and-suspenders: 's row starts at exactly 0 # Builds either SinusoidalPositionalEncoding or # LearnedPositionalEncoding depending on cfg.pos_encoding — see # build_positional_encoding() above. self.pos_enc = build_positional_encoding(cfg.pos_encoding, cfg.d_model, cfg.max_len, cfg.dropout) # nn.ModuleList: a Python-list-like container that PyTorch still # recognizes as holding submodules (a plain Python list would NOT # register these layers' parameters, and .to(device)/state_dict() # would silently miss them). self.encoder_layers = nn.ModuleList([ EncoderLayer(cfg.d_model, cfg.n_heads, cfg.d_ff, cfg.dropout) for _ in range(cfg.n_encoder_layers)]) self.decoder_layers = nn.ModuleList([ DecoderLayer(cfg.d_model, cfg.n_heads, cfg.d_ff, cfg.dropout) for _ in range(cfg.n_decoder_layers)]) # Pre-LN needs a final norm at the top of each stack, otherwise the # output of the last residual branch is never normalized. self.enc_norm = nn.LayerNorm(cfg.d_model) # applied at the end of encode() self.dec_norm = nn.LayerNorm(cfg.d_model) # applied at the end of decode(), before output_proj self.output_proj = nn.Linear(cfg.d_model, cfg.vocab_size, bias=False) # ^ turns each decoder position's d_model-dim vector into a # vocab_size-dim vector of LOGITS (unnormalized scores per token) — # this is what decode()'s return value actually is. if cfg.tie_embeddings: # Weight tying: the output layer reuses the embedding matrix, so # logit[v] = . Saves vocab*d_model # parameters (16000*256 = 4.1M here, a large share of the model) # and usually helps a low-resource model generalize. # NOTE: this line REPLACES self.output_proj.weight with a # reference to the SAME tensor object as self.embed.weight — from # this point on, training self.embed.weight via backprop also # changes self.output_proj's behavior, and vice versa, because # they are literally one tensor with two names. self.output_proj.weight = self.embed.weight # -- encoder ------------------------------------------------------------ def encode(self, src, src_key_padding_mask=None): """src: (B, S) -> memory: (B, S, d) Called from: TransformerTranslator.forward() below (training path), AND directly by decoding.py's greedy_decode()/beam_search_decode() (generation path) — those call it ONCE per source sentence, then reuse the returned `memory` across every decode() step, which is the whole computational payoff of splitting encoder/decoder into separate methods instead of only exposing forward(). """ # Scale by sqrt(d_model) as in the paper: embeddings are initialized # with std d^-0.5, so this puts them on the same scale as the # positional encoding (which has unit-ish amplitude) rather than being # swamped by it. x = self.embed(src) * math.sqrt(self.d_model) # (B, S) ids -> (B, S, d) vectors, rescaled x = self.pos_enc(x) # add positional information, (B, S, d) for layer in self.encoder_layers: x = layer(x, src_key_padding_mask=src_key_padding_mask) # ^ EncoderLayer.forward() — see that class for ATTENTION CALL 1. # Output of layer i becomes input to layer i+1. return self.enc_norm(x) # final pre-LN normalization required by the pre-LN convention # -- decoder ------------------------------------------------------------ def decode(self, tgt_in, memory, memory_key_padding_mask=None, tgt_key_padding_mask=None): """tgt_in: (B, T) decoder INPUT (already shifted, starts with ) memory: (B, S, d) -> logits (B, T, vocab) Called from: TransformerTranslator.forward() below (training path, ONE call covering the whole target sequence at once thanks to teacher forcing + the causal mask), AND from decoding.py's greedy_decode()/beam_search_decode() in a LOOP — each generation step re-runs decode() over the WHOLE prefix generated so far (see decoding.py's module docstring "Efficiency note" for why this is O(T^2) and deliberately not KV-cached). """ T = tgt_in.size(1) x = self.embed(tgt_in) * math.sqrt(self.d_model) # SAME embedding table as encode() used x = self.pos_enc(x) # SAME positional-encoding module as encode() cm = causal_mask(T, tgt_in.device) # (T, T), built fresh each call — see causal_mask() above for layer in self.decoder_layers: x = layer(x, memory, causal_mask=cm, tgt_key_padding_mask=tgt_key_padding_mask, memory_key_padding_mask=memory_key_padding_mask) # ^ DecoderLayer.forward() — see that class for ATTENTION CALLS 2 & 3. # `memory` is passed unchanged to every layer; only `x` accumulates. x = self.dec_norm(x) # final pre-LN normalization return self.output_proj(x) # (B, T, d) -> (B, T, vocab_size) logits # -- both --------------------------------------------------------------- def forward(self, src, tgt_in, src_key_padding_mask=None, tgt_key_padding_mask=None): """Training forward pass. src : (B, S) source token ids tgt_in : (B, T) decoder input = target shifted right ( w1 w2 w3) returns: (B, T, vocab) logits, where logits[:, i] predicts the token that should follow tgt_in[:, i] Called from train.py's run_epoch()/validate_loss() as `logits = model(src, tgt_in, src_key_padding_mask=src_mask, tgt_key_padding_mask=tgt_mask)` — this is the ONLY place forward() is used; generation (decoding.py) calls encode()/decode() directly instead, precisely to avoid recomputing `memory` at every generation step (see encode()'s docstring). """ memory = self.encode(src, src_key_padding_mask) return self.decode(tgt_in, memory, memory_key_padding_mask=src_key_padding_mask, # ^ reuses the SAME mask tensor for both "encoder # self-attn keys" (inside encode(), call 1) and # "cross-attn keys" (here, call 3) — both are # masking the same source padding pattern. tgt_key_padding_mask=tgt_key_padding_mask) def num_parameters(self, trainable_only=True): """Called by train.py/evaluate.py purely for the printed "X trainable parameters" log line.""" ps = self.parameters() return sum(p.numel() for p in ps if p.requires_grad or not trainable_only) def build_model(cfg: ModelConfig, pad_id=0, device="cpu"): """Thin factory wrapping TransformerTranslator(cfg).to(device). Called from: train.py's main(), evaluate.py's main(), sanity_checks.py's _model(), and this file's own __main__ block below. Centralizing construction+device-placement here means every caller gets identical behavior instead of each remembering to call .to(device) itself. """ model = TransformerTranslator(cfg, pad_id=pad_id).to(device) return model # --------------------------------------------------------------------------- # Quick self-check: python model.py # --------------------------------------------------------------------------- if __name__ == "__main__": # Runs only when this file is executed directly. Builds a tiny toy model # (small vocab/d_model so it's instant) and runs one forward pass with # RANDOM token ids (no real tokenizer/data needed) just to confirm the # architecture is wired correctly end to end — shapes match, weight # tying is really the same tensor. For actual correctness properties # (causality, mask polarity) see sanity_checks.py instead. torch.manual_seed(0) cfg = ModelConfig(vocab_size=500, d_model=64, n_heads=4, n_encoder_layers=2, n_decoder_layers=2, d_ff=128, dropout=0.0) m = build_model(cfg) m.eval() # disables dropout for a deterministic check B, S, T = 3, 7, 5 src = torch.randint(4, 500, (B, S)) # random "sentences", ids 4..499 (avoiding specials 0-3) tgt_in = torch.randint(4, 500, (B, T)) logits = m(src, tgt_in) # exercises forward() -> encode() -> decode() print("logits:", tuple(logits.shape), "(expected", (B, T, cfg.vocab_size), ")") print("parameters:", f"{m.num_parameters():,}") print("embedding tied to output:", m.output_proj.weight is m.embed.weight) # ^ `is` checks OBJECT IDENTITY, not just equal values — confirms the # weight-tying assignment in __init__ really did share one tensor. print("\nRun `python sanity_checks.py` for the causality / masking tests.")