Upload src/xscript/model.py with huggingface_hub

src/xscript/model.py

@@ -0,0 +1,186 @@
+"""Llama-style decoder-only transformer.
+
+RMSNorm, rotary position embeddings, SwiGLU MLP, untied input/output
+embeddings. Deliberately small and dependency-light (just torch) so it runs
+unchanged on the GH200 nodes and on a login-node CPU smoke test.
+
+Configs live in configs/*.yaml; ModelConfig mirrors the yaml `model:` block.
+"""
+from dataclasses import dataclass
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+@dataclass
+class ModelConfig:
+    vocab_size: int = 65536
+    dim: int = 2048
+    n_layers: int = 16
+    n_heads: int = 16
+    n_kv_heads: int | None = None      # None -> = n_heads (no GQA)
+    ffn_dim: int = 5632
+    max_seq_len: int = 2048
+    rope_theta: float = 10000.0
+    norm_eps: float = 1e-5
+
+    @property
+    def kv_heads(self) -> int:
+        return self.n_kv_heads or self.n_heads
+
+    @property
+    def head_dim(self) -> int:
+        return self.dim // self.n_heads
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x):
+        dt = x.dtype
+        x = x.float()
+        x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
+        return (x * self.weight.float()).to(dt)
+
+
+def _rope_cache(seq_len: int, head_dim: int, theta: float, device, dtype):
+    inv = 1.0 / (theta ** (torch.arange(0, head_dim, 2, device=device).float() / head_dim))
+    t = torch.arange(seq_len, device=device).float()
+    freqs = torch.outer(t, inv)                       # (T, head_dim/2)
+    return torch.cos(freqs).to(dtype), torch.sin(freqs).to(dtype)
+
+
+def _apply_rope(x, cos, sin):
+    # x: (B, H, T, D). split even/odd halves (rotate-half convention)
+    x1, x2 = x[..., ::2], x[..., 1::2]
+    cos = cos[None, None, :, :]
+    sin = sin[None, None, :, :]
+    o1 = x1 * cos - x2 * sin
+    o2 = x1 * sin + x2 * cos
+    out = torch.empty_like(x)
+    out[..., ::2] = o1
+    out[..., 1::2] = o2
+    return out
+
+
+class Attention(nn.Module):
+    def __init__(self, cfg: ModelConfig):
+        super().__init__()
+        self.n_heads = cfg.n_heads
+        self.kv_heads = cfg.kv_heads
+        self.head_dim = cfg.head_dim
+        self.wq = nn.Linear(cfg.dim, cfg.n_heads * cfg.head_dim, bias=False)
+        self.wk = nn.Linear(cfg.dim, cfg.kv_heads * cfg.head_dim, bias=False)
+        self.wv = nn.Linear(cfg.dim, cfg.kv_heads * cfg.head_dim, bias=False)
+        self.wo = nn.Linear(cfg.n_heads * cfg.head_dim, cfg.dim, bias=False)
+
+    def forward(self, x, cos, sin):
+        B, T, _ = x.shape
+        q = self.wq(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
+        k = self.wk(x).view(B, T, self.kv_heads, self.head_dim).transpose(1, 2)
+        v = self.wv(x).view(B, T, self.kv_heads, self.head_dim).transpose(1, 2)
+        q = _apply_rope(q, cos, sin)
+        k = _apply_rope(k, cos, sin)
+        if self.kv_heads != self.n_heads:
+            rep = self.n_heads // self.kv_heads
+            k = k.repeat_interleave(rep, dim=1)
+            v = v.repeat_interleave(rep, dim=1)
+        out = F.scaled_dot_product_attention(q, k, v, is_causal=True)
+        out = out.transpose(1, 2).contiguous().view(B, T, -1)
+        return self.wo(out)
+
+
+class SwiGLU(nn.Module):
+    def __init__(self, cfg: ModelConfig):
+        super().__init__()
+        self.w1 = nn.Linear(cfg.dim, cfg.ffn_dim, bias=False)   # gate
+        self.w3 = nn.Linear(cfg.dim, cfg.ffn_dim, bias=False)   # up
+        self.w2 = nn.Linear(cfg.ffn_dim, cfg.dim, bias=False)   # down
+
+    def forward(self, x):
+        return self.w2(F.silu(self.w1(x)) * self.w3(x))
+
+
+class Block(nn.Module):
+    def __init__(self, cfg: ModelConfig):
+        super().__init__()
+        self.attn_norm = RMSNorm(cfg.dim, cfg.norm_eps)
+        self.attn = Attention(cfg)
+        self.ffn_norm = RMSNorm(cfg.dim, cfg.norm_eps)
+        self.ffn = SwiGLU(cfg)
+
+    def forward(self, x, cos, sin):
+        x = x + self.attn(self.attn_norm(x), cos, sin)
+        x = x + self.ffn(self.ffn_norm(x))
+        return x
+
+
+class Transformer(nn.Module):
+    def __init__(self, cfg: ModelConfig):
+        super().__init__()
+        self.cfg = cfg
+        self.tok_emb = nn.Embedding(cfg.vocab_size, cfg.dim)
+        self.layers = nn.ModuleList(Block(cfg) for _ in range(cfg.n_layers))
+        self.norm = RMSNorm(cfg.dim, cfg.norm_eps)
+        self.lm_head = nn.Linear(cfg.dim, cfg.vocab_size, bias=False)  # untied
+        self._rope = None
+        self.apply(self._init)
+        # scale residual-projection inits by depth (GPT-2/Llama convention)
+        for name, p in self.named_parameters():
+            if name.endswith("wo.weight") or name.endswith("w2.weight"):
+                nn.init.normal_(p, mean=0.0, std=0.02 / (2 * cfg.n_layers) ** 0.5)
+
+    def _init(self, m):
+        if isinstance(m, nn.Linear):
+            nn.init.normal_(m.weight, mean=0.0, std=0.02)
+        elif isinstance(m, nn.Embedding):
+            nn.init.normal_(m.weight, mean=0.0, std=0.02)
+
+    def _rope_for(self, T, device, dtype):
+        if self._rope is None or self._rope[0].shape[0] < T or self._rope[0].device != device:
+            self._rope = _rope_cache(self.cfg.max_seq_len, self.cfg.head_dim,
+                                     self.cfg.rope_theta, device, dtype)
+        cos, sin = self._rope
+        return cos[:T], sin[:T]
+
+    def forward(self, idx, targets=None):
+        B, T = idx.shape
+        x = self.tok_emb(idx)
+        cos, sin = self._rope_for(T, idx.device, x.dtype)
+        for layer in self.layers:
+            x = layer(x, cos, sin)
+        x = self.norm(x)
+        if targets is None:
+            return self.lm_head(x[:, -1:, :])
+        logits = self.lm_head(x)
+        loss = F.cross_entropy(logits.view(-1, logits.size(-1)),
+                               targets.reshape(-1), ignore_index=-100)
+        return logits, loss
+
+    @torch.no_grad()
+    def layer_reps(self, idx):
+        """Per-layer hidden states for representation analysis (MEXA).
+
+        Returns a tensor (n_layers+1, B, T, dim): index 0 is the embedding
+        output, index i>=1 is the output of block i. Causal attention means
+        right-padding never contaminates real positions, so callers can pool
+        over a length mask safely.
+        """
+        B, T = idx.shape
+        x = self.tok_emb(idx)
+        cos, sin = self._rope_for(T, idx.device, x.dtype)
+        reps = [x]
+        for layer in self.layers:
+            x = layer(x, cos, sin)
+            reps.append(x)
+        return torch.stack(reps, dim=0)
+
+    def num_params(self, embedding: bool = True) -> int:
+        n = sum(p.numel() for p in self.parameters())
+        if not embedding:
+            n -= self.tok_emb.weight.numel() + self.lm_head.weight.numel()
+        return n