| """Dense decoder-only Transformer adapted from the Part 1 Model 1 architecture.""" |
|
|
| from __future__ import annotations |
|
|
| from dataclasses import dataclass |
|
|
| import torch |
| from torch import nn |
| from torch.nn import functional as F |
|
|
|
|
| @dataclass(frozen=True) |
| class ModelConfig: |
| vocab_size: int |
| d_model: int = 512 |
| n_heads: int = 8 |
| n_layers: int = 6 |
| mlp_hidden_size: int = 2048 |
| max_sequence_length: int = 256 |
| dropout: float = 0.1 |
| tie_embeddings: bool = True |
|
|
| @classmethod |
| def from_mapping(cls, values: dict, vocab_size: int | None = None) -> ModelConfig: |
| clean = dict(values) |
| if vocab_size is not None: |
| clean["vocab_size"] = vocab_size |
| return cls(**clean) |
|
|
| def __post_init__(self) -> None: |
| if self.d_model % self.n_heads: |
| raise ValueError("d_model must be divisible by n_heads") |
| if self.d_model // self.n_heads % 2: |
| raise ValueError("attention head size must be even for rotary embeddings") |
| if self.max_sequence_length < 2: |
| raise ValueError("max_sequence_length must be at least 2") |
|
|
|
|
| class DenseMLP(nn.Module): |
| """Part 1 Model 1's standard two-layer GELU feed-forward network.""" |
|
|
| def __init__(self, config: ModelConfig) -> None: |
| super().__init__() |
| self.up = nn.Linear(config.d_model, config.mlp_hidden_size, bias=False) |
| self.down = nn.Linear(config.mlp_hidden_size, config.d_model, bias=False) |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| return self.down(F.gelu(self.up(x))) |
|
|
|
|
| class CausalSelfAttention(nn.Module): |
| def __init__(self, config: ModelConfig) -> None: |
| super().__init__() |
| self.n_heads = config.n_heads |
| self.head_size = config.d_model // config.n_heads |
| self.dropout = config.dropout |
| self.qkv = nn.Linear(config.d_model, 3 * config.d_model, bias=False) |
| self.out = nn.Linear(config.d_model, config.d_model, bias=False) |
| positions = torch.arange(config.max_sequence_length) |
| frequencies = 10000.0 ** (-torch.arange(0, self.head_size, 2).float() / self.head_size) |
| angles = torch.outer(positions, frequencies) |
| self.register_buffer("rope_cos", angles.cos(), persistent=False) |
| self.register_buffer("rope_sin", angles.sin(), persistent=False) |
|
|
| def _apply_rope(self, tensor: torch.Tensor, position_offset: int) -> torch.Tensor: |
| length = tensor.shape[-2] |
| stop = position_offset + length |
| cosine = self.rope_cos[position_offset:stop].to(tensor.dtype)[None, None] |
| sine = self.rope_sin[position_offset:stop].to(tensor.dtype)[None, None] |
| even, odd = tensor[..., ::2], tensor[..., 1::2] |
| return torch.stack( |
| (even * cosine - odd * sine, even * sine + odd * cosine), dim=-1 |
| ).flatten(-2) |
|
|
| def forward( |
| self, |
| x: torch.Tensor, |
| cache: tuple[torch.Tensor, torch.Tensor] | None = None, |
| ) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]: |
| batch, length, width = x.shape |
| query, key, value = ( |
| self.qkv(x) |
| .view(batch, length, 3, self.n_heads, self.head_size) |
| .permute(2, 0, 3, 1, 4) |
| .unbind(0) |
| ) |
| cached_length = 0 if cache is None else cache[0].shape[-2] |
| query = self._apply_rope(query, cached_length) |
| key = self._apply_rope(key, cached_length) |
| if cache is not None: |
| if length != 1: |
| raise ValueError("cached attention accepts one new token at a time") |
| key = torch.cat((cache[0], key), dim=-2) |
| value = torch.cat((cache[1], value), dim=-2) |
| attended = F.scaled_dot_product_attention( |
| query, |
| key, |
| value, |
| is_causal=cache is None, |
| dropout_p=self.dropout if self.training else 0.0, |
| ) |
| output = self.out(attended.transpose(1, 2).reshape(batch, length, width)) |
| return output, (key, value) |
|
|
|
|
| class TransformerBlock(nn.Module): |
| def __init__(self, config: ModelConfig) -> None: |
| super().__init__() |
| self.attention_norm = nn.LayerNorm(config.d_model) |
| self.mlp_norm = nn.LayerNorm(config.d_model) |
| self.attention = CausalSelfAttention(config) |
| self.mlp = DenseMLP(config) |
| self.dropout = nn.Dropout(config.dropout) |
|
|
| def forward( |
| self, |
| x: torch.Tensor, |
| cache: tuple[torch.Tensor, torch.Tensor] | None = None, |
| ) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]: |
| attended, next_cache = self.attention(self.attention_norm(x), cache) |
| x = x + self.dropout(attended) |
| return x + self.dropout(self.mlp(self.mlp_norm(x))), next_cache |
|
|
|
|
| class DenseTransformer(nn.Module): |
| def __init__(self, config: ModelConfig) -> None: |
| super().__init__() |
| self.config = config |
| self.embedding = nn.Embedding(config.vocab_size, config.d_model) |
| self.dropout = nn.Dropout(config.dropout) |
| self.blocks = nn.ModuleList([TransformerBlock(config) for _ in range(config.n_layers)]) |
| self.final_norm = nn.LayerNorm(config.d_model) |
| self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False) |
| self.apply(self._initialize) |
| if config.tie_embeddings: |
| self.lm_head.weight = self.embedding.weight |
|
|
| @staticmethod |
| def _initialize(module: nn.Module) -> None: |
| if isinstance(module, (nn.Linear, nn.Embedding)): |
| nn.init.normal_(module.weight, mean=0.0, std=0.02) |
| if isinstance(module, nn.Linear) and module.bias is not None: |
| nn.init.zeros_(module.bias) |
|
|
| def forward( |
| self, |
| input_ids: torch.Tensor, |
| labels: torch.Tensor | None = None, |
| *, |
| cache: list[tuple[torch.Tensor, torch.Tensor]] | None = None, |
| use_cache: bool = False, |
| ) -> dict[str, object]: |
| if input_ids.ndim != 2: |
| raise ValueError("input_ids must have shape [batch, sequence]") |
| cached_length = 0 if cache is None else cache[0][0].shape[-2] |
| if cached_length + input_ids.shape[1] > self.config.max_sequence_length: |
| raise ValueError("input exceeds configured maximum sequence length") |
| if cache is not None and len(cache) != len(self.blocks): |
| raise ValueError("cache must contain one entry per Transformer block") |
| hidden = self.dropout(self.embedding(input_ids)) |
| next_cache = [] |
| for index, block in enumerate(self.blocks): |
| hidden, block_cache = block(hidden, None if cache is None else cache[index]) |
| if use_cache: |
| next_cache.append(block_cache) |
| logits = self.lm_head(self.final_norm(hidden)) |
| loss = None |
| if labels is not None: |
| if labels.shape != input_ids.shape: |
| raise ValueError("labels and input_ids must have the same shape") |
| loss = F.cross_entropy( |
| logits.reshape(-1, logits.shape[-1]).float(), |
| labels.reshape(-1), |
| ignore_index=-100, |
| ) |
| return {"logits": logits, "loss": loss, "cache": next_cache if use_cache else None} |
|
|
| @torch.no_grad() |
| def greedy_generate(self, prompt: torch.Tensor, max_new_tokens: int) -> torch.Tensor: |
| """Generate with the model's forward pass and deterministic greedy decoding.""" |
| if prompt.ndim != 2: |
| raise ValueError("prompt must have shape [batch, sequence]") |
| if prompt.shape[1] + max_new_tokens > self.config.max_sequence_length: |
| raise ValueError("prompt and continuation exceed the model context window") |
| generated = prompt |
| output = self(prompt, use_cache=True) |
| for step in range(max_new_tokens): |
| next_token = output["logits"][:, -1].argmax(dim=-1, keepdim=True) |
| generated = torch.cat((generated, next_token), dim=1) |
| if step + 1 < max_new_tokens: |
| output = self(next_token, cache=output["cache"], use_cache=True) |
| return generated |
|
|
|
|
| def parameter_count(model: nn.Module) -> int: |
| return sum(parameter.numel() for parameter in model.parameters()) |
|
|
