train.py

"""Training script for TRM (Transformer Reasoning Model)."""

import math
import os
from dataclasses import dataclass
from pathlib import Path

import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from tqdm import tqdm
import trackio as wandb

from puzzle_dataset import PuzzleDataset, PuzzleDatasetConfig
from model import TRM
from losses import compute_loss


@dataclass
class PretrainConfig:
    """Configuration for pretraining."""
    seed: int
    data_path: str
    global_batch_size: int


def create_dataloader(config: PretrainConfig, split: str, **kwargs):
    """Create a dataloader for training or evaluation.

    Args:
        config: Pretraining configuration
        split: Dataset split ("train" or "test")
        **kwargs: Additional arguments passed to PuzzleDatasetConfig
            - test_set_mode: bool (required)
            - epochs_per_iter: int (required)
            - global_batch_size: int (optional, defaults to config.global_batch_size)

    Returns:
        tuple: (dataloader, metadata)
    """
    # Fixed to single GPU training (no distributed)
    rank = 0
    world_size = 1

    # Create dataset with provided configuration
    dataset = PuzzleDataset(
        PuzzleDatasetConfig(
            seed=config.seed,
            dataset_path=config.data_path,
            rank=rank,
            num_replicas=world_size,
            **kwargs
        ),
        split=split
    )

    # Create dataloader with specified settings
    dataloader = DataLoader(
        dataset,
        batch_size=None,  # Dataset handles batching internally
        num_workers=1,
        prefetch_factor=8,
        pin_memory=True,
        persistent_workers=True
    )

    return dataloader, dataset.metadata


class EMA:
    """Exponential Moving Average of model parameters.

    Maintains shadow copies of model parameters that are updated using
    exponential moving average. This helps prevent training collapse
    as described in Section 4.7 of the paper.

    Hardcoded parameters:
    - decay = 0.999
    """

    def __init__(self, model: nn.Module):
        """Initialize EMA with model parameters.

        Args:
            model: The model whose parameters to track
        """
        self.decay = 0.999
        self.shadow = {}
        self.backup = {}

        # Clone all model parameters as shadow parameters
        for name, param in model.named_parameters():
            if param.requires_grad:
                self.shadow[name] = param.data.clone()

    def update(self, model: nn.Module):
        """Update shadow parameters with current model parameters.

        Args:
            model: The model with updated parameters
        """
        for name, param in model.named_parameters():
            if param.requires_grad:
                assert name in self.shadow
                new_average = (
                    self.decay * self.shadow[name].data +
                    (1.0 - self.decay) * param.data
                )
                self.shadow[name].data = new_average

    def apply_shadow(self, model: nn.Module):
        """Replace model parameters with shadow parameters.

        Args:
            model: The model to apply shadow parameters to
        """
        for name, param in model.named_parameters():
            if param.requires_grad:
                assert name in self.shadow
                self.backup[name] = param.data.clone()
                param.data = self.shadow[name]

    def restore(self, model: nn.Module):
        """Restore original model parameters.

        Args:
            model: The model to restore parameters to
        """
        for name, param in model.named_parameters():
            if param.requires_grad:
                assert name in self.backup
                param.data = self.backup[name]
        self.backup = {}


def setup_device() -> torch.device:
    """Select and return the best available device.

    Returns:
        torch.device: Selected device (MPS, CUDA, or CPU)
    """
    if torch.backends.mps.is_available():
        device = torch.device('mps')
    elif torch.cuda.is_available():
        device = torch.device('cuda')
    else:
        device = torch.device('cpu')

    print(f"Using device: {device}")
    return device


