""" Same architecture family as before (RMSNorm, RoPE, grouped-query attention via F.scaled_dot_product_attention, SwiGLU, tied embeddings) -- only the CONFIG changed (see configs/config.py): small custom vocab, shorter context, sized to land ~18.9M params at a genuine ~20:1 token:param ratio. Run directly to print exact param count + smoke test: python model.py """ import torch import torch.nn as nn import torch.nn.functional as F from configs.config import ModelConfig class RMSNorm(nn.Module): def __init__(self, dim: int, eps: float = 1e-5): super().__init__() self.eps = eps self.weight = nn.Parameter(torch.ones(dim)) def forward(self, x): norm = x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps) return norm * self.weight def precompute_rope(head_dim, seq_len, theta, device, dtype=torch.float32): freqs = 1.0 / (theta ** (torch.arange(0, head_dim, 2, device=device, dtype=dtype) / head_dim)) t = torch.arange(seq_len, device=device, dtype=dtype) freqs = torch.outer(t, freqs) return torch.cos(freqs), torch.sin(freqs) def apply_rope(x, cos, sin): x1, x2 = x[..., 0::2], x[..., 1::2] cos = cos[None, None, :, :] sin = sin[None, None, :, :] r1 = x1 * cos - x2 * sin r2 = x1 * sin + x2 * cos return torch.stack([r1, r2], dim=-1).flatten(-2).to(x.dtype) class GQAttention(nn.Module): def __init__(self, cfg: ModelConfig): super().__init__() assert cfg.d_model % cfg.n_head == 0 assert cfg.n_head % cfg.n_kv_head == 0 self.n_head = cfg.n_head self.n_kv_head = cfg.n_kv_head self.head_dim = cfg.d_model // cfg.n_head self.n_rep = cfg.n_head // cfg.n_kv_head self.q_proj = nn.Linear(cfg.d_model, cfg.n_head * self.head_dim, bias=False) self.k_proj = nn.Linear(cfg.d_model, cfg.n_kv_head * self.head_dim, bias=False) self.v_proj = nn.Linear(cfg.d_model, cfg.n_kv_head * self.head_dim, bias=False) self.o_proj = nn.Linear(cfg.n_head * self.head_dim, cfg.d_model, bias=False) def forward(self, x, cos, sin): b, t, _ = x.shape q = self.q_proj(x).view(b, t, self.n_head, self.head_dim).transpose(1, 2) k = self.k_proj(x).view(b, t, self.n_kv_head, self.head_dim).transpose(1, 2) v = self.v_proj(x).view(b, t, self.n_kv_head, self.head_dim).transpose(1, 2) q = apply_rope(q, cos, sin) k = apply_rope(k, cos, sin) if self.n_rep > 1: k = k.repeat_interleave(self.n_rep, dim=1) v = v.repeat_interleave(self.n_rep, dim=1) out = F.scaled_dot_product_attention(q, k, v, is_causal=True) out = out.transpose(1, 2).contiguous().view(b, t, self.n_head * self.head_dim) return self.o_proj(out) class SwiGLU(nn.Module): def __init__(self, cfg: ModelConfig): super().__init__() self.gate_proj = nn.Linear(cfg.d_model, cfg.d_ff, bias=False) self.up_proj = nn.Linear(cfg.d_model, cfg.d_ff, bias=False) self.down_proj = nn.Linear(cfg.d_ff, cfg.d_model, bias=False) def forward(self, x): return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x)) class Block(nn.Module): def __init__(self, cfg: ModelConfig): super().__init__() self.attn_norm = RMSNorm(cfg.d_model) self.attn = GQAttention(cfg) self.mlp_norm = RMSNorm(cfg.d_model) self.mlp = SwiGLU(cfg) self.dropout = nn.Dropout(cfg.dropout) def forward(self, x, cos, sin): x = x + self.dropout(self.attn(self.attn_norm(x), cos, sin)) x = x + self.dropout(self.mlp(self.mlp_norm(x))) return x class TinyTransformer(nn.Module): def __init__(self, cfg: ModelConfig): super().__init__() self.cfg = cfg self.tok_emb = nn.Embedding(cfg.vocab_size, cfg.d_model) self.blocks = nn.ModuleList([Block(cfg) for _ in range(cfg.n_layer)]) self.final_norm = RMSNorm(cfg.d_model) self.lm_head = nn.Linear(cfg.d_model, cfg.vocab_size, bias=False) if cfg.tie_embeddings: self.lm_head.weight = self.tok_emb.weight self.head_dim = cfg.d_model // cfg.n_head self.apply(self._init_weights) def _init_weights(self, module): if isinstance(module, nn.Linear): nn.init.normal_(module.weight, mean=0.0, std=0.02) if module.bias is not None: nn.init.zeros_(module.bias) elif isinstance(module, nn.Embedding): nn.init.normal_(module.weight, mean=0.0, std=0.02) def forward(self, idx, targets=None): b, t = idx.shape assert t <= self.cfg.context_len, f"seq len {t} exceeds context_len {self.cfg.context_len}" cos, sin = precompute_rope(self.head_dim, t, self.cfg.rope_theta, idx.device) cos, sin = cos.to(self.tok_emb.weight.dtype), sin.to(self.tok_emb.weight.dtype) x = self.tok_emb(idx) for block in self.blocks: x = block(x, cos, sin) x = self.final_norm(x) logits = self.lm_head(x) loss = None if targets is not None: loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1) return logits, loss @torch.no_grad() def generate(self, idx, max_new_tokens, temperature=1.0, top_k=None): for _ in range(max_new_tokens): idx_cond = idx if idx.size(1) <= self.cfg.context_len else idx[:, -self.cfg.context_len:] logits, _ = self(idx_cond) logits = logits[:, -1, :] / max(temperature, 1e-5) if top_k is not None: v, _ = torch.topk(logits, min(top_k, logits.size(-1))) logits[logits < v[:, [-1]]] = -float("inf") probs = F.softmax(logits, dim=-1) next_id = torch.multinomial(probs, num_samples=1) idx = torch.cat([idx, next_id], dim=1) return idx def num_params(self, non_embedding=False): n = sum(p.numel() for p in self.parameters()) if non_embedding: n -= self.tok_emb.weight.numel() return n if __name__ == "__main__": cfg = ModelConfig() model = TinyTransformer(cfg) n = model.num_params() n_emb = model.tok_emb.weight.numel() print(f"Config: vocab={cfg.vocab_size} d_model={cfg.d_model} n_layer={cfg.n_layer} " f"n_head={cfg.n_head} n_kv_head={cfg.n_kv_head} d_ff={cfg.d_ff} context_len={cfg.context_len}") print(f"Total parameters: {n:,} (~{n/1e6:.2f}M)") print(f"Embedding: {n_emb:,} ({100*n_emb/n:.0f}% of total)") x = torch.randint(0, cfg.vocab_size, (2, 64)) y = torch.randint(0, cfg.vocab_size, (2, 64)) logits, loss = model(x, y) assert logits.shape == (2, 64, cfg.vocab_size) loss.backward() n_missing = sum(1 for p in model.parameters() if p.grad is None) print(f"Forward/backward OK. loss={loss.item():.3f} params_without_grad={n_missing}")