| 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)) |
|
|
