import math import torch import torch.nn as nn import torch.nn.functional as F from config_300M import ModelConfig class SinusoidalPositionalEncoding(nn.Module): def __init__(self, d_model: int, max_len: int = 2048): super().__init__() pe = torch.zeros(max_len, d_model) pos = torch.arange(0, max_len).unsqueeze(1).float() div = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)) pe[:, 0::2] = torch.sin(pos * div) pe[:, 1::2] = torch.cos(pos * div) self.register_buffer("pe", pe.unsqueeze(0)) # (1, max_len, d_model) def forward(self, x: torch.Tensor) -> torch.Tensor: return x + self.pe[:, :x.size(1)] class MultiHeadAttention(nn.Module): def __init__(self, d_model: int, n_heads: int, dropout: float = 0.1): super().__init__() assert d_model % n_heads == 0 self.d_k = d_model // n_heads self.n_heads = n_heads self.W_q = nn.Linear(d_model, d_model) self.W_k = nn.Linear(d_model, d_model) self.W_v = nn.Linear(d_model, d_model) self.W_o = nn.Linear(d_model, d_model) self.dropout = nn.Dropout(dropout) def forward(self, q, k, v, mask=None): B, T, _ = q.shape q = self.W_q(q).view(B, -1, self.n_heads, self.d_k).transpose(1, 2) k = self.W_k(k).view(B, -1, self.n_heads, self.d_k).transpose(1, 2) v = self.W_v(v).view(B, -1, self.n_heads, self.d_k).transpose(1, 2) out = F.scaled_dot_product_attention(q, k, v, attn_mask=mask, dropout_p=self.dropout.p if self.training else 0.0) out = out.transpose(1, 2).contiguous().view(B, -1, self.n_heads * self.d_k) return self.W_o(out) class FeedForward(nn.Module): def __init__(self, d_model: int, d_ff: int, dropout: float = 0.1): super().__init__() self.w1 = nn.Linear(d_model, d_ff) self.w2 = nn.Linear(d_ff, d_model) self.dropout = nn.Dropout(dropout) def forward(self, x): return self.w2(self.dropout(F.gelu(self.w1(x)))) class EncoderLayer(nn.Module): def __init__(self, d_model: int, n_heads: int, d_ff: int, dropout: float): super().__init__() self.self_attn = MultiHeadAttention(d_model, n_heads, dropout) self.ff = FeedForward(d_model, d_ff, dropout) self.norm1 = nn.LayerNorm(d_model) self.norm2 = nn.LayerNorm(d_model) self.dropout = nn.Dropout(dropout) def forward(self, x, src_mask=None): x = x + self.dropout(self.self_attn(self.norm1(x), self.norm1(x), self.norm1(x), src_mask)) x = x + self.dropout(self.ff(self.norm2(x))) return x class DecoderLayer(nn.Module): def __init__(self, d_model: int, n_heads: int, d_ff: int, dropout: float): super().__init__() self.self_attn = MultiHeadAttention(d_model, n_heads, dropout) self.cross_attn = MultiHeadAttention(d_model, n_heads, dropout) self.ff = FeedForward(d_model, d_ff, dropout) self.norm1 = nn.LayerNorm(d_model) self.norm2 = nn.LayerNorm(d_model) self.norm3 = nn.LayerNorm(d_model) self.dropout = nn.Dropout(dropout) def forward(self, x, enc_out, src_mask=None, tgt_mask=None): x = x + self.dropout(self.self_attn(self.norm1(x), self.norm1(x), self.norm1(x), tgt_mask)) n2 = self.norm2(x) x = x + self.dropout(self.cross_attn(n2, enc_out, enc_out, src_mask)) x = x + self.dropout(self.ff(self.norm3(x))) return x class TranslationModel(nn.Module): def __init__(self, cfg: ModelConfig): super().__init__() self.cfg = cfg # shared embedding (encoder + decoder + output projection) self.shared_emb = nn.Embedding(cfg.vocab_size, cfg.d_model, padding_idx=cfg.pad_id) self.emb_scale = math.sqrt(cfg.d_model) self.pos_enc = SinusoidalPositionalEncoding(cfg.d_model, cfg.max_seq_len) self.emb_dropout = nn.Dropout(cfg.dropout) # language tag embeddings (one per supported language) self.lang_emb = nn.Embedding(cfg.n_langs, cfg.d_model) self.encoder = nn.ModuleList([ EncoderLayer(cfg.d_model, cfg.n_heads, cfg.d_ff, cfg.dropout) for _ in range(cfg.encoder_layers) ]) self.encoder_norm = nn.LayerNorm(cfg.d_model) self.decoder = nn.ModuleList([ DecoderLayer(cfg.d_model, cfg.n_heads, cfg.d_ff, cfg.dropout) for _ in range(cfg.decoder_layers) ]) self.decoder_norm = nn.LayerNorm(cfg.d_model) # output projection shares weights with embedding self.output_proj = nn.Linear(cfg.d_model, cfg.vocab_size, bias=False) self.output_proj.weight = self.shared_emb.weight self._init_weights() def _init_weights(self): for p in self.parameters(): if p.dim() > 1: nn.init.xavier_uniform_(p) def encode(self, src_ids, src_lang_id, src_mask=None): x = self.shared_emb(src_ids) * self.emb_scale x = x + self.lang_emb(src_lang_id).unsqueeze(1) x = self.emb_dropout(self.pos_enc(x)) for layer in self.encoder: x = layer(x, src_mask) return self.encoder_norm(x) def decode(self, tgt_ids, enc_out, tgt_lang_id, src_mask=None, tgt_mask=None): x = self.shared_emb(tgt_ids) * self.emb_scale x = x + self.lang_emb(tgt_lang_id).unsqueeze(1) x = self.emb_dropout(self.pos_enc(x)) for layer in self.decoder: x = layer(x, enc_out, src_mask, tgt_mask) return self.decoder_norm(x) def forward(self, src_ids, tgt_ids, src_lang_id, tgt_lang_id, src_mask=None, tgt_mask=None): enc_out = self.encode(src_ids, src_lang_id, src_mask) dec_out = self.decode(tgt_ids, enc_out, tgt_lang_id, src_mask, tgt_mask) return self.output_proj(dec_out) @staticmethod def make_src_mask(src_ids, pad_id=0): # (B, 1, 1, S) return (src_ids != pad_id).unsqueeze(1).unsqueeze(2) @staticmethod def make_tgt_mask(tgt_ids, pad_id=0): B, T = tgt_ids.shape pad_mask = (tgt_ids != pad_id).unsqueeze(1).unsqueeze(2) # (B, 1, 1, T) causal = torch.tril(torch.ones(T, T, device=tgt_ids.device)).bool() # (T, T) return pad_mask & causal.unsqueeze(0).unsqueeze(0) def count_params(self): return sum(p.numel() for p in self.parameters()) if __name__ == "__main__": cfg = ModelConfig() model = TranslationModel(cfg) total = model.count_params() print(f"Total parameters: {total:,} ({total/1e6:.1f}M)") src = torch.randint(3, 100, (2, 32)) tgt = torch.randint(3, 100, (2, 24)) src_mask = model.make_src_mask(src) tgt_mask = model.make_tgt_mask(tgt) out = model(src, tgt, src_lang_id=0, tgt_lang_id=1, src_mask=src_mask, tgt_mask=tgt_mask) print(f"Output shape: {out.shape}") # (2, 24, vocab_size)