Publish trained storyteller checkpoint

config.json

@@ -6,20 +6,20 @@
     "AutoConfig": "configuration_tiny_gpt.TinyGPTConfig",
     "AutoModelForCausalLM": "modeling_tiny_gpt.TinyGPTForCausalLM"
   },
-  "context_size": 512,
-  "d_model": 256,
-  "dropout": 0.1,
+  "context_size": 1024,
+  "d_model": 384,
+  "dropout": 0.0,
   "dtype": "float32",
   "eos_token_id": 0,
-  "hidden_size": 256,
-  "max_position_embeddings": 512,
+  "hidden_size": 384,
+  "max_position_embeddings": 1024,
   "model_type": "tiny_gpt",
-  "n_heads": 8,
-  "n_layers": 6,
-  "num_attention_heads": 8,
-  "num_hidden_layers": 6,
+  "n_heads": 6,
+  "n_layers": 10,
+  "num_attention_heads": 6,
+  "num_hidden_layers": 10,
   "pad_token_id": 0,
   "tie_word_embeddings": true,
   "transformers_version": "5.12.1",
-  "vocab_size": 10000
+  "vocab_size": 16000
 }

generation_config.json

@@ -1,4 +1,6 @@
 {
   "_from_model_config": true,
+  "output_attentions": false,
+  "output_hidden_states": false,
   "transformers_version": "5.12.1"
 }

model.py

@@ -1,24 +1,18 @@
-import math
 import torch
+import torch.nn.functional as F
 from torch import nn
 
-"""
-token embeddings
-learned positional embeddings
-causal self-attention
-feed-forward network
-custom LayerNorm
-residual connections
-stacked transformer blocks
-GELU
-dropout
-Pre-LayerNorm
-tied token embedding/output projection weights
-multi-head attention
-
-"""
-class Model(nn.Module):
-    def __init__(self, context_size, vocab_size, d_model, n_layers, n_heads, dropout=0.1):
+
+class GPTModel(nn.Module):
+    def __init__(
+        self,
+        context_size,
+        vocab_size,
+        d_model,
+        n_layers,
+        n_heads,
+        dropout=0.1,
+    ):
         super().__init__()
 
         self.context_size = context_size
@@ -31,136 +25,166 @@ class Model(nn.Module):
         self.token_embedding = nn.Embedding(vocab_size, d_model)
         self.position_embedding = nn.Embedding(context_size, d_model)
         self.transformer_blocks = nn.ModuleList(
-            [TransformerBlock(d_model, n_heads, dropout) for _ in range(n_layers)]
+            [
+                TransformerBlock(d_model, n_heads, dropout)
+                for _ in range(n_layers)
+            ]
         )
         self.linear = nn.Linear(d_model, vocab_size, bias=False)
         self.dropout = nn.Dropout(dropout)
-        self.final_layer_norm = LayerNorm(d_model)
+        self.final_layer_norm = nn.LayerNorm(d_model, eps=1e-6)
 
     def forward(self, x):
-        B, T = x.shape
-
-        assert T <= self.context_size, "Input sequence is longer than context_size"
+        _, sequence_length = x.shape
 
-        positions = torch.arange(T, device=x.device)
+        assert sequence_length <= self.context_size, (
+            "Input sequence is longer than context_size"
+        )
 
+        positions = torch.arange(sequence_length, device=x.device)
         position = self.position_embedding(positions)
         token = self.token_embedding(x)
 
-        x = token + position
-        x = self.dropout(x)
-
+        x = self.dropout(token + position)
         for block in self.transformer_blocks:
             x = block(x)
 
         x = self.final_layer_norm(x)
-        logits = self.linear(x)
+        return self.linear(x)
 
-        return logits
 
 class FeedForward(nn.Module):
     def __init__(self, d_model):
         super().__init__()
         self.ff1 = nn.Linear(d_model, 4 * d_model)
-        self.ff2 = nn.Linear(d_model * 4, d_model)
+        self.ff2 = nn.Linear(4 * d_model, d_model)
 
     def forward(self, x):
