train_and_benchmark.py

#!/usr/bin/env python3
"""
train_and_benchmark.py — Train Baseline vs Tiered HRM, then benchmark both.

Usage:
    source venv/bin/activate

    # Quick smoke test (tiny dataset, 2 epochs)
    python train_and_benchmark.py --data-path data/sudoku-1k --epochs 2 --batch-size 384

    # Full run
    python train_and_benchmark.py --data-path data/sudoku-1k --epochs 1000 --batch-size 384

    # Benchmark only (skip training)
    python train_and_benchmark.py --benchmark-only

Outputs:
    - Console: live training metrics + benchmark table
    - benchmark_results/results.json
    - benchmark_results/benchmark_comparison.png
    - benchmark_results/memory_analysis.png
"""

import argparse
import json
import os
import sys
import time

import torch

# ── Project imports ──────────────────────────────────────────────────
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from omegaconf import OmegaConf
from pretrain import PretrainConfig, init_train_state, train_batch, evaluate, create_dataloader, save_train_state
from benchmark import (
    benchmark_tiered_model,
    benchmark_baseline_model,
    compare_models,
    print_results_table,
    generate_plots,
)


# =====================================================================
#  PART 1: TRAINING
# =====================================================================

def run_training(arch_name: str, args, device):
    """Train one architecture variant and return final metrics."""
    print(f"\n{'='*60}")
    print(f"  Training: {arch_name}")
    print(f"  Data:     {args.data_path}")
    print(f"  Epochs:   {args.epochs}  |  Batch: {args.batch_size}")
    print(f"{'='*60}\n")

    # Build Hydra-style config dict
    cfg = OmegaConf.create({
        "arch": arch_name,
        "data_path": args.data_path,
        "global_batch_size": args.batch_size,
        "epochs": args.epochs,
        "lr": 7e-5,
        "lr_min_ratio": 0.0,
        "lr_warmup_steps": 100,
        "weight_decay": 1.0,
        "beta1": 0.9,
        "beta2": 0.95,
        "puzzle_emb_lr": 7e-5,
        "puzzle_emb_weight_decay": 1.0,
        "seed": 0,
        "skip_eval": False,
        "eval_interval": max(1, args.epochs // 5),  # eval 5 times
    })
    # Merge with the arch yaml
    from hydra import compose, initialize_config_dir
    from hydra.core.global_hydra import GlobalHydra
    GlobalHydra.instance().clear()
    with initialize_config_dir(config_dir=os.path.join(os.path.abspath("."), "config"), version_base=None):
        hydra_cfg = compose(config_name="cfg_pretrain", overrides=[
            f"arch={arch_name}",
            f"data_path={args.data_path}",
            f"global_batch_size={args.batch_size}",
            f"epochs={args.epochs}",
        ])

    config = PretrainConfig(**OmegaConf.to_container(hydra_cfg, resolve=True))
    config.eval_interval = max(1, args.epochs // 5)

    # Dataloaders
    train_loader, train_meta = create_dataloader(
        config, "train", 0, 1,
        test_set_mode=False, epochs_per_iter=1,
        global_batch_size=config.global_batch_size,
    )
    eval_loader, eval_meta = create_dataloader(
        config, "test", 0, 1,
        test_set_mode=True, epochs_per_iter=1,
        global_batch_size=config.global_batch_size,
    )

    # Model
    train_state = init_train_state(config, train_meta, world_size=1)
    param_count = sum(p.numel() for p in train_state.model.parameters())
    print(f"  Parameters: {param_count:,} ({param_count/1e6:.1f}M)")

    # Training loop
    best_acc = 0.0
    epoch_times = []

    for epoch in range(1, config.epochs + 1):
        train_state.model.train()
        t0 = time.perf_counter()
        last_metrics = None

        for set_name, batch, gbs in train_loader:
            metrics = train_batch(config, train_state, batch, gbs, rank=0, world_size=1)
            if metrics:
                last_metrics = metrics

        dt = time.perf_counter() - t0
        epoch_times.append(dt)

        # Log
        if last_metrics and epoch % max(1, args.epochs // 20) == 0:
            loss = last_metrics.get("train/total_loss", 0)
            acc  = last_metrics.get("train/exact_accuracy", 0)
            print(f"  Epoch {epoch:>5}/{config.epochs} | Loss: {loss:.4f} | Acc: {acc:.2%} | {dt:.1f}s")

        # Eval
        if config.eval_interval and epoch % config.eval_interval == 0:
            train_state.model.eval()
            eval_results = evaluate(config, train_state, eval_loader, eval_meta, rank=0, world_size=1)
            if eval_results:
                for s_name, s_metrics in eval_results.items():
                    ea = s_metrics.get("exact_accuracy", 0)
                    print(f"    eval/{s_name}: exact_accuracy={ea:.2%}")
                    best_acc = max(best_acc, ea)

    # Summary
    avg_epoch = sum(epoch_times) / len(epoch_times) if epoch_times else 0
    print(f"\n  ── {arch_name} Done ──")
    print(f"  Best Eval Accuracy: {best_acc:.2%}")
    print(f"  Avg Epoch Time:     {avg_epoch:.2f}s")
    print(f"  GPU Peak Memory:    {torch.cuda.max_memory_allocated()/1e9:.2f} GB")

    return {"arch": arch_name, "best_acc": best_acc, "avg_epoch_s": avg_epoch}


# =====================================================================
#  PART 2: HARDWARE BENCHMARK
# =====================================================================

def run_benchmark(args):
    """Compare Tiered vs Baseline on synthetic data."""
    print(f"\n{'='*60}")
    print(f"  Hardware Benchmark: Tiered vs Baseline")
    print(f"  Batch sizes: {args.bench_batch_sizes}")
    print(f"  Seq lengths:  {args.bench_seq_lens}")
    print(f"{'='*60}")

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    bs_list = [int(x) for x in args.bench_batch_sizes.split(",")]
    sl_list = [int(x) for x in args.bench_seq_lens.split(",")]

    results = compare_models(
        batch_sizes=bs_list, seq_lens=sl_list,
        hidden_size=512, warmup=5, iterations=args.bench_iters,
        device=device,
    )

    # Save
    os.makedirs(args.output_dir, exist_ok=True)
    out_path = os.path.join(args.output_dir, "results.json")
    with open(out_path, "w") as f:
        json.dump(results, f, indent=2, default=str)
    print(f"\n  Results saved → {out_path}")

    # Plots
    try:
        generate_plots(results, output_dir=args.output_dir)
    except Exception as e:
        print(f"  (Plots skipped: {e})")

    return results


# =====================================================================
#  MAIN
# =====================================================================

def main():
    parser = argparse.ArgumentParser(description="HRM: Train + Benchmark")
    # Training
    parser.add_argument("--data-path", type=str, default="data/sudoku-1k")
    parser.add_argument("--epochs", type=int, default=100)
    parser.add_argument("--batch-size", type=int, default=384)
    # Benchmark
    parser.add_argument("--bench-batch-sizes", type=str, default="1,8,32")
    parser.add_argument("--bench-seq-lens", type=str, default="64,128")
    parser.add_argument("--bench-iters", type=int, default=20)
    parser.add_argument("--output-dir", type=str, default="benchmark_results")
    # Flow control
    parser.add_argument("--benchmark-only", action="store_true", help="Skip training, only run benchmark")
    parser.add_argument("--train-only", action="store_true", help="Skip benchmark, only run training")
    args = parser.parse_args()

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"\n  Device: {device}")
    if torch.cuda.is_available():
        print(f"  GPU:    {torch.cuda.get_device_name(0)}")

    # ── Training ──
    train_summary = {}
    if not args.benchmark_only:
        for arch in ["hrm_v1", "hrm_tiered"]:
            try:
                result = run_training(arch, args, device)
                train_summary[arch] = result
            except Exception as e:
                print(f"\n  ⚠ Training failed for {arch}: {e}")
                train_summary[arch] = {"error": str(e)}
            torch.cuda.empty_cache()

        # Print training comparison
        print(f"\n{'='*60}")
        print(f"  Training Comparison")
        print(f"{'='*60}")
        print(f"  {'Arch':<15} {'Best Acc':>10} {'Avg Epoch':>12}")
        print(f"  {'-'*37}")
        for arch, r in train_summary.items():
            if "error" in r:
                print(f"  {arch:<15} {'FAILED':>10} {'':>12}")
            else:
                print(f"  {arch:<15} {r['best_acc']:>10.2%} {r['avg_epoch_s']:>10.2f}s")

    # ── Benchmark ──
    if not args.train_only:
        run_benchmark(args)

    print(f"\n  All done!\n")


if __name__ == "__main__":
    main()