File size: 4,732 Bytes
f5f443b | 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 | from dataclasses import asdict, dataclass
import torch
from torch import nn
from torch.nn import functional as functional
from torch.utils.checkpoint import checkpoint
@dataclass
class ModelConfig:
vocab_size: int = 8192
context_length: int = 1024
layers: int = 8
hidden_size: int = 512
heads: int = 8
intermediate_size: int = 1408
def __post_init__(self):
if any(value <= 0 for value in asdict(self).values()):
raise ValueError("All model dimensions must be positive")
if self.hidden_size % self.heads or (self.hidden_size // self.heads) % 2:
raise ValueError("Head dimension must be an even integer")
class Attention(nn.Module):
def __init__(self, config):
super().__init__()
self.heads = config.heads
self.head_size = config.hidden_size // config.heads
self.qkv = nn.Linear(config.hidden_size, 3 * config.hidden_size, bias=False)
self.output = nn.Linear(config.hidden_size, config.hidden_size, bias=False)
frequencies = 1.0 / (10000 ** (torch.arange(0, self.head_size, 2).float() / self.head_size))
angles = torch.outer(torch.arange(config.context_length).float(), frequencies)
self.register_buffer("cosines", angles.cos(), persistent=False)
self.register_buffer("sines", angles.sin(), persistent=False)
def rotate(self, tensor):
length = tensor.size(-2)
cosines = self.cosines[:length].to(tensor.dtype)
sines = self.sines[:length].to(tensor.dtype)
even, odd = tensor[..., 0::2], tensor[..., 1::2]
return torch.stack((even * cosines - odd * sines, even * sines + odd * cosines), dim=-1).flatten(-2)
def forward(self, hidden):
batch_size, length, width = hidden.shape
query, key, value = self.qkv(hidden).reshape(batch_size, length, 3, self.heads, self.head_size).unbind(2)
attended = functional.scaled_dot_product_attention(
self.rotate(query.transpose(1, 2)),
self.rotate(key.transpose(1, 2)),
value.transpose(1, 2),
is_causal=True,
)
return self.output(attended.transpose(1, 2).contiguous().view(batch_size, length, width))
class Block(nn.Module):
def __init__(self, config):
super().__init__()
self.attention_norm = nn.RMSNorm(config.hidden_size, eps=1e-5)
self.attention = Attention(config)
self.ffn_norm = nn.RMSNorm(config.hidden_size, eps=1e-5)
self.gate = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
self.up = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
self.down = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)
def forward(self, hidden):
hidden = hidden + self.attention(self.attention_norm(hidden))
normalized = self.ffn_norm(hidden)
return hidden + self.down(functional.silu(self.gate(normalized)) * self.up(normalized))
class LanguageModel(nn.Module):
def __init__(self, config):
super().__init__()
self.config = config
self.embedding = nn.Embedding(config.vocab_size, config.hidden_size)
self.blocks = nn.ModuleList([Block(config) for _ in range(config.layers)])
self.norm = nn.RMSNorm(config.hidden_size, eps=1e-5)
self.gradient_checkpointing = False
self.apply(self.initialize)
def set_gradient_checkpointing(self, enabled):
self.gradient_checkpointing = enabled
@staticmethod
def initialize(module):
if isinstance(module, (nn.Linear, nn.Embedding)):
nn.init.normal_(module.weight, mean=0.0, std=0.02)
def forward(self, tokens, targets=None, return_logits=False, logit_positions=None):
if tokens.size(1) > self.config.context_length:
raise ValueError("Input exceeds context length")
hidden = self.embedding(tokens)
for block in self.blocks:
if self.gradient_checkpointing and self.training and torch.is_grad_enabled():
hidden = checkpoint(block, hidden, use_reentrant=False)
else:
hidden = block(hidden)
hidden = self.norm(hidden)
if logit_positions is not None:
hidden = hidden[:, logit_positions, :]
if targets is None and not return_logits:
return functional.linear(hidden[:, -1:, :], self.embedding.weight)
logits = functional.linear(hidden, self.embedding.weight)
if targets is None:
return logits
return functional.cross_entropy(logits.reshape(-1, self.config.vocab_size), targets.reshape(-1))
|