train.py

"""
TinyThinker — train.py
Model definition and training loop.

This is the file autoresearch modifies. Keep it self-contained.
"""

import os
import sys
import time
import math
import json
from pathlib import Path

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader

from prepare import (
    prepare_data, collate_fn, evaluate_model,
    WordTokenizer, ReasoningDataset, DATA_DIR,
)

# ---------------------------------------------------------------------------
# Config
# ---------------------------------------------------------------------------

class Config:
    # Model
    n_layer: int = 6
    n_head: int = 8
    n_embd: int = 256
    dropout: float = 0.0  # small models are undertrained, not overtrained
    max_seq_len: int = 256

    # Training
    batch_size: int = 64
    learning_rate: float = 3e-4
    weight_decay: float = 0.01
    max_steps: int = 5000
    warmup_steps: int = 200
    eval_interval: int = 250
    eval_samples: int = 200  # how many test examples to eval on
    grad_clip: float = 1.0
    early_stop_window: int = 10  # rolling average window size
    early_stop_min_evals: int = 20  # minimum evals before early stopping can trigger
    early_stop_threshold: float = 0.005  # stop when rolling avg improves less than this over a window
    max_gen_len: int = 0  # 0 = auto-detect from trace mode

    # Data
    trace_mode: str = "power"  # "power" or "verbose"
    num_train: int = 10000
    num_val: int = 500
    num_test: int = 500
    exception_prob: float = 0.0  # probability of exception/trap examples
    varied_vocab: bool = False  # use varied phrasing templates
    math_ratio: float = 0.0  # fraction of examples that are arithmetic

    # System
    device: str = "cuda" if torch.cuda.is_available() else "cpu"
    compile_model: bool = torch.cuda.is_available()
    seed: int = 42


cfg = Config()


# ---------------------------------------------------------------------------
# Model
# ---------------------------------------------------------------------------

class CausalSelfAttention(nn.Module):
    def __init__(self, config):
        super().__init__()
        assert config.n_embd % config.n_head == 0
        self.n_head = config.n_head
        self.head_dim = config.n_embd // config.n_head
        self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd, bias=False)
        self.c_proj = nn.Linear(config.n_embd, config.n_embd, bias=False)
        self.dropout = config.dropout

    def forward(self, x):
        B, T, C = x.size()
        qkv = self.c_attn(x)
        q, k, v = qkv.split(C, dim=2)
        q = q.view(B, T, self.n_head, self.head_dim).transpose(1, 2)
        k = k.view(B, T, self.n_head, self.head_dim).transpose(1, 2)
        v = v.view(B, T, self.n_head, self.head_dim).transpose(1, 2)
        # Flash attention via PyTorch scaled_dot_product_attention
        y = F.scaled_dot_product_attention(q, k, v, is_causal=True,
                                           dropout_p=self.dropout if self.training else 0.0)
        y = y.transpose(1, 2).contiguous().view(B, T, C)
        return self.c_proj(y)


