File size: 6,280 Bytes
7da1a03 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 | """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):
# Accumulate normalization in fp32 for stability.
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)
# Explicit head repeats for compatibility with PyTorch versions without enable_gqa.
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)
# Small residual projections improve deep-transformer stability.
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:
# Shift inputs and labels in the dataset, so no shift happens here.
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)")
|