Add training script

train_story10m.py

@@ -0,0 +1,330 @@
+#!/usr/bin/env python3
+"""
+StoryLM-10M: ~10.5M param LLaMA-style language model trained on TinyStories.
+
+Architecture:
+  d_model=256, n_heads=8, n_kv_heads=8 (MHA), n_layers=8
+  SwiGLU FFN (4x), RoPE, RMSNorm, vocab 8192, tied embed/head, ctx 512
+  Total: ~10.5M learnable parameters
+
+Training:
+  ~2B tokens from TinyStories, AdamW 3e-4, cosine + warmup
+  batch 64 (effective), seq 512, ~6,100 steps
+
+Usage:
+  python3 train_story10m.py --stage prepare   # download + tokenize
+  python3 train_story10m.py --stage train     # train
+  python3 train_story10m.py --stage eval      # evaluate
+  python3 train_story10m.py --stage all
+"""
+import os, sys, math, json, time, argparse, glob
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+# ============================================================================
+# Config
+# ============================================================================
+VOCAB = 8192
+D_MODEL = 256
+N_HEADS = 8
+N_KV_HEADS = 8
+N_LAYERS = 8
+HEAD_DIM = D_MODEL // N_HEADS  # 32
+KV_DIM = N_KV_HEADS * HEAD_DIM  # 256
+FFN_DIM = D_MODEL * 4  # 1024
+SEQ_LEN = 512
+BATCH = 4
+LR = 3e-4
+WARMUP_STEPS = 500
+TOTAL_STEPS = 24400
+CHECKPOINT_EVERY = 500
+OUT_DIR = "/work/story10m"
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+TOKENIZER_PATH = "/work/tokenizer.json"
+
+os.makedirs(OUT_DIR, exist_ok=True)
+
+# ============================================================================
+# Architecture
+# ============================================================================
+class RMSNorm(nn.Module):
+    def __init__(self, dim, eps=1e-5):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+    def forward(self, x):
+        norm = x.float().pow(2).mean(-1, keepdim=True).add(self.eps).rsqrt()
+        return (x * norm).type_as(x) * self.weight
+
+def precompute_rope(dim, max_seq, base=10000.0):
+    freqs = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
+    t = torch.arange(max_seq).float()
+    angles = torch.outer(t, freqs)
+    return torch.stack([angles.cos(), angles.sin()], dim=-1)
+
+def apply_rope(x, cos, sin):
+    # x: (B, H, T, D)
+    B, H, T, D = x.shape
+    x1 = x[..., 0::2]
+    x2 = x[..., 1::2]
+    c = cos[:T].unsqueeze(0).unsqueeze(0)  # (1, 1, T, D/2)
+    s = sin[:T].unsqueeze(0).unsqueeze(0)
+    out = torch.empty_like(x)
+    out[..., 0::2] = x1 * c - x2 * s
+    out[..., 1::2] = x1 * s + x2 * c
+    return out
+
+class Attention(nn.Module):
+    def __init__(self, d_model, n_heads, n_kv_heads, head_dim):
+        super().__init__()
+        self.n_heads = n_heads
+        self.n_kv_heads = n_kv_heads
+        self.head_dim = head_dim
+        self.qkv = nn.Linear(d_model, (n_heads + 2 * n_kv_heads) * head_dim, bias=False)
+        self.o = nn.Linear(n_heads * head_dim, d_model, bias=False)
+
+    def forward(self, x, cos, sin):
+        B, T, _ = x.shape
+        qkv = self.qkv(x)
+        q, k, v = qkv.split([self.n_heads * self.head_dim,
+                             self.n_kv_heads * self.head_dim,
+                             self.n_kv_heads * self.head_dim], dim=-1)
+        q = q.view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
+        k = k.view(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2)
+        v = v.view(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2)
+        if self.n_kv_heads != self.n_heads:
+            rep = self.n_heads // self.n_kv_heads
+            k = k.repeat_interleave(rep, dim=1)
+            v = v.repeat_interleave(rep, dim=1)
+        q = apply_rope(q, cos, sin)
+        k = apply_rope(k, cos, sin)
+        attn = F.scaled_dot_product_attention(q, k, v, is_causal=True)
+        attn = attn.transpose(1, 2).contiguous().view(B, T, -1)
+        return self.o(attn)
+
+class SwiGLU(nn.Module):
+    def __init__(self, d_model, ffn_dim):
+        super().__init__()
+        self.gate = nn.Linear(d_model, ffn_dim, bias=False)
+        self.up = nn.Linear(d_model, ffn_dim, bias=False)
+        self.down = nn.Linear(ffn_dim, d_model, bias=False)
+    def forward(self, x):
+        return self.down(F.silu(self.gate(x)) * self.up(x))
+
+class Block(nn.Module):
+    def __init__(self, d_model, n_heads, n_kv_heads, head_dim, ffn_dim):
+        super().__init__()
+        self.norm1 = RMSNorm(d_model)
+        self.attn = Attention(d_model, n_heads, n_kv_heads, head_dim)
+        self.norm2 = RMSNorm(d_model)
+        self.ffn = SwiGLU(d_model, ffn_dim)
+    def forward(self, x, cos, sin):
+        x = x + self.attn(self.norm1(x), cos, sin)
+        x = x + self.ffn(self.norm2(x))
+        return x
+
+class StoryLM(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.embed = nn.Embedding(VOCAB, D_MODEL)
+        self.layers = nn.ModuleList([
+            Block(D_MODEL, N_HEADS, N_KV_HEADS, HEAD_DIM, FFN_DIM)
+            for _ in range(N_LAYERS)
+        ])
+        self.norm = RMSNorm(D_MODEL)
+        self.lm_head = nn.Linear(D_MODEL, VOCAB, bias=False)
+        self.lm_head.weight = self.embed.weight  # tied
+        rope = precompute_rope(HEAD_DIM, SEQ_LEN)
+        self.register_buffer('rope_cos', rope[..., 0])
+        self.register_buffer('rope_sin', rope[..., 1])
+
+    def forward(self, x, targets=None):
+        h = self.embed(x)
+        for block in self.layers:
+            h = block(h, self.rope_cos, self.rope_sin)
+        h = self.norm(h)
+        logits = self.lm_head(h)
+        loss = None
+        if targets is not None:
+            loss = F.cross_entropy(logits.view(-1, VOCAB), targets.view(-1))
+        return logits, loss
+
+# ============================================================================
+# Data
+# ============================================================================
+def prepare_data():
+    """Download TinyStories and tokenize."""
+    print("Loading tokenizer...")
+    from tokenizers import Tokenizer
+    tok = Tokenizer.from_file(TOKENIZER_PATH)
+    
+    print("Downloading TinyStories...")
+    import urllib.request
+    stories_path = os.path.join(OUT_DIR, "tinystories.txt")
+    if not os.path.exists(stories_path):
+        url = "https://huggingface.co/datasets/roneneldan/TinyStories/resolve/main/TinyStoriesv2_2M.txt"
+        urllib.request.urlretrieve(url, stories_path)
+    
+    with open(stories_path, 'r') as f:
+        stories = f.readlines()
+    print(f"  {len(stories)} stories")
+    
+    print("Tokenizing...")
+    all_tokens = []
+    for i, story in enumerate(stories):
+        story = story.strip()
+        if len(story) < 20:
+            continue
+        ids = tok.encode(story, add_special_tokens=False).ids
+        all_tokens.extend(ids)
+        if (i+1) % 100000 == 0:
+            print(f"  {i+1}/{len(stories)} stories, {len(all_tokens)/1e6:.1f}M tokens so far")
+    
+    tokens = np.array(all_tokens, dtype=np.uint16)
+    np.save(os.path.join(OUT_DIR, "tokens.npy"), tokens)
+    print(f"  Total: {len(tokens)/1e6:.1f}M tokens")
+    print(f"  Saved to {OUT_DIR}/tokens.npy")
+
+def load_data():
+    tokens = np.load(os.path.join(OUT_DIR, "tokens.npy"))
+    print(f"Loaded {len(tokens)/1e6:.1f}M tokens")
+    # 99% train, 1% val
+    split = int(len(tokens) * 0.99)
+    train = tokens[:split]
+    val = tokens[split:]
+    print(f"Train: {len(train)/1e6:.1f}M, Val: {len(val)/1e6:.1f}M")
+    return train, val
+
+def get_batch(data, batch_size, seq_len):
+    ix = torch.randint(len(data) - seq_len - 1, (batch_size,))
+    x = torch.stack([torch.from_numpy((data[i:i+seq_len]).astype(np.int64)) for i in ix])
+    y = torch.stack([torch.from_numpy((data[i+1:i+1+seq_len]).astype(np.int64)) for i in ix])
+    return x.to(DEVICE), y.to(DEVICE)
+
+# ============================================================================
+# Training
+# ============================================================================
+def train():
+    train_data, val_data = load_data()
+    model = StoryLM().to(DEVICE)
+    total_params = sum(p.numel() for p in model.parameters())
+    print(f"Parameters: {total_params:,}")
+    
+    optimizer = torch.optim.AdamW(model.parameters(), lr=LR, weight_decay=0.0, betas=(0.9, 0.95))
+    
+    best_val = float('inf')
+    t0 = time.time()
+    
+    for step in range(1, TOTAL_STEPS + 1):
+        # LR schedule: warmup + cosine
+        if step < WARMUP_STEPS:
+            lr = LR * step / WARMUP_STEPS
+        else:
+            progress = (step - WARMUP_STEPS) / (TOTAL_STEPS - WARMUP_STEPS)
+            lr = LR * 0.5 * (1 + math.cos(math.pi * progress))
+        for param_group in optimizer.param_groups:
+            param_group['lr'] = lr
+        
+        x, y = get_batch(train_data, BATCH, SEQ_LEN)
+        _, loss = model(x, y)
+        optimizer.zero_grad()
+        loss.backward()
+        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+        optimizer.step()
+        
+        if step % 100 == 0:
+            elapsed = time.time() - t0
+            tok_per_s = (step * BATCH * SEQ_LEN) / elapsed
+            print(f"step {step}/{TOTAL_STEPS} | loss {loss.item():.4f} | lr {lr:.2e} | {tok_per_s/1000:.1f}k tok/s | {elapsed:.0f}s", flush=True)
+        
+        if step % CHECKPOINT_EVERY == 0:
+            # Eval
+            model.eval()
+            val_losses = []
+            for _ in range(20):
+                vx, vy = get_batch(val_data, BATCH, SEQ_LEN)
+                _, vloss = model(vx, vy)
+                val_losses.append(vloss.item())
+            val_loss = sum(val_losses) / len(val_losses)
+            model.train()
+            print(f"  val_loss: {val_loss:.4f}")
+            
+            ckpt_path = os.path.join(OUT_DIR, f"ckpt_{step}.pt")
+            torch.save(model.state_dict(), ckpt_path)
+            print(f"  checkpoint: {ckpt_path}")
+            
+            if val_loss < best_val:
+                best_val = val_loss
+                torch.save(model.state_dict(), os.path.join(OUT_DIR, "best.pt"))
+                print(f"  new best: {val_loss:.4f}")
+    
+    # Save final
+    torch.save(model.state_dict(), os.path.join(OUT_DIR, "final.pt"))
+    print(f"\nTraining done. Best val_loss: {best_val:.4f}")
+    print(f"Final model: {OUT_DIR}/final.pt")
+
+# ============================================================================
+# Evaluation
+# ============================================================================
+def evaluate():
+    model = StoryLM().to(DEVICE)
+    ckpt = torch.load(os.path.join(OUT_DIR, "best.pt"), map_location=DEVICE)
+    model.load_state_dict(ckpt)
+    model.eval()
+    
+    _, val_data = load_data()
+    
+    # Perplexity
+    val_losses = []
+    for _ in range(50):
+        vx, vy = get_batch(val_data, BATCH, SEQ_LEN)
+        _, vloss = model(vx, vy)
+        val_losses.append(vloss.item())
+    avg_loss = sum(val_losses) / len(val_losses)
+    ppl = math.exp(avg_loss)
+    print(f"Val perplexity: {ppl:.2f} (loss {avg_loss:.4f})")
+    
+    # Generation samples
+    from tokenizers import Tokenizer
+    tok = Tokenizer.from_file(TOKENIZER_PATH)
+    prompts = [
+        "Once upon a time, there was a little",
+        "The cat sat on the",
+        "In the beginning, the world was",
+        "A small robot named",
+        "Every morning, the sun",
+    ]
+    print("\n=== Generation Samples ===")
+    for prompt in prompts:
+        ids = tok.encode(prompt, add_special_tokens=False).ids
+        x = torch.tensor([ids], dtype=torch.long, device=DEVICE)
+        with torch.no_grad():
+            for _ in range(30):
+                out, _ = model(x)
+                next_id = out[0, -1].argmax().item()
+                x = torch.cat([x, torch.tensor([[next_id]], device=DEVICE)], dim=1)
+        full_ids = ids + [next_id]
+        text = tok.decode(full_ids, skip_special_tokens=True)
+        print(f"Prompt: {prompt}")
+        print(f"Output: {text}")
+        print()
+
+# ============================================================================
+# Main
+# ============================================================================
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--stage", choices=["prepare", "train", "eval", "all"], default="all")
+    args = parser.parse_args()
+    
+    if args.stage in ("prepare", "all"):
+        prepare_data()
+    if args.stage in ("train", "all"):
+        train()
+    if args.stage in ("eval", "all"):
+        evaluate()
+
+if __name__ == "__main__":
+    main()