class MLP(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.c_fc = nn.Linear(config.n_embd, 4 * config.n_embd, bias=False)
        self.c_proj = nn.Linear(4 * config.n_embd, config.n_embd, bias=False)

    def forward(self, x):
        return self.c_proj(F.gelu(self.c_fc(x)))


class Block(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.ln_1 = nn.LayerNorm(config.n_embd)
        self.attn = CausalSelfAttention(config)
        self.ln_2 = nn.LayerNorm(config.n_embd)
        self.mlp = MLP(config)

    def forward(self, x):
        x = x + self.attn(self.ln_1(x))
        x = x + self.mlp(self.ln_2(x))
        return x


class TinyThinker(nn.Module):
    def __init__(self, config, vocab_size):
        super().__init__()
        self.config = config
        self.tok_emb = nn.Embedding(vocab_size, config.n_embd)
        self.pos_emb = nn.Embedding(config.max_seq_len, config.n_embd)
        self.blocks = nn.ModuleList([Block(config) for _ in range(config.n_layer)])
        self.ln_f = nn.LayerNorm(config.n_embd)
        # Weight tying: output projection shares weights with token embedding
        self.lm_head = nn.Linear(config.n_embd, vocab_size, bias=False)
        self.lm_head.weight = self.tok_emb.weight

        self.apply(self._init_weights)
        # Scale residual projections per GPT-2 paper
        for pn, p in self.named_parameters():
            if pn.endswith('c_proj.weight'):
                torch.nn.init.normal_(p, mean=0.0,
                                      std=0.02 / math.sqrt(2 * config.n_layer))

    def _init_weights(self, module):
        if isinstance(module, nn.Linear):
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
            if module.bias is not None:
                torch.nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)

    def forward(self, idx):
        B, T = idx.size()
        pos = torch.arange(0, T, dtype=torch.long, device=idx.device)
        x = self.tok_emb(idx) + self.pos_emb(pos)
        for block in self.blocks:
            x = block(x)
        x = self.ln_f(x)
        logits = self.lm_head(x)
        return logits

    def count_parameters(self):
        return sum(p.numel() for p in self.parameters() if p.requires_grad)


# ---------------------------------------------------------------------------
# Training
# ---------------------------------------------------------------------------

def get_lr(step, warmup_steps, max_steps, max_lr, min_lr=1e-5):
    """Cosine learning rate schedule with warmup."""
    if step < warmup_steps:
        return max_lr * (step + 1) / warmup_steps
    if step >= max_steps:
        return min_lr
    progress = (step - warmup_steps) / (max_steps - warmup_steps)
    return min_lr + 0.5 * (max_lr - min_lr) * (1 + math.cos(math.pi * progress))


def compute_loss(model, batch, pad_id):
    """Autoregressive loss masked to answer portion only.
    The model sees the full sequence but only trains on tokens after 'A:'.
    This focuses gradient signal on learning to reason, not memorize premises."""
    input_ids = batch["input_ids"]
    loss_mask = batch["loss_mask"]
    # Shift: predict next token from current
    x = input_ids[:, :-1]
    y = input_ids[:, 1:]
    mask = loss_mask[:, 1:]  # align mask with targets
    logits = model(x)
    # Compute per-token loss
    per_token_loss = F.cross_entropy(
        logits.reshape(-1, logits.size(-1)), y.reshape(-1), reduction="none"
    )
    per_token_loss = per_token_loss.view(y.shape)
    # Apply mask: only count loss on answer tokens
    masked_loss = (per_token_loss * mask.float()).sum()
    num_tokens = mask.sum().clamp(min=1)
    return masked_loss / num_tokens


def train():
    torch.manual_seed(cfg.seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed(cfg.seed)

    # Prepare data
    tokenizer, train_dataset, val_dataset, test_examples = prepare_data(
        trace_mode=cfg.trace_mode,
        num_train=cfg.num_train,
        num_val=cfg.num_val,
        num_test=cfg.num_test,
        max_seq_len=cfg.max_seq_len,
        exception_prob=cfg.exception_prob,
        varied_vocab=cfg.varied_vocab,
        math_ratio=cfg.math_ratio,
    )

    train_loader = DataLoader(
        train_dataset, batch_size=cfg.batch_size, shuffle=True,
        collate_fn=collate_fn, num_workers=2, pin_memory=True,
    )

    # Build model
    model = TinyThinker(cfg, tokenizer.vocab_size).to(cfg.device)
    n_params = model.count_parameters()
    print(f"\nModel: {n_params:,} parameters")
    print(f"Config: {cfg.n_layer}L {cfg.n_head}H {cfg.n_embd}D")
    print(f"Vocab: {tokenizer.vocab_size} tokens")

    if cfg.compile_model:
        model = torch.compile(model)

    # Optimizer
    optimizer = torch.optim.AdamW(
        model.parameters(), lr=cfg.learning_rate,
        weight_decay=cfg.weight_decay, betas=(0.9, 0.95),
    )

    # Training loop
    step = 0
    best_accuracy = 0.0
    best_step = 0
    eval_history = []  # rolling window for convergence detection
    data_iter = iter(train_loader)
    t0 = time.time()

    print(f"\nTraining for {cfg.max_steps} steps...")
    print(f"Device: {cfg.device}")
    print("-" * 60)

    while step < cfg.max_steps:
        model.train()

        # Get batch (cycle through data)
        try:
            batch = next(data_iter)
        except StopIteration:
            data_iter = iter(train_loader)
            batch = next(data_iter)

        batch = {k: v.to(cfg.device) for k, v in batch.items()}

        # Forward + backward
        lr = get_lr(step, cfg.warmup_steps, cfg.max_steps, cfg.learning_rate)
        for param_group in optimizer.param_groups:
            param_group['lr'] = lr

        loss = compute_loss(model, batch, tokenizer.pad_id)
        loss.backward()

        if cfg.grad_clip > 0:
            torch.nn.utils.clip_grad_norm_(model.parameters(), cfg.grad_clip)

        optimizer.step()
        optimizer.zero_grad(set_to_none=True)

        # Logging
        if step % 50 == 0:
            elapsed = time.time() - t0
            print(f"step {step:5d} | loss {loss.item():.4f} | "
                  f"lr {lr:.2e} | {elapsed:.1f}s")

        # Evaluation
        if step > 0 and step % cfg.eval_interval == 0:
            eval_examples = test_examples[:cfg.eval_samples]
            raw_model = model._orig_mod if hasattr(model, '_orig_mod') else model
            gen_len = cfg.max_gen_len if cfg.max_gen_len > 0 else 100
            results = evaluate_model(raw_model, tokenizer, eval_examples,
                                     cfg.device, max_gen_len=gen_len)
            accuracy = results["accuracy"]
            print(f"\n  EVAL step {step}: accuracy={accuracy:.3f} "
                  f"({results['correct']}/{results['total']}) "
                  f"parse_fail={results['parse_failures']}")

            # Track best single eval for checkpoint saving
            if accuracy > best_accuracy:
                best_accuracy = accuracy
                best_step = step
                save_dir = Path(__file__).parent / "checkpoints"
                save_dir.mkdir(exist_ok=True)
                torch.save(raw_model.state_dict(), save_dir / f"best_{cfg.trace_mode}.pt")
                with open(save_dir / f"config_{cfg.trace_mode}.json", "w") as f:
                    json.dump({
                        "n_layer": cfg.n_layer, "n_head": cfg.n_head,
                        "n_embd": cfg.n_embd, "max_seq_len": cfg.max_seq_len,
                        "vocab_size": tokenizer.vocab_size,
                        "trace_mode": cfg.trace_mode,
                    }, f, indent=2)
                print(f"  New best accuracy: {best_accuracy:.3f}")

            # Rolling average convergence detection
            eval_history.append(accuracy)
            w = cfg.early_stop_window
            if len(eval_history) >= w:
                current_avg = sum(eval_history[-w:]) / w
                print(f"  Rolling avg ({w}): {current_avg:.3f}")

            print("-" * 60)

            # Early stopping: detect plateau by measuring improvement rate
            # Compare current rolling avg to the rolling avg from w evals ago.
            # If the improvement is less than threshold, we've captured 95%+ of gains.
            # Minimum eval floor prevents premature stopping during early oscillations.
            if (cfg.early_stop_threshold > 0
                    and len(eval_history) >= cfg.early_stop_min_evals
                    and len(eval_history) >= w * 2):
                current_avg = sum(eval_history[-w:]) / w
                prior_avg = sum(eval_history[-w*2:-w]) / w
                improvement = current_avg - prior_avg

                if improvement < cfg.early_stop_threshold:
                    print(f"\n  Converged: rolling avg {current_avg:.3f}, "
                          f"improvement over last window: {improvement:+.4f} "
                          f"(threshold: {cfg.early_stop_threshold})")
                    print(f"  Best accuracy: {best_accuracy:.3f} at step {best_step}")
                    break

        step += 1

    # Final evaluation
    elapsed = time.time() - t0
    print(f"\nTraining complete in {elapsed:.1f}s ({elapsed/60:.1f}m)")

    raw_model = model._orig_mod if hasattr(model, '_orig_mod') else model
    gen_len = 200 if cfg.trace_mode == "verbose" else 100
    results = evaluate_model(raw_model, tokenizer, test_examples,
                             cfg.device, max_gen_len=gen_len)
    print(f"\nFinal test accuracy: {results['accuracy']:.3f} "
          f"({results['correct']}/{results['total']})")
    print(f"Parse failures: {results['parse_failures']}")
    print(f"Best accuracy during training: {best_accuracy:.3f}")

    # Save final model
    save_dir = Path(__file__).parent / "checkpoints"
    save_dir.mkdir(exist_ok=True)
    torch.save(raw_model.state_dict(), save_dir / f"final_{cfg.trace_mode}.pt")

    # Print a few generated examples for inspection
    print("\n--- Generated examples ---")
    model.eval()
    for i in range(min(5, len(test_examples))):
        ex = test_examples[i]
        parts = ex["text"].split(" A: ")
        if len(parts) != 2:
            continue
        prompt = parts[0] + " A:"
        prompt_ids = tokenizer.encode(prompt)
        input_ids = torch.tensor([prompt_ids], dtype=torch.long, device=cfg.device)

        from prepare import generate
        ex_gen_len = cfg.max_gen_len if cfg.max_gen_len > 0 else 100
        gen_ids = generate(raw_model, input_ids, ex_gen_len, tokenizer.eos_id, tokenizer.pad_id)
        gen_text = tokenizer.decode(gen_ids[0].tolist())

        print(f"\n[{i}] Expected answer: {'yes' if ex['answer'] else 'no'}")
        print(f"    Generated: {gen_text}")

    # Print the key metric for autoresearch
    print(f"\n>>> val_accuracy: {results['accuracy']:.4f}")

    return results["accuracy"]


if __name__ == "__main__":
    train()