train_blochwalk_200k.py

#!/usr/bin/env python3
"""
Bloch Walk training with 200K sequences and 5 epochs.
"""
import matplotlib
matplotlib.use('Agg')

import numpy as np
import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
from pathlib import Path

from belief_geometry.processes import generate_blochwalk_sequence
from belief_geometry.model import ResidualHookedTransformer


class BlochwalkDataset(Dataset):
    """Dataset for Bloch Walk."""
    def __init__(self, num_sequences, max_ctx, seed=42):
        self.num_sequences = num_sequences
        self.max_ctx = max_ctx
        self.rng = np.random.default_rng(seed)
        self._generate()
    
    def _generate(self):
        L = self.max_ctx
        N = self.num_sequences
        tokens = np.zeros((N, L), dtype=np.int64)
        beliefs = np.zeros((N, L, 3), dtype=np.float32)
        
        print(f"  Generating {N} sequences of length {L}...")
        for i in range(N):
            if (i + 1) % 10000 == 0:
                print(f"    {i+1}/{N}")
            s = generate_blochwalk_sequence(L, self.rng)
            tokens[i] = s["tokens"]
            beliefs[i] = s["beliefs"].astype(np.float32)
        
        x = torch.from_numpy(tokens[:, :-1].copy())
        y = torch.from_numpy(tokens[:, 1:].copy())
        
        self.samples = {
            "x": x,
            "y": y,
            "beliefs": torch.from_numpy(beliefs[:, 1:].copy()),
        }
    
    def __len__(self):
        return len(self.samples["x"])
    
    def __getitem__(self, idx):
        return {k: v[idx] for k, v in self.samples.items()}


def train():
    """Train transformer on Bloch Walk."""
    print("=" * 70)
    print("Bloch Walk Training: 200K sequences, 5 epochs")
    print("=" * 70)
    print()
    
    # Config
    num_sequences = 200_000
    max_ctx = 8
    vocab_size = 4
    batch_size = 256
    lr = 3e-4
    epochs = 5
    device = "cuda" if torch.cuda.is_available() else "cpu"
    out_dir = Path("runs/blochwalk_200k")
    out_dir.mkdir(parents=True, exist_ok=True)
    
    print(f"Device: {device}")
    print(f"Output: {out_dir}")
    print()
    
    # Generate data
    print("Generating dataset...")
    full_ds = BlochwalkDataset(num_sequences, max_ctx, seed=42)
    
    # Split
    n_train = int(0.90 * len(full_ds))
    n_val = int(0.05 * len(full_ds))
    n_test = len(full_ds) - n_train - n_val
    
    train_ds, val_ds, test_ds = torch.utils.data.random_split(
        full_ds, [n_train, n_val, n_test],
        generator=torch.Generator().manual_seed(42)
    )
    
    print(f"  Train: {len(train_ds)}, Val: {len(val_ds)}, Test: {len(test_ds)}")
    print()
    
    # Model
    print("Creating model...")
    model = ResidualHookedTransformer(
        vocab_size=vocab_size,
        d_model=128,
        n_heads=4,
        d_ff=512,
        n_layers=2,
        max_ctx=max_ctx - 1,
        dropout=0.0,
    ).to(device)
    
    print(f"  Parameters: {sum(p.numel() for p in model.parameters()):,}")
    print()
    
    # Training
    optimizer = torch.optim.Adam(model.parameters(), lr=lr)
    loss_fn = nn.CrossEntropyLoss()
    
    train_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True)
    val_loader = DataLoader(val_ds, batch_size=batch_size, shuffle=False)
    
    print("Training...")
    for epoch in range(epochs):
        # Train
        model.train()
        train_loss = 0.0
        for batch in train_loader:
            x = batch["x"].to(device)
            y = batch["y"].to(device)
            
            optimizer.zero_grad()
            logits, _ = model(x)
            loss = loss_fn(logits.reshape(-1, vocab_size), y.reshape(-1))
            loss.backward()
            optimizer.step()
            
            train_loss += loss.item()
        
        train_loss /= len(train_loader)
        
        # Val
        model.eval()
        val_loss = 0.0
        with torch.no_grad():
            for batch in val_loader:
                x = batch["x"].to(device)
                y = batch["y"].to(device)
                logits, _ = model(x)
                loss = loss_fn(logits.reshape(-1, vocab_size), y.reshape(-1))
                val_loss += loss.item()
        
        val_loss /= len(val_loader)
        print(f"  Epoch {epoch+1}/{epochs}: train_loss={train_loss:.4f}, val_loss={val_loss:.4f}")
    
    print()
    print("Saving model...")
    torch.save(model.state_dict(), out_dir / "model.pt")
    
    # Export analysis batch
    print("Exporting analysis batch...")
    model.eval()
    export_samples = min(2048, len(val_ds))
    loader = DataLoader(val_ds, batch_size=256, shuffle=False)
    
    residuals = []
    beliefs = []
    tokens_list = []
    
    with torch.no_grad():
        for batch in loader:
            x = batch["x"].to(device)
            logits, res = model(x)
            residuals.append(res.cpu())
            beliefs.append(batch["beliefs"])
            tokens_list.append(batch["x"])
            if sum(r.size(0) for r in residuals) >= export_samples:
                break
    
    residuals = torch.cat(residuals, dim=0)[:export_samples]
    beliefs = torch.cat(beliefs, dim=0)[:export_samples]
    tokens_list = torch.cat(tokens_list, dim=0)[:export_samples]
    
    torch.save({
        "residuals": residuals,
        "beliefs": beliefs,
        "tokens": tokens_list,
    }, out_dir / "analysis_batch.pt")
    
    print(f"  Saved {export_samples} samples")
    print()
    print("✓ Training complete!")


if __name__ == "__main__":
    train()