-        x = self.ff1(x)
-        x = nn.functional.gelu(x)
-        x = self.ff2(x)
-        return x
+        return self.ff2(F.gelu(self.ff1(x)))
 
-class LayerNorm(nn.Module):
-    def __init__(self, d_model):
+
+class TransformerBlock(nn.Module):
+    def __init__(self, d_model, n_heads, dropout=0.1):
         super().__init__()
-        self.gamma = nn.Parameter(torch.ones(d_model))
-        self.beta = nn.Parameter(torch.zeros(d_model))
+
+        self.feed_forward = FeedForward(d_model)
+        self.layer_norm1 = nn.LayerNorm(d_model, eps=1e-6)
+        self.layer_norm2 = nn.LayerNorm(d_model, eps=1e-6)
+        self.dropout = nn.Dropout(dropout)
+        self.multi_head_attention = MultiHeadAttention(
+            d_model=d_model,
+            n_heads=n_heads,
+            dropout=dropout,
+        )
 
     def forward(self, x):
-        mean = x.mean(dim=-1, keepdim=True)
-        diff = (x - mean)
-        variance = (diff * diff).mean(dim=-1, keepdim=True)
-        normalized = diff / torch.sqrt(variance + 1e-6)
-        return self.gamma * normalized + self.beta
+        attention = self.multi_head_attention(self.layer_norm1(x))
+        x = x + self.dropout(attention)
+
+        feed_forward = self.feed_forward(self.layer_norm2(x))
+        return x + self.dropout(feed_forward)
+
 
 class MultiHeadAttention(nn.Module):
     def __init__(self, d_model, n_heads, dropout=0.1):
         super().__init__()
 
-        assert d_model % n_heads == 0, "d_model must be divisible by n_heads"
+        assert d_model % n_heads == 0, (
+            "d_model must be divisible by n_heads"
+        )
 
         self.n_heads = n_heads
         self.head_dim = d_model // n_heads
-        self.scale = math.sqrt(self.head_dim)
-
-        self.query = nn.Linear(d_model, d_model)
-        self.key = nn.Linear(d_model, d_model)
-        self.value = nn.Linear(d_model, d_model)
+        self.dropout_p = dropout
 
+        self.qkv = nn.Linear(d_model, 3 * d_model)
         self.head_proj = nn.Linear(d_model, d_model)
 
     def forward(self, x):
