rewrite tinystories example

examples/tinystories-llm/train.py

@@ -1,6 +1,7 @@
 import argparse
 import random
 import urllib.request
+from dataclasses import dataclass
 from pathlib import Path
 
 import torchmlx as torch
@@ -8,125 +9,6 @@ from torchmlx import nn, optim
 import torchmlx.nn.functional as F
 
 
-DATA_URL = "https://huggingface.co/datasets/roneneldan/TinyStories/resolve/main/TinyStoriesV2-GPT4-valid.txt"
-DATA_PATH = Path(__file__).with_name("TinyStoriesV2-GPT4-valid.txt")
-
-
-class Attention(nn.Module):
-    def __init__(self, width, heads):
-        super().__init__()
-        self.heads = heads
-        self.head_width = width // heads
-        self.qkv = nn.Linear(width, width * 3)
-        self.output = nn.Linear(width, width)
-
-    def forward(self, x):
-        batch, length, width = x.shape
-        q, k, v = torch.chunk(self.qkv(x), 3, dim=-1)
-        q = torch.transpose(
-            q.reshape(batch, length, self.heads, self.head_width), 1, 2
-        )
-        k = torch.transpose(
-            k.reshape(batch, length, self.heads, self.head_width), 1, 2
-        )
-        v = torch.transpose(
-            v.reshape(batch, length, self.heads, self.head_width), 1, 2
-        )
-        x = F.scaled_dot_product_attention(q, k, v, is_causal=True)
-        x = torch.transpose(x, 1, 2).reshape(batch, length, width)
-        return self.output(x)
-
-
-class Block(nn.Module):
-    def __init__(self, width, heads):
-        super().__init__()
-        self.attention_norm = nn.LayerNorm(width)
-        self.attention = Attention(width, heads)
-        self.feed_forward_norm = nn.LayerNorm(width)
-        self.feed_forward = nn.Sequential(
-            nn.Linear(width, width * 4),
-            nn.GELU(),
-            nn.Linear(width * 4, width),
-        )
-
-    def forward(self, x):
-        x = x + self.attention(self.attention_norm(x))
-        return x + self.feed_forward(self.feed_forward_norm(x))
-
-
-class TinyStoriesModel(nn.Module):
-    def __init__(self, vocab_size, context, width=64, heads=4, layers=2):
-        super().__init__()
-        self.context = context
-        self.token_embedding = nn.Embedding(vocab_size, width)
-        self.position_embedding = nn.Embedding(context, width)
-        self.blocks = nn.ModuleList([Block(width, heads) for _ in range(layers)])
-        self.norm = nn.LayerNorm(width)
-        self.output = nn.Linear(width, vocab_size)
-
-    def forward(self, tokens):
-        positions = torch.arange(tokens.shape[1], device=device)
-        x = self.token_embedding(tokens) + self.position_embedding(positions)
-        for block in self.blocks:
-            x = block(x)
-        return self.output(self.norm(x))
-
-
-def download_data():
-    if not DATA_PATH.exists():
-        print(f"downloading TinyStories to {DATA_PATH}")
-        urllib.request.urlretrieve(DATA_URL, DATA_PATH)
-    return DATA_PATH.read_text(encoding="utf-8")
-
-
-def batch(encoded, batch_size, context):
-    starts = [random.randrange(len(encoded) - context - 1) for _ in range(batch_size)]
-    x = [encoded[start : start + context] for start in starts]
-    y = [encoded[start + 1 : start + context + 1] for start in starts]
-    return torch.tensor(x, dtype=torch.int64, device=device), torch.tensor(
-        y, dtype=torch.int64, device=device
-    )
-
-
-def generate(model, prompt, encode, decode, length):
-    tokens = encode(prompt)
-    for _ in range(length):
-        x = torch.tensor([tokens[-model.context :]], dtype=torch.int64, device=device)
-        logits = model(x)[:, -1, :]
-        token = torch.categorical(logits).item()
-        tokens.append(token)
-    return decode(tokens)
-
-
-def main():
-    parser = argparse.ArgumentParser()
-    parser.add_argument("steps", nargs="?", type=int, default=20)
-    parser.add_argument("--no-compile", action="store_true")
-    args = parser.parse_args()
-    text = download_data()
-    characters = sorted(set(text))
-    to_id = {character: index for index, character in enumerate(characters)}
-    encode = lambda value: [to_id[character] for character in value]
-    decode = lambda values: "".join(characters[int(value)] for value in values)
-    encoded = encode(text)
-    context = 64
-    model = TinyStoriesModel(len(characters), context).to(device)
-    optimizer = optim.AdamW(model.parameters(), lr=3e-4)
-    loss_fn = lambda logits, targets: F.cross_entropy(
-        logits.reshape(-1, len(characters)), targets.reshape(-1)
-    )
-    trainer = torch.Trainer(
-        model, optimizer, loss_fn, compile=not args.no_compile
-    )
-    interval = max(1, args.steps // 10)
-    for step in range(args.steps):
-        x, y = batch(encoded, 8, context)
-        loss = trainer.step(x, y)
-        if step % interval == 0 or step == args.steps - 1:
-            print(f"step {step + 1}: loss {loss.item():.4f}")
-    print(generate(model, "Once upon a time", encode, decode, 120))
-
-
 device = torch.device(
     "mps"
     if torch.backends.mps.is_available()
@@ -135,6 +17,231 @@ device = torch.device(
     else "cpu"
 )
 
+DATA_URL = "https://huggingface.co/datasets/roneneldan/TinyStories/resolve/main/TinyStoriesV2-GPT4-valid.txt"
+DATA_PATH = Path(__file__).with_name("TinyStoriesV2-GPT4-valid.txt")
 
-if __name__ == "__main__":
+
-    main()
+@dataclass
+class TransformerConfig:
+    vocabulary_size: int
+    context_length: int = 64
+    model_dimension: int = 64
+    head_count: int = 4
+    layer_count: int = 2
+    feed_forward_dimension: int = 256
+
+
+class CharacterTokenizer:
+    def __init__(self, text):
+        self.characters = sorted(set(text))
+        self.token_by_character = {
+            character: token for token, character in enumerate(self.characters)
+        }
+
+    def encode(self, text):
+        return [self.token_by_character[character] for character in text]
+
+    def decode(self, tokens):
+        return "".join(self.characters[int(token)] for token in tokens)
+
+    def __len__(self):
+        return len(self.characters)
+
+
+class TokenEmbedding(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.embedding = nn.Embedding(
+            config.vocabulary_size, config.model_dimension
+        )
+
+    def forward(self, tokens):
+        return self.embedding(tokens)
+
+
+class PositionalEmbedding(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.embedding = nn.Embedding(
+            config.context_length, config.model_dimension
+        )
+
+    def forward(self, tokens):
+        positions = torch.arange(tokens.shape[1], device=tokens.device)
+        return self.embedding(positions)
+
+
+class CausalSelfAttention(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.head_count = config.head_count
+        self.head_dimension = config.model_dimension // config.head_count
+        self.query = nn.Linear(config.model_dimension, config.model_dimension)
+        self.key = nn.Linear(config.model_dimension, config.model_dimension)
+        self.value = nn.Linear(config.model_dimension, config.model_dimension)
+        self.output = nn.Linear(config.model_dimension, config.model_dimension)
+
+    def split_heads(self, hidden_states):
+        batch_size, sequence_length, _ = hidden_states.shape
+        hidden_states = hidden_states.reshape(
+            batch_size, sequence_length, self.head_count, self.head_dimension
+        )
+        return hidden_states.transpose(1, 2)
+
+    def merge_heads(self, hidden_states):
+        batch_size, _, sequence_length, _ = hidden_states.shape
+        hidden_states = hidden_states.transpose(1, 2).contiguous()
+        return hidden_states.reshape(
+            batch_size, sequence_length, self.head_count * self.head_dimension
+        )
+
+    def forward(self, hidden_states):
+        queries = self.split_heads(self.query(hidden_states))
+        keys = self.split_heads(self.key(hidden_states))
+        values = self.split_heads(self.value(hidden_states))
+
+        attention_scores = torch.matmul(queries, keys.transpose(-2, -1))
+        attention_scores = attention_scores / self.head_dimension**0.5
+
+        sequence_length = hidden_states.shape[1]
+        future_positions = torch.triu(
+            torch.ones(
+                sequence_length,
+                sequence_length,
+                dtype=torch.bool,
+                device=hidden_states.device,
+            ),
+            diagonal=1,
+        )
+        attention_scores = attention_scores.masked_fill(
+            future_positions, float("-inf")
+        )
+        attention_weights = F.softmax(attention_scores, dim=-1)
+        attended_values = torch.matmul(attention_weights, values)
+        return self.output(self.merge_heads(attended_values))
+
+
+class FeedForward(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.layers = nn.Sequential(
+            nn.Linear(config.model_dimension, config.feed_forward_dimension),
+            nn.GELU(),
+            nn.Linear(config.feed_forward_dimension, config.model_dimension),
+        )
+
+    def forward(self, hidden_states):
+        return self.layers(hidden_states)
+
+
+class Unembedding(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.output = nn.Linear(
+            config.model_dimension, config.vocabulary_size, bias=False
+        )
+
+    def forward(self, hidden_states):
+        return self.output(hidden_states)
+
+
+class TransformerBlock(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.attention_norm = nn.LayerNorm(config.model_dimension)
+        self.attention = CausalSelfAttention(config)
+        self.feed_forward_norm = nn.LayerNorm(config.model_dimension)
+        self.feed_forward = FeedForward(config)
+
+    def forward(self, hidden_states):
+        hidden_states = hidden_states + self.attention(
+            self.attention_norm(hidden_states)
+        )
+        return hidden_states + self.feed_forward(
+            self.feed_forward_norm(hidden_states)
+        )
+
+
+class TinyStoriesTransformer(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.token_embedding = TokenEmbedding(config)
+        self.position_embedding = PositionalEmbedding(config)
+        self.blocks = nn.ModuleList(
+            [TransformerBlock(config) for _ in range(config.layer_count)]
+        )
+        self.final_norm = nn.LayerNorm(config.model_dimension)
+        self.unembedding = Unembedding(config)
+
+    def forward(self, tokens):
+        hidden_states = self.token_embedding(tokens)
+        hidden_states = hidden_states + self.position_embedding(tokens)
+        for block in self.blocks:
+            hidden_states = block(hidden_states)
+        return self.unembedding(self.final_norm(hidden_states))
+
+
+def load_tiny_stories():
+    if not DATA_PATH.exists():
+        print(f"downloading TinyStories to {DATA_PATH}")
+        urllib.request.urlretrieve(DATA_URL, DATA_PATH)
+    return DATA_PATH.read_text(encoding="utf-8")
+
+
+def create_batch(encoded_text, batch_size, context_length):
+    starts = [
+        random.randrange(len(encoded_text) - context_length - 1)
+        for _ in range(batch_size)
+    ]
+    input_tokens = [
+        encoded_text[start : start + context_length] for start in starts
+    ]
+    target_tokens = [
+        encoded_text[start + 1 : start + context_length + 1]
+        for start in starts
+    ]
+    return torch.tensor(input_tokens, dtype=torch.long, device=device), torch.tensor(
+        target_tokens, dtype=torch.long, device=device
+    )
+
+
+def language_model_loss(logits, target_tokens):
+    vocabulary_size = logits.shape[-1]
+    return F.cross_entropy(
+        logits.reshape(-1, vocabulary_size), target_tokens.reshape(-1)
+    )
+
+
+def generate_text(model, tokenizer, prompt, token_count):
+    generated_tokens = tokenizer.encode(prompt)
+    model.eval()
+    for _ in range(token_count):
+        context = generated_tokens[-model.config.context_length :]
+        input_tokens = torch.tensor([context], dtype=torch.long, device=device)
+        next_token_logits = model(input_tokens)[:, -1, :]
+        next_token = torch.categorical(next_token_logits).item()
+        generated_tokens.append(next_token)
+    return tokenizer.decode(generated_tokens)
+
+
+parser = argparse.ArgumentParser()
+parser.add_argument("steps", nargs="?", type=int, default=20)
+arguments = parser.parse_args()
+
+training_text = load_tiny_stories()
+tokenizer = CharacterTokenizer(training_text)
+config = TransformerConfig(vocabulary_size=len(tokenizer))
+encoded_text = tokenizer.encode(training_text)
+model = TinyStoriesTransformer(config).to(device)
+optimizer = optim.AdamW(model.parameters(), lr=3e-4)
+trainer = torch.Trainer(model, optimizer, language_model_loss, compile=True)
+
+for step in range(arguments.steps):
+    input_tokens, target_tokens = create_batch(
+        encoded_text, batch_size=8, context_length=config.context_length
+    )
+    loss = trainer.step(input_tokens, target_tokens)
+    print(f"step {step + 1}: loss {loss.item():.4f}")
+
+print(generate_text(model, tokenizer, "Once upon a time", token_count=120))