def create_training_components(device: torch.device):
    """Initialize model, optimizer, scheduler, and EMA.

    Args:
        device: Device to place model on

    Returns:
        tuple: (model, optimizer, scheduler, ema)
    """
    # Hardcoded hyperparameters from Page 11
    warmup_iterations = 2000

    # Initialize model
    print("Initializing TRM model...")
    model = TRM().to(device)

    # Count parameters
    total_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
    print(f"Model parameters: {total_params:,} ({total_params/1e6:.1f}M)")

    # Initialize optimizer (AdamW)
    optimizer = torch.optim.AdamW(
        model.parameters(),
        lr=1e-4,
        betas=(0.9, 0.95),
        weight_decay=1.0
    )

    # Learning rate warmup scheduler
    def lr_lambda(step):
        """Linear warmup for first 2000 steps, then constant."""
        if step < warmup_iterations:
            return step / warmup_iterations
        return 1.0

    scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)

    # Initialize EMA
    print("Initializing EMA...")
    ema = EMA(model)

    return model, optimizer, scheduler, ema



def save_checkpoint(
    checkpoint_dir: Path,
    epoch: int,
    global_step: int,
    model: nn.Module,
    optimizer: torch.optim.Optimizer,
    scheduler: torch.optim.lr_scheduler.LambdaLR,
    ema: EMA,
    loss: float
):
    """Save training checkpoint.

    Args:
        checkpoint_dir: Directory to save checkpoint
        epoch: Current epoch number
        global_step: Current global step
        model: TRM model
        optimizer: Optimizer
        scheduler: Learning rate scheduler
        ema: EMA tracker
        loss: Current loss value
    """
    checkpoint_path = checkpoint_dir / f"checkpoint_epoch_{epoch}.pt"
    torch.save({
        'epoch': epoch,
        'global_step': global_step,
        'model_state_dict': model.state_dict(),
        'optimizer_state_dict': optimizer.state_dict(),
        'scheduler_state_dict': scheduler.state_dict(),
        'ema_shadow': ema.shadow,
        'loss': loss,
    }, checkpoint_path)
    print(f"Saved checkpoint: {checkpoint_path}")


def save_ema_checkpoint(
    checkpoint_dir: Path,
    epoch: int,
    global_step: int,
    model: nn.Module,
    ema: EMA,
    loss: float
):
    """Save EMA model checkpoint.

    Args:
        checkpoint_dir: Directory to save checkpoint
        epoch: Current epoch number
        global_step: Current global step
        model: TRM model
        ema: EMA tracker
        loss: Current loss value
    """
    ema_checkpoint_path = checkpoint_dir / f"checkpoint_ema_epoch_{epoch}.pt"

    # Temporarily apply EMA weights
    ema.apply_shadow(model)
    torch.save({
        'epoch': epoch,
        'global_step': global_step,
        'model_state_dict': model.state_dict(),
        'loss': loss,
    }, ema_checkpoint_path)
    ema.restore(model)

    print(f"Saved EMA checkpoint: {ema_checkpoint_path}")


