Upload model.py

model.py

@@ -0,0 +1,196 @@
+"""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())