| """Eather-500M: custom decoder-only Transformer in pure PyTorch. |
| RoPE, grouped-query attention, RMSNorm, SwiGLU, tied token embeddings. |
| """ |
| import json |
| from dataclasses import dataclass |
| from pathlib import Path |
| import torch |
| from torch import nn |
| import torch.nn.functional as F |
|
|
|
|
| @dataclass |
| class ModelConfig: |
| vocab_size: int = 32000 |
| d_model: int = 1280 |
| n_layers: int = 26 |
| n_heads: int = 20 |
| n_kv_heads: int = 5 |
| ffn_hidden: int = 3584 |
| max_seq_len: int = 2048 |
| rope_theta: float = 10000.0 |
| dropout: float = 0.0 |
| tie_embeddings: bool = True |
|
|
| @classmethod |
| def load(cls, path): |
| return cls(**json.loads(Path(path).read_text())) |
|
|
|
|
| def rms_norm(x, weight, eps=1e-6): |
| |
| var = x.float().pow(2).mean(dim=-1, keepdim=True) |
| return (x.float() * torch.rsqrt(var + eps)).to(x.dtype) * weight |
|
|
|
|
| class RMSNorm(nn.Module): |
| def __init__(self, size): |
| super().__init__() |
| self.weight = nn.Parameter(torch.ones(size)) |
|
|
| def forward(self, x): |
| return rms_norm(x, self.weight) |
|
|
|
|
| def rope(q, k, theta, offset=0): |
| d = q.shape[-1] |
| positions = torch.arange(offset, offset + q.shape[-2], device=q.device, dtype=torch.float32) |
| freqs = 1.0 / (theta ** (torch.arange(0, d, 2, device=q.device, dtype=torch.float32) / d)) |
| angles = torch.outer(positions, freqs) |
| cos = angles.cos()[None, None, :, :] |
| sin = angles.sin()[None, None, :, :] |
|
|
| def rotate(x): |
| a, b = x.float()[..., ::2], x.float()[..., 1::2] |
| return torch.stack((a * cos - b * sin, a * sin + b * cos), dim=-1).flatten(-2).to(x.dtype) |
|
|
| return rotate(q), rotate(k) |
|
|
|
|
| class Attention(nn.Module): |
| def __init__(self, cfg): |
| super().__init__() |
| assert cfg.d_model % cfg.n_heads == 0 |
| assert cfg.n_heads % cfg.n_kv_heads == 0 |
| self.head_dim = cfg.d_model // cfg.n_heads |
| self.n_heads = cfg.n_heads |
| self.n_kv_heads = cfg.n_kv_heads |
| self.theta = cfg.rope_theta |
| self.dropout = cfg.dropout |
| self.q = nn.Linear(cfg.d_model, cfg.n_heads * self.head_dim, bias=False) |
| self.k = nn.Linear(cfg.d_model, cfg.n_kv_heads * self.head_dim, bias=False) |
| self.v = nn.Linear(cfg.d_model, cfg.n_kv_heads * self.head_dim, bias=False) |
| self.o = nn.Linear(cfg.d_model, cfg.d_model, bias=False) |
|
|
| def forward(self, x): |
| b, t, _ = x.shape |
| q = self.q(x).view(b, t, self.n_heads, self.head_dim).transpose(1, 2) |
| k = self.k(x).view(b, t, self.n_kv_heads, self.head_dim).transpose(1, 2) |
| v = self.v(x).view(b, t, self.n_kv_heads, self.head_dim).transpose(1, 2) |
| q, k = rope(q, k, self.theta) |
| |
| repeats = self.n_heads // self.n_kv_heads |
| k = k.repeat_interleave(repeats, dim=1) |
| v = v.repeat_interleave(repeats, dim=1) |
| out = F.scaled_dot_product_attention( |
| q, k, v, is_causal=True, |
| dropout_p=self.dropout if self.training else 0.0, |
| ) |
| out = out.transpose(1, 2).contiguous().view(b, t, -1) |
| return self.o(out) |
|
|
|
|
| class Block(nn.Module): |
| def __init__(self, cfg): |
| super().__init__() |
| self.attn_norm = RMSNorm(cfg.d_model) |
| self.attn = Attention(cfg) |
| self.ff_norm = RMSNorm(cfg.d_model) |
| self.up = nn.Linear(cfg.d_model, cfg.ffn_hidden, bias=False) |
| self.gate = nn.Linear(cfg.d_model, cfg.ffn_hidden, bias=False) |
| self.down = nn.Linear(cfg.ffn_hidden, cfg.d_model, bias=False) |
|
|
| def forward(self, x): |
| x = x + self.attn(self.attn_norm(x)) |
| y = self.ff_norm(x) |
| return x + self.down(F.silu(self.gate(y)) * self.up(y)) |
|
|
|
|
| class Eather500M(nn.Module): |
| def __init__(self, cfg: ModelConfig): |
| super().__init__() |
| self.cfg = cfg |
| self.embed = nn.Embedding(cfg.vocab_size, cfg.d_model) |
| self.blocks = nn.ModuleList(Block(cfg) for _ in range(cfg.n_layers)) |
| self.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.embed.weight |
| self.apply(self._init_weights) |
| |
| for layer in self.blocks: |
| nn.init.normal_(layer.attn.o.weight, mean=0, std=0.02 / (2 * cfg.n_layers) ** 0.5) |
| nn.init.normal_(layer.down.weight, mean=0, std=0.02 / (2 * cfg.n_layers) ** 0.5) |
| self.gradient_checkpointing = False |
|
|
| @staticmethod |
| def _init_weights(module): |
| if isinstance(module, (nn.Linear, nn.Embedding)): |
| nn.init.normal_(module.weight, mean=0, std=0.02) |
|
|
| def forward(self, input_ids, labels=None, loss_mask=None): |
| if input_ids.shape[1] > self.cfg.max_seq_len: |
| raise ValueError(f"Context exceeds {self.cfg.max_seq_len} tokens") |
| x = self.embed(input_ids) |
| for block in self.blocks: |
| if self.gradient_checkpointing and self.training: |
| x = torch.utils.checkpoint.checkpoint(block, x, use_reentrant=False) |
| else: |
| x = block(x) |
| logits = self.lm_head(self.norm(x)) |
| loss = None |
| if labels is not None: |
| |
| per_token = F.cross_entropy( |
| logits.float().reshape(-1, logits.size(-1)), labels.reshape(-1), |
| reduction="none", |
| ) |
| if loss_mask is not None: |
| mask = loss_mask.reshape(-1).float() |
| loss = (per_token * mask).sum() / mask.sum().clamp_min(1) |
| else: |
| loss = per_token.mean() |
| return logits, loss |
|
|
|
|
| if __name__ == "__main__": |
| import argparse |
| p = argparse.ArgumentParser() |
| p.add_argument("--config", default="config.json") |
| p.add_argument("--device", default="meta", help="Use 'meta' to count without allocating RAM") |
| args = p.parse_args() |
| with torch.device(args.device): |
| model = Eather500M(ModelConfig.load(args.config)) |
| n = sum(p.numel() for p in model.parameters()) |
| print(f"Parameters: {n:,} ({n / 1e6:.2f}M)") |
|
|
