Add train.py

train.py

@@ -0,0 +1,424 @@
+"""
+Training script for Resonance 200M.
+ClimbMix data, own BPE tokenizer (Rust backend), AdamW optimizer.
+Shows BOTH train loss AND val loss. Always.
+"""
+
+import os
+import sys
+import time
+import math
+import struct
+import argparse
+import numpy as np
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.amp import autocast, GradScaler
+
+from model import Resonance, RESONANCE_200M
+from bpe_tokenizer import BPETokenizer
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# Data
+# ─────────────────────────────────────────────────────────────────────────────
+
+def download_climbmix_shards(data_dir, n_shards=100):
+    """Download ClimbMix parquet shards from HuggingFace."""
+    os.makedirs(data_dir, exist_ok=True)
+
+    try:
+        import pyarrow.parquet as pq
+    except ImportError:
+        print("pip install pyarrow pandas")
+        sys.exit(1)
+
+    base_url = "https://huggingface.co/datasets/karpathy/climbmix-400b-shuffle/resolve/main"
+    texts_path = os.path.join(data_dir, "texts.txt")
+
+    if os.path.exists(texts_path):
+        size = os.path.getsize(texts_path)
+        print(f"  [Data] texts.txt exists ({size/1e9:.2f} GB), skipping download")
+        return texts_path
+
+    import urllib.request
+    import ssl
+    ctx = ssl.create_default_context()
+    ctx.check_hostname = False
+    ctx.verify_mode = ssl.CERT_NONE
+
+    total_bytes = 0
+    with open(texts_path, 'w', encoding='utf-8') as out:
+        for i in range(n_shards):
+            shard_name = f"shard_{i:05d}.parquet"
+            shard_path = os.path.join(data_dir, shard_name)
+            url = f"{base_url}/{shard_name}"
+
+            if not os.path.exists(shard_path):
+                print(f"  [Data] Downloading shard {i+1}/{n_shards}...", end=" ", flush=True)
+                try:
+                    urllib.request.urlretrieve(url, shard_path)
+                    print("OK")
+                except Exception as e:
+                    print(f"FAIL: {e}")
+                    continue
+
+            # Extract text
+            try:
+                table = pq.read_table(shard_path, columns=['text'])
+                texts = table.column('text').to_pylist()
+                for text in texts:
+                    if text and len(text) > 100:
+                        out.write(text + '\n')
+                        total_bytes += len(text)
+                # Remove parquet to save disk
+                os.remove(shard_path)
+            except Exception as e:
+                print(f"  [Data] Error reading shard {i}: {e}")
+                continue
+
+            if (i + 1) % 10 == 0:
+                print(f"  [Data] {i+1}/{n_shards} shards, {total_bytes/1e9:.2f} GB text")
+
+    print(f"  [Data] Total: {total_bytes/1e9:.2f} GB text from {n_shards} shards")
+    return texts_path
+
+
+def tokenize_data(texts_path, tokenizer, data_dir, context_len):
+    """Tokenize text into binary shards (uint16 for vocab < 65536).
+    Streams to disk — no OOM on 16GB+ corpora."""
+    train_path = os.path.join(data_dir, "train.bin")
+    val_path = os.path.join(data_dir, "val.bin")
+
+    if os.path.exists(train_path) and os.path.exists(val_path):
+        train_tokens = os.path.getsize(train_path) // 2
+        val_tokens = os.path.getsize(val_path) // 2
+        print(f"  [Data] Tokenized data exists: train={train_tokens:,} val={val_tokens:,}")
+        return train_tokens, val_tokens
+
+    print(f"  [Data] Tokenizing...")
+    tmp_path = os.path.join(data_dir, "tokens_all.bin")
+    total_tokens = 0
+    t0 = time.time()
+
+    with open(texts_path, 'r', encoding='utf-8', errors='replace') as f_in, \
+         open(tmp_path, 'wb') as f_out:
+        chunk_size = 10_000_000  # 10MB chunks
+        total_chars = 0
+        while True:
+            text = f_in.read(chunk_size)
+            if not text:
+                break
+            ids = tokenizer.encode(text)
+            arr = np.array(ids, dtype=np.uint16)
+            f_out.write(arr.tobytes())
+            total_tokens += len(ids)
+            total_chars += len(text)
+            if total_chars % 100_000_000 < chunk_size:
+                elapsed = time.time() - t0
+                rate = total_chars / elapsed / 1e6
+                print(f"  [Data] {total_chars/1e9:.2f} GB text → {total_tokens:,} tokens "
+                      f"({rate:.1f} MB/s, {elapsed:.0f}s)")
+
+    elapsed = time.time() - t0
+    print(f"  [Data] Tokenized {total_chars/1e9:.2f} GB → {total_tokens:,} tokens in {elapsed:.0f}s")
+
+    # Split 95/5 train/val — stream from memmap to avoid loading all into RAM
+    split = int(total_tokens * 0.95)
+    print(f"  [Data] Splitting: train={split:,} val={total_tokens - split:,}")
+
+    all_data = np.memmap(tmp_path, dtype=np.uint16, mode='r')
+
+    # Write train split in chunks
+    chunk = 50_000_000  # 50M tokens per chunk
+    with open(train_path, 'wb') as f:
+        for start in range(0, split, chunk):
+            end = min(start + chunk, split)
+            f.write(all_data[start:end].tobytes())
+
+    # Write val split
+    with open(val_path, 'wb') as f:
+        for start in range(split, total_tokens, chunk):
+            end = min(start + chunk, total_tokens)
+            f.write(all_data[start:end].tobytes())
+
+    del all_data
+    os.remove(tmp_path)
+
+    train_tokens = split
+    val_tokens = total_tokens - split
+    print(f"  [Data] train: {train_tokens:,} tokens ({train_tokens*2/1e9:.2f} GB)")
+    print(f"  [Data] val:   {val_tokens:,} tokens ({val_tokens*2/1e9:.2f} GB)")
+    return train_tokens, val_tokens
+
+
+class DataLoader:
+    """Simple random-chunk dataloader from mmap'd binary file."""
+
+    def __init__(self, path, context_len, batch_size, device):
+        self.data = np.memmap(path, dtype=np.uint16, mode='r')
+        self.context_len = context_len
+        self.batch_size = batch_size
+        self.device = device
+        self.n_tokens = len(self.data)
+
+    def get_batch(self):
+        T = self.context_len
+        B = self.batch_size
+        ix = torch.randint(0, self.n_tokens - T - 1, (B,))
+        x = torch.stack([torch.from_numpy(self.data[i:i+T].astype(np.int64)) for i in ix])
+        y = torch.stack([torch.from_numpy(self.data[i+1:i+T+1].astype(np.int64)) for i in ix])
+        return x.to(self.device), y.to(self.device)
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# Training
+# ─────────────────────────────────────────────────────────────────────────────
+
+def get_lr(step, warmup_steps, total_steps, max_lr, min_lr=0.0):
+    """WSD schedule: warmup → stable → linear decay."""
+    if step < warmup_steps:
+        return max_lr * (step + 1) / warmup_steps
+    decay_start = total_steps // 2
+    if step < decay_start:
+        return max_lr
+    # Linear decay
+    progress = (step - decay_start) / (total_steps - decay_start)
+    return max_lr * (1.0 - progress) + min_lr * progress
+
+
+@torch.no_grad()
+def evaluate(model, val_loader, n_batches=50):
+    """Evaluate val loss. Returns average loss."""
+    model.eval()
+    losses = []
+    for _ in range(n_batches):
+        x, y = val_loader.get_batch()
+        with autocast('cuda', dtype=torch.bfloat16):
+            _, loss = model(x, y)
+        losses.append(loss.item())
+    model.train()
+    return sum(losses) / len(losses)
+
+
+def save_checkpoint(model, optimizer, step, train_loss, val_loss, config, path):
+    """Save PyTorch checkpoint."""
+    torch.save({
+        'model': model.state_dict(),
+        'optimizer': optimizer.state_dict(),
+        'step': step,
+        'train_loss': train_loss,
+        'val_loss': val_loss,
+        'config': config,
+    }, path)
+
+
+def save_c_weights(model, tokenizer, config, path):
+    """Save weights in C-compatible binary format for resonance-bpe.c."""
+    with open(path, 'wb') as f:
+        # Header: magic + config
+        f.write(struct.pack('<I', 0x52533032))  # "RS02"
+        f.write(struct.pack('<9i',
+            config['n_embd'], config['n_layer'], config['context_len'],
+            config['n_head'], config['head_dim'], config['rrpram_rank'],
+            config['ffn_dim'], config['vocab_size'], config['n_head']))  # kv_heads = n_head (MHA)
+
+        # BPE merges
+        f.write(struct.pack('<I', len(tokenizer.merges)))
+        for a, b, new_id in tokenizer.merges:
+            f.write(struct.pack('<III', a, b, new_id))
+
+        # All parameters in order
+        for name, param in model.named_parameters():
+            data = param.detach().float().cpu().numpy()
+            f.write(data.tobytes())
+
+    size_mb = os.path.getsize(path) / 1e6
+    print(f"  [Save] C weights: {path} ({size_mb:.1f} MB)")
+
+
+def train(args):
+    device = 'cuda' if torch.cuda.is_available() else 'cpu'
+    print(f"Device: {device}")
+
+    config = RESONANCE_200M.copy()
+    if args.vocab_size:
+        config['vocab_size'] = args.vocab_size
+
+    data_dir = args.data_dir
+    os.makedirs(data_dir, exist_ok=True)
+    os.makedirs(args.save_dir, exist_ok=True)
+
+    # Step 1: Download ClimbMix
+    print("\n[1] Data...")
+    texts_path = download_climbmix_shards(data_dir, n_shards=args.n_shards)
+
+    # Step 2: Train BPE tokenizer
+    print("\n[2] BPE tokenizer...")
+    tok_path = os.path.join(args.save_dir, "tokenizer.bin")
+    tokenizer = BPETokenizer(max_merges=config['vocab_size'] - 256)
+
+    if os.path.exists(tok_path):
+        tokenizer.load(tok_path)
+    else:
+        # Train on first 200MB of text
+        with open(texts_path, 'rb') as f:
+            sample = f.read(200_000_000)
+        tokenizer.train(sample, num_merges=config['vocab_size'] - 256, report_every=2000)
+        tokenizer.save_copies(tok_path, n=3)
+
+    config['vocab_size'] = tokenizer.vocab_size
+
+    # Step 3: Tokenize data
+    print("\n[3] Tokenizing data...")
+    n_train, n_val = tokenize_data(texts_path, tokenizer, data_dir, config['context_len'])
+
+    # Step 4: Build model
+    print("\n[4] Model...")
+    model = Resonance(config).to(device)
+    model.set_gradient_checkpointing(True)
+    model = torch.compile(model)
+    print(f"  Gradient checkpointing: ON, torch.compile: ON")
+
+    # Step 5: Optimizer
+    print("\n[5] Optimizer...")
+    # Separate param groups: decay vs no-decay
+    decay_params = []
+    no_decay_params = []
+    for name, p in model.named_parameters():
+        if p.dim() >= 2:
+            decay_params.append(p)
+        else:
+            no_decay_params.append(p)
+
+    optimizer = torch.optim.AdamW([
+        {'params': decay_params, 'weight_decay': args.weight_decay},
+        {'params': no_decay_params, 'weight_decay': 0.0},
+    ], lr=args.lr, betas=(0.9, 0.95), eps=1e-8)
+
+    scaler = GradScaler('cuda')
+
+    # Step 6: Data loaders (micro-batch for gradient accumulation)
+    T = config['context_len']
+    micro_B = args.micro_batch // T  # sequences per micro-batch
+    grad_accum = args.batch_size // args.micro_batch
+    print(f"\n[6] DataLoader: effective_batch={args.batch_size} tokens "
+          f"({grad_accum} x {args.micro_batch} micro), {micro_B} seq x {T} ctx")
+
+    train_loader = DataLoader(os.path.join(data_dir, "train.bin"), T, micro_B, device)
+    val_loader = DataLoader(os.path.join(data_dir, "val.bin"), T, micro_B, device)
+
+    total_steps = n_train // args.batch_size
+    print(f"  Total steps: {total_steps:,}")
+
+    # Step 7: Train loop
+    print(f"\n[7] Training resonance-200m...")
+    print(f"  {'step':>8} | {'train_loss':>10} | {'val_loss':>10} | {'lr':>10} | {'tok/s':>10} | {'time':>8}")
+    print("  " + "-" * 75)
+
+    best_val_loss = float('inf')
+    running_loss = 0.0
+    t0 = time.time()
+    tokens_seen = 0
+
+    model.train()
+    for step in range(total_steps):
+        # LR schedule
+        lr = get_lr(step, args.warmup_steps, total_steps, args.lr)
+        for pg in optimizer.param_groups:
+            pg['lr'] = lr
+
+        # Gradient accumulation: grad_accum micro-batches per optimizer step
+        optimizer.zero_grad(set_to_none=True)
+        step_loss = 0.0
+        for micro_step in range(grad_accum):
+            x, y = train_loader.get_batch()
+            with autocast('cuda', dtype=torch.bfloat16):
+                _, loss = model(x, y)
+            loss = loss / grad_accum
+            scaler.scale(loss).backward()
+            step_loss += loss.item() * grad_accum
+
+        scaler.unscale_(optimizer)
+        nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
+        scaler.step(optimizer)
+        scaler.update()
+
+        train_loss = step_loss / grad_accum
+        running_loss += train_loss
+        tokens_seen += args.batch_size
+
+        # Log every N steps
+        if (step + 1) % args.log_every == 0:
+            avg_train = running_loss / args.log_every
+            running_loss = 0.0
+            elapsed = time.time() - t0
+            tok_per_sec = tokens_seen / elapsed
+
+            # Val loss
+            val_loss = evaluate(model, val_loader, n_batches=args.val_batches)
+
+            print(f"  {step+1:>8} | {avg_train:>10.4f} | {val_loss:>10.4f} | "
+                  f"{lr:>10.2e} | {tok_per_sec/1000:>8.1f}k | {elapsed:>7.0f}s")
+
+            # Save best
+            if val_loss < best_val_loss:
+                best_val_loss = val_loss
+                save_checkpoint(model, optimizer, step, avg_train, val_loss, config,
+                              os.path.join(args.save_dir, "best.pt"))
+
+        # Checkpoint every N steps
+        if (step + 1) % args.save_every == 0:
+            save_checkpoint(model, optimizer, step, train_loss, val_loss if 'val_loss' in dir() else 0,
+                          config, os.path.join(args.save_dir, f"step_{step+1}.pt"))
+            save_c_weights(model, tokenizer, config,
+                         os.path.join(args.save_dir, f"resonance_200m_step{step+1}.bin"))
+
+        # Gate monitoring every N steps
+        if (step + 1) % (args.log_every * 5) == 0:
+            gates = []
+            for block in model._orig_mod.blocks if hasattr(model, '_orig_mod') else model.blocks:
+                g = torch.sigmoid(block.gate).detach().cpu().numpy()
+                gates.append(g.mean())
+            gate_str = " ".join(f"{g:.2f}" for g in gates)
+            print(f"  [gates] {gate_str}")
+
+    # Final save
+    elapsed = time.time() - t0
+    print(f"\n  Training complete. {elapsed/3600:.1f} hours, {tokens_seen:,} tokens")
+
+    save_checkpoint(model, optimizer, total_steps, train_loss, best_val_loss, config,
+                  os.path.join(args.save_dir, "final.pt"))
+    save_c_weights(model, tokenizer, config,
+                 os.path.join(args.save_dir, "resonance_200m_final.bin"))
+
+    # Re-save tokenizer (paranoia)
+    tokenizer.save_copies(os.path.join(args.save_dir, "tokenizer.bin"), n=3)
+
+    print(f"\n  Best val loss: {best_val_loss:.4f}")
+    print(f"  resonance is unbreakable.")
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--data-dir', type=str, default='data/')
+    parser.add_argument('--save-dir', type=str, default='checkpoints/')
+    parser.add_argument('--n-shards', type=int, default=65,
+                        help='Number of ClimbMix shards to download (~65 for ~4B tokens)')
+    parser.add_argument('--vocab-size', type=int, default=None,
+                        help='Override vocab size (default: 16384)')
+    parser.add_argument('--batch-size', type=int, default=131072,
+                        help='Effective batch size in tokens (default: 131072)')
+    parser.add_argument('--micro-batch', type=int, default=65536,
+                        help='Micro-batch size in tokens for grad accum (default: 65536)')
+    parser.add_argument('--lr', type=float, default=3e-4)
+    parser.add_argument('--warmup-steps', type=int, default=800)
+    parser.add_argument('--weight-decay', type=float, default=0.1)
+    parser.add_argument('--grad-clip', type=float, default=1.0)
+    parser.add_argument('--log-every', type=int, default=100)
+    parser.add_argument('--save-every', type=int, default=2000)
+    parser.add_argument('--val-batches', type=int, default=50)
+    args = parser.parse_args()
+    train(args)