Upload model.py with huggingface_hub

model.py

@@ -0,0 +1,173 @@
+"""
+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}")