def train_one_batch(
    model: nn.Module,
    batch: dict,
    optimizer: torch.optim.Optimizer,
    scheduler: torch.optim.lr_scheduler.LambdaLR,
    device: torch.device,
    global_batch_size: int,
    max_supervision_steps: int
) -> float:
    """Train on a single batch with deep supervision.

    Implements the deep supervision loop from Algorithm 3 (Figure 3, Page 5).
    Accumulates gradients across multiple supervision steps, then applies
    single optimizer update per batch.

    Args:
        model: TRM model
        batch: Dictionary with 'inputs' and 'labels' keys
        optimizer: AdamW optimizer
        scheduler: Learning rate scheduler
        device: Device for computation
        global_batch_size: Batch size for loss scaling
        max_supervision_steps: Maximum supervision steps (16)

    Returns:
        float: Total batch loss (sum across all supervision steps)
    """
    # Extract input and target from batch
    x_input = batch['inputs'].to(device)  # [B, 81]
    y_true = batch['labels'].to(device)  # [B, 81]

    # Initialize states
    batch_loss = 0.0
    batch_pred_loss = 0.0
    batch_halt_loss = 0.0
    y, z = None, None  # Start with fresh initialization

    # Clear gradients before accumulation
    optimizer.zero_grad()

    for step in range(max_supervision_steps):
        # Forward pass: embed input, update states, get predictions
        (y, z), y_hat, q_hat = model(x_input, y, z)

        # --- NaN diagnostics: model outputs ---
        if not torch.isfinite(y_hat).all():
            n_nan = (~torch.isfinite(y_hat)).sum().item()
            raise RuntimeError(
                f"Non-finite values in y_hat at supervision step {step}: "
                f"{n_nan}/{y_hat.numel()} elements. "
                f"min={y_hat.min().item()}, max={y_hat.max().item()}"
            )
        if not torch.isfinite(q_hat).all():
            raise RuntimeError(
                f"Non-finite values in q_hat at supervision step {step}: q_hat={q_hat}"
            )

        # Compute loss (returns tuple of total_loss, pred_loss, halt_loss)
        loss, pred_loss, halt_loss = compute_loss(y_hat, y_true, q_hat)

        # --- NaN diagnostics: loss components ---
        if not torch.isfinite(loss):
            raise RuntimeError(
                f"Non-finite loss at supervision step {step}: "
                f"total={loss.item()}, pred={pred_loss.item()}, halt={halt_loss.item()}. "
                f"y_hat stats: min={y_hat.min().item():.4f}, max={y_hat.max().item():.4f}, "
                f"mean={y_hat.mean().item():.4f}, std={y_hat.std().item():.4f}"
            )

        # Scale and accumulate gradients
        # Divide by batch_size AND max_supervision_steps to average gradients
        # across both batch samples and supervision steps
        (loss / global_batch_size).backward()

        # Accumulate losses for logging
        batch_loss += loss.item()
        batch_pred_loss += pred_loss.item()
        batch_halt_loss += halt_loss.item()

    # Note: No early stopping in simplified version
    # Reference uses per-sample halted flags and breaks when all samples halted
    # For now, we run all max_supervision_steps for consistency

    # Apply accumulated gradients (single optimizer step per batch)
    optimizer.step()

    # Update learning rate schedule (once per batch)
    scheduler.step()

    # Compute final metrics for logging
    # y_hat and q_hat are the final predictions from last supervision step
    with torch.no_grad():
        # Puzzle accuracy: fraction of correct token predictions
        pred_tokens = y_hat.argmax(dim=-1)  # [B, L]
        accuracy = (pred_tokens == y_true).float().mean()

        # Halting statistics
        halting_mean = q_hat.mean()
        halting_std = q_hat.std()

    # Build metrics dict and filter non-finite values so trackio's JSON
    # encoder doesn't crash on NaN/Inf (e.g. q_hat.std() with B=1).
    raw_metrics = {
        'train/batch_loss': batch_loss,
        'train/prediction_loss': batch_pred_loss,
        'train/halting_loss': batch_halt_loss,
        'train/puzzle_accuracy': accuracy.item(),
        'train/halting_mean': halting_mean.item(),
        'train/halting_std': halting_std.item(),
    }
    bad = {k: v for k, v in raw_metrics.items() if not math.isfinite(v)}
    if bad:
        # Loud diagnostic, but don't crash logging.
        print(f"[WARN] Dropping non-finite metrics from log: {bad}")
    safe_metrics = {k: v for k, v in raw_metrics.items() if math.isfinite(v)}
    wandb.log(safe_metrics)

    return batch_loss


