File size: 6,537 Bytes
4719692 | 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 | """100M-family decoder: RMSNorm, RoPE, SwiGLU and tied embeddings by default.
Keep 14 layers, width 768, 12 MHA heads, context 256, projection biases and
vanilla weight initialization. architecture_version=1 reproduces the historical
LayerNorm/learned-position/GELU model only for legacy checkpoint compatibility.
"""
from dataclasses import dataclass
import torch
from model_10m import model as vanilla
from torch import nn
from torch.nn import functional as F
__all__ = ["GPT", "Attention", "Block", "ModelConfig", "RMSNorm", "SwiGLU", "causal_attention"]
MODEL_DEPENDENCIES = {"model_10m_definition": vanilla.__file__}
@dataclass(frozen=True)
class ModelConfig(vanilla.ModelConfig):
n_layer: int = 14
n_head: int = 12
n_embd: int = 768
architecture_version: int = 2
def __post_init__(self):
if type(self.architecture_version) is not int or self.architecture_version not in (1, 2):
raise ValueError("architecture_version must be 1 (legacy) or 2 (RMSNorm/RoPE/SwiGLU)")
super().__post_init__()
if self.architecture_version == 2 and (self.n_embd // self.n_head) % 2:
raise ValueError("RoPE requires an even head dimension")
def causal_attention(query, key, value):
"""Standard causal SDPA; CUDA selects FlashAttention where supported.
Same scale, mask and zero dropout as the explicit 10M reference. The legacy
architecture retains its original implementation in model_10m.
"""
return F.scaled_dot_product_attention(query, key, value, is_causal=True, dropout_p=0.0)
class RMSNorm(nn.Module):
def __init__(self, width, eps=1e-5):
super().__init__()
self.weight = nn.Parameter(torch.ones(width))
self.eps = eps # Match the historical LayerNorm epsilon.
def forward(self, x):
values = x.float() if x.dtype in (torch.float16, torch.bfloat16) else x
normalized = values * torch.rsqrt(values.square().mean(-1, keepdim=True) + self.eps)
return normalized.to(x.dtype) * self.weight.to(x.dtype)
def rotary(x, offset=0):
"""Rotate adjacent query/key feature pairs using position and theta 10,000."""
dtype = torch.float64 if x.dtype == torch.float64 else torch.float32
frequencies = 10000.0 ** (
-torch.arange(0, x.size(-1), 2, device=x.device, dtype=dtype) / x.size(-1)
)
positions = torch.arange(offset, offset + x.size(-2), device=x.device, dtype=dtype)
angles = torch.outer(positions, frequencies)
cos, sin = angles.cos().to(x.dtype), angles.sin().to(x.dtype)
even, odd = x[..., ::2], x[..., 1::2]
return torch.stack((even * cos - odd * sin, even * sin + odd * cos), -1).flatten(-2)
class Attention(vanilla.Attention):
def cache_rotary(self, length):
width = self.qkv.in_features // self.n_head
frequencies = 10000.0 ** (
-torch.arange(0, width, 2, device=self.qkv.weight.device).float() / width
)
angles = torch.outer(torch.arange(length, device=frequencies.device).float(), frequencies)
self.register_buffer("rope_cos", angles.cos(), persistent=False)
self.register_buffer("rope_sin", angles.sin(), persistent=False)
def rotate(self, x):
if not hasattr(self, "rope_cos"):
return rotary(x)
cos = self.rope_cos[: x.size(-2)].to(x.dtype)
sin = self.rope_sin[: x.size(-2)].to(x.dtype)
even, odd = x[..., ::2], x[..., 1::2]
return torch.stack((even * cos - odd * sin, even * sin + odd * cos), -1).flatten(-2)
def forward(self, x):
batch, length, width = x.shape
query, key, value = self.qkv(x).chunk(3, -1)
query, key, value = [
tensor.view(batch, length, self.n_head, width // self.n_head).transpose(1, 2)
for tensor in (query, key, value)
]
attended = causal_attention(self.rotate(query), self.rotate(key), value)
return self.projection(attended.transpose(1, 2).contiguous().view(batch, length, width))
class SwiGLU(nn.Module):
def __init__(self, width):
super().__init__()
self.hidden_width = (8 * width + 2) // 3
self.gate = nn.Linear(width, self.hidden_width)
self.up = nn.Linear(width, self.hidden_width)
self.down = nn.Linear(self.hidden_width, width)
def forward(self, x):
return self.down(F.silu(self.gate(x)) * self.up(x))
class Block(vanilla.Block):
def __init__(self, config):
super().__init__(config)
self.attention_norm = RMSNorm(config.n_embd)
self.attention = Attention(config)
self.mlp_norm = RMSNorm(config.n_embd)
self.mlp = SwiGLU(config.n_embd)
class GPT(vanilla.GPT):
def __init__(self, config):
if config.architecture_version == 1:
super().__init__(config)
return
nn.Module.__init__(self)
self.config = config
self.token_embedding = nn.Embedding(config.vocab_size, config.n_embd)
self.blocks = nn.Sequential(*(Block(config) for _ in range(config.n_layer)))
self.final_norm = RMSNorm(config.n_embd)
self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
self.apply(self.initialize)
self.lm_head.weight = self.token_embedding.weight
def forward(self, tokens, labels=None, loss_chunk_size=0):
if self.config.architecture_version == 1:
return super().forward(tokens)
if not 0 < tokens.size(1) <= self.config.block_size:
raise ValueError("sequence length must be within the model's block_size")
hidden = self.final_norm(self.blocks(self.token_embedding(tokens)))
if labels is None:
return self.lm_head(hidden)
if not loss_chunk_size:
return F.cross_entropy(
self.lm_head(hidden).flatten(0, 1),
labels.flatten(),
ignore_index=-100,
reduction="sum",
)
# Recompute each projection during backward, bounding vocabulary-logit memory.
from torch.utils.checkpoint import checkpoint
def chunk_loss(features, targets):
return F.cross_entropy(
self.lm_head(features), targets, ignore_index=-100, reduction="sum"
)
return sum(
checkpoint(chunk_loss, features, targets, use_reentrant=False)
for features, targets in zip(
hidden.flatten(0, 1).split(loss_chunk_size), labels.flatten().split(loss_chunk_size)
)
)
|