-        query = self.split_heads(self.query(x))
-        key = self.split_heads(self.key(x))
-        value = self.split_heads(self.value(x))
-
-        scores = torch.matmul(query, key.transpose(-2, -1))
-        scores = scores / self.scale
-
-        context_size = query.shape[2]
-
-        mask = torch.tril(
-            torch.ones(context_size, context_size, device=query.device)
+        query, key, value = self.qkv(x).chunk(3, dim=-1)
+        query = self.split_heads(query)
+        key = self.split_heads(key)
+        value = self.split_heads(value)
+
+        attended = F.scaled_dot_product_attention(
+            query,
+            key,
+            value,
+            dropout_p=self.dropout_p if self.training else 0.0,
+            is_causal=True,
         )
-        mask = mask.view(1, 1, context_size, context_size)
-
-        scores = scores.masked_fill(mask == 0, float("-inf"))
-
-        weights = torch.nn.functional.softmax(scores, dim=-1)
-
-        attended = torch.matmul(weights, value)
-
-        attended = self.combine_heads(attended)
-
-        attended = self.head_proj(attended)
-
-        return attended
+        return self.head_proj(self.combine_heads(attended))
 
     def split_heads(self, x):
-        batch_size, seq_len, d_model = x.shape
-
-        x = x.reshape(batch_size, seq_len, self.n_heads, self.head_dim)
-        x = x.transpose(1, 2)
-
-        return x
+        batch_size, sequence_length, _ = x.shape
+        x = x.reshape(
+            batch_size,
+            sequence_length,
+            self.n_heads,
+            self.head_dim,
+        )
+        return x.transpose(1, 2)
 
     def combine_heads(self, x):
-        batch_size, n_heads, seq_len, head_dim = x.shape
-
+        batch_size, n_heads, sequence_length, head_dim = x.shape
         x = x.transpose(1, 2)
-        x = x.contiguous().view(batch_size, seq_len, n_heads * head_dim)
-
-        return x
-
-class TransformerBlock(nn.Module):
-    def __init__(self, d_model, n_heads, dropout=0.1):
-        super().__init__()
-
-        self.feed_forward = FeedForward(d_model)
-        self.layer_norm1 = LayerNorm(d_model)
-        self.layer_norm2 = LayerNorm(d_model)
-        self.dropout = nn.Dropout(dropout)
-        self.multi_head_attention = MultiHeadAttention(d_model, n_heads, dropout)
-
-    def forward(self, x):
-        attention = self.multi_head_attention(self.layer_norm1(x))
-        attention = self.dropout(attention)
+        return x.contiguous().view(
+            batch_size,
+            sequence_length,
+            n_heads * head_dim,
+        )
 
-        x = x + attention
 
-        feed_forward = self.feed_forward(self.layer_norm2(x))
-        feed_forward = self.dropout(feed_forward)
+Model = GPTModel
+
+
+def convert_reference_state_dict(state_dict):
+    """Convert reference Q/K/V and LayerNorm keys to the optimized layout."""
+    converted = dict(state_dict)
+
+    for key in list(converted):
+        if key.endswith(".gamma"):
+            converted[key.removesuffix(".gamma") + ".weight"] = converted.pop(
+                key
+            )
+        elif key.endswith(".beta"):
+            converted[key.removesuffix(".beta") + ".bias"] = converted.pop(
+                key
+            )
+
+    attention_suffix = ".multi_head_attention.query.weight"
+    query_weight_keys = [
+        key for key in converted if key.endswith(attention_suffix)
+    ]
+    for query_weight_key in query_weight_keys:
+        prefix = query_weight_key.removesuffix("query.weight")
+        qkv_weight_key = prefix + "qkv.weight"
+        qkv_bias_key = prefix + "qkv.bias"
+
+        converted[qkv_weight_key] = torch.cat(
+            [
+                converted.pop(prefix + "query.weight"),
+                converted.pop(prefix + "key.weight"),
+                converted.pop(prefix + "value.weight"),
+            ],
+            dim=0,
+        )
+        converted[qkv_bias_key] = torch.cat(
+            [
+                converted.pop(prefix + "query.bias"),
+                converted.pop(prefix + "key.bias"),
+                converted.pop(prefix + "value.bias"),
+            ],
+            dim=0,
+        )
 
-        x = x + feed_forward
+    return converted
 
-        return x
+__all__ = [
+    "convert_reference_state_dict",
+    "FeedForward",
+    "GPTModel",
+    "Model",
+    "MultiHeadAttention",
+    "TransformerBlock",
+]

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:9e9f8e95c7fa765cda29765de8ccd1c89d0f08545d42c72943b9c546665db657
-size 39973360
+oid sha256:fcd36c7f8d81cbb66a20f2b28a8b1afe03c99d3f57e110845755ee510aa35a07
+size 97146024

modeling_tiny_gpt.py

@@ -4,16 +4,17 @@ from transformers.generation import GenerationMixin
 from transformers.modeling_outputs import CausalLMOutput
 
 from .configuration_tiny_gpt import TinyGPTConfig
-from .model import Model
+from .model import GPTModel
 
 
 class TinyGPTForCausalLM(PreTrainedModel, GenerationMixin):
     config_class = TinyGPTConfig
     main_input_name = "input_ids"
+    _tied_weights_keys = {"core_model.linear.weight": "core_model.token_embedding.weight"}
 
     def __init__(self, config):
         super().__init__(config)
-        self.core_model = Model(
+        self.core_model = GPTModel(
             context_size=config.context_size,
             vocab_size=config.vocab_size,
             d_model=config.d_model,

tokenizer_config.json

@@ -1,7 +1,7 @@
 {
   "backend": "tokenizers",
   "eos_token": "<EOS>",
-  "model_max_length": 512,
+  "model_max_length": 1024,
   "pad_token": "<EOS>",
   "tokenizer_class": "TokenizersBackend"
 }