def train(
    max_epochs=60000,
    batch_size=256,  # Reduced from 768 (paper spec) for MPS memory constraints
    data_path="./data/sudoku-extreme-1k-aug-1000"
):
    """Main training function implementing Algorithm 3 from the paper.

    All hyperparameters are hardcoded as specified in Page 11:
    - learning_rate = 1e-4
    - weight_decay = 1.0
    - beta1 = 0.9
    - beta2 = 0.95
    - warmup_iterations = 2000
    - max_epochs = 60000
    - max_supervision_steps = 16
    - batch_size = 768 (handled by dataset)

    Training follows Algorithm 3 (Figure 3, Page 5):
    1. Initialize model, optimizer, EMA
    2. For each batch:
        - Deep supervision loop (up to 16 steps)
        - Single optimizer update per batch
    3. Periodic evaluation and checkpointing

    Args:
        max_epochs: Maximum number of training epochs (default: 60000)
        batch_size: Batch size (default: 64, paper spec is 768)
        data_path: Path to training data
    """
    # Hardcoded hyperparameters from Page 11
    max_supervision_steps = 16
    learning_rate = 1e-4
    weight_decay = 1.0
    warmup_iterations = 2000
    ema_decay = 0.999

    # Setup
    device = setup_device()
    model, optimizer, scheduler, ema = create_training_components(device)

    # Get total model parameters for config
    total_params = sum(p.numel() for p in model.parameters() if p.requires_grad)

    # Trackio configuration
    trackio_config = {
        # Optimization hyperparameters
        "learning_rate": learning_rate,
        "weight_decay": weight_decay,
        "beta1": 0.9,
        "beta2": 0.95,
        "warmup_iterations": warmup_iterations,
        # Training hyperparameters
        "max_epochs": max_epochs,
        "batch_size": batch_size,
        "max_supervision_steps": max_supervision_steps,
        "ema_decay": ema_decay,
        # Architecture parameters
        "hidden_size": 512,
        "num_layers": 2,
        "context_length": 81,
        "vocab_size": 11,
        "total_parameters": total_params,
    }

    # Load training data
    print("Loading training data...")
    config = PretrainConfig(
        seed=42,
        data_path=data_path,
        global_batch_size=batch_size
    )
    train_loader, train_metadata = create_dataloader(
        config,
        "train",
        test_set_mode=False,
        epochs_per_iter=1,
        global_batch_size=batch_size
    )
    print(f"Train metadata: {train_metadata}")

    # Create checkpoint directory
    checkpoint_dir = Path("checkpoints")
    checkpoint_dir.mkdir(exist_ok=True)

    # Calculate estimated batches per epoch for progress tracking
    est_batches_per_epoch = train_metadata.total_groups // batch_size

    # Training loop with trackio
    print(f"\nStarting training for {max_epochs} epochs...")

    # Initialize trackio
    wandb.init(project="trm-sudoku", config=trackio_config)

    global_step = 0

    for epoch in range(max_epochs):
        model.train()
        epoch_loss = 0.0
        num_batches = 0

        pbar = tqdm(train_loader, desc=f"Epoch {epoch + 1}/{max_epochs}")
        for __, batch, global_batch_size in pbar:
            # Train on single batch with deep supervision
            batch_loss = train_one_batch(
                model, batch, optimizer, scheduler, device, global_batch_size, max_supervision_steps
            )

            # Update EMA (once per batch)
            ema.update(model)

            # Track statistics
            epoch_loss += batch_loss
            num_batches += 1
            global_step += 1

            # Log step-level metrics to trackio
            wandb.log({
                'train/learning_rate': scheduler.get_last_lr()[0],
                'train/global_step': global_step,
                'train/epoch': epoch,
            })

            # Update progress bar
            pbar.set_postfix({
                'batch': f'{num_batches}/{est_batches_per_epoch}',
                'loss': f'{batch_loss:.4f}',
                'lr': f'{scheduler.get_last_lr()[0]}'
            })

        # Epoch statistics
        avg_epoch_loss = epoch_loss / num_batches if num_batches > 0 else 0.0

        # Periodic logging
        if (epoch + 1) % 100 == 0:
            print(f"Epoch {epoch+1}/{max_epochs} - Loss: {avg_epoch_loss:.4f} - Step: {global_step}")

        # Periodic checkpointing and evaluation
        if (epoch + 1) % 1000 == 0:
            print(f"\nSaving checkpoint at epoch {epoch+1}...")
            save_checkpoint(checkpoint_dir, epoch + 1, global_step, model, optimizer, scheduler, ema, avg_epoch_loss)
            save_ema_checkpoint(checkpoint_dir, epoch + 1, global_step, model, ema, avg_epoch_loss)

    print("\nTraining complete!")

    # Save final checkpoints
    print("Saving final checkpoints...")
    save_checkpoint(checkpoint_dir, max_epochs, global_step, model, optimizer, scheduler, ema, 0.0)
    save_ema_checkpoint(checkpoint_dir, max_epochs, global_step, model, ema, 0.0)
    print(f"Final checkpoints saved to {checkpoint_dir}")

    # Finish trackio run
    wandb.finish()


if __name__ == "__main__":
    train()