train_full.py

"""
Unattended training script for ailmo.

Runs training to completion with:
- Auto-resume from latest checkpoint
- Graceful signal handling (SIGINT/SIGTERM → save and exit)
- Rolling checkpoint window (keeps last N)
- JSONL logging to results/
- Periodic validation with results saved

Usage:
    python train_full.py                                    # defaults
    python train_full.py --max-steps 50000 --lr 3e-4        # custom
    python train_full.py --data-dir data/fineweb/           # FineWeb-Edu data
    # If interrupted, re-run the same command — it auto-resumes
"""

import argparse
import json
import math
import os
import signal
import sys
import time
import torch
from torch.utils.data import DataLoader

from configs import ModelConfig, TrainConfig
from model import LLM
from data import TokenDataset, download_tiny_shakespeare, prepare_dataset


# ---------------------------------------------------------------------------
# Graceful shutdown
# ---------------------------------------------------------------------------
_shutdown_requested = False


def _signal_handler(signum, frame):
    global _shutdown_requested
    _shutdown_requested = True
    sig_name = signal.Signals(signum).name
    print(f"\n[!] Received {sig_name} — will save checkpoint and exit after this step...")


# ---------------------------------------------------------------------------
# Checkpoint management
# ---------------------------------------------------------------------------

def find_latest_checkpoint(ckpt_dir: str) -> str | None:
    """Find the highest-step checkpoint in the directory."""
    if not os.path.exists(ckpt_dir):
        return None
    files = [f for f in os.listdir(ckpt_dir) if f.startswith("step_") and f.endswith(".pt")]
    if not files:
        return None
    # Sort by step number
    files.sort(key=lambda f: int(f.replace("step_", "").replace(".pt", "")))
    return os.path.join(ckpt_dir, files[-1])


def save_checkpoint(path: str, step: int, model, optimizer, model_config, train_config, metrics: dict):
    """Save a full checkpoint with all state needed for resume."""
    os.makedirs(os.path.dirname(path), exist_ok=True)
    torch.save({
        "step": step,
        "model_state_dict": model.state_dict(),
        "optimizer_state_dict": optimizer.state_dict(),
        "model_config": model_config,
        "train_config": train_config,
        "metrics": metrics,
    }, path)


def cleanup_old_checkpoints(ckpt_dir: str, keep_n: int):
    """Delete oldest checkpoints, keeping the last N plus any 'final.pt'."""
    if not os.path.exists(ckpt_dir):
        return
    files = [f for f in os.listdir(ckpt_dir) if f.startswith("step_") and f.endswith(".pt")]
    files.sort(key=lambda f: int(f.replace("step_", "").replace(".pt", "")))
    # Delete oldest files beyond the keep window
    to_delete = files[:-keep_n] if len(files) > keep_n else []
    for f in to_delete:
        os.remove(os.path.join(ckpt_dir, f))
        print(f"  [cleanup] Removed old checkpoint: {f}")


# ---------------------------------------------------------------------------
# Learning rate schedule
# ---------------------------------------------------------------------------

def get_lr(step: int, warmup_steps: int, max_steps: int, lr: float, min_lr: float) -> float:
    if step < warmup_steps:
        return lr * (step + 1) / warmup_steps
    progress = (step - warmup_steps) / max(1, max_steps - warmup_steps)
    cosine = 0.5 * (1 + math.cos(math.pi * progress))
    return min_lr + (lr - min_lr) * cosine


# ---------------------------------------------------------------------------
# Evaluation
# ---------------------------------------------------------------------------

@torch.no_grad()
def evaluate(model, val_loader, device, dtype, max_batches: int = 50) -> float:
    model.eval()
    total_loss = 0.0
    n = 0
    for x, y in val_loader:
        if n >= max_batches:
            break
        x, y = x.to(device), y.to(device)
        with torch.autocast(device_type="cuda", dtype=dtype):
            logits = model(x)
            loss = torch.nn.functional.cross_entropy(
                logits.view(-1, logits.size(-1)), y.view(-1)
            )
        total_loss += loss.item()
        n += 1
    model.train()
    return total_loss / max(n, 1)


# ---------------------------------------------------------------------------
# Main training function
# ---------------------------------------------------------------------------

def train(args):
    # Register signal handlers for graceful shutdown
    signal.signal(signal.SIGINT, _signal_handler)
    signal.signal(signal.SIGTERM, _signal_handler)

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    dtype = torch.bfloat16
    print(f"Device: {device} | Precision: bf16")

    # --- Config ---
    model_config = ModelConfig()
    train_config = TrainConfig(
        learning_rate=args.lr,
        max_steps=args.max_steps,
        batch_size=args.batch_size,
        gradient_accumulation_steps=args.grad_accum,
        eval_interval=args.eval_interval,
        save_interval=args.save_interval,
        log_interval=args.log_interval,
        dataset_path=os.path.join(args.data_dir, "train.bin"),
        val_dataset_path=os.path.join(args.data_dir, "val.bin"),
        checkpoint_dir=args.checkpoint_dir,
    )

    # --- Data ---
    if not os.path.exists(train_config.dataset_path):
        print("No training data found. Preparing TinyShakespeare...")
        text_path = download_tiny_shakespeare(args.data_dir)
        prepare_dataset(text_path, args.data_dir)

    train_ds = TokenDataset(train_config.dataset_path, model_config.context_length)
    val_ds = TokenDataset(train_config.val_dataset_path, model_config.context_length)

    train_loader = DataLoader(
        train_ds, batch_size=train_config.batch_size, shuffle=True,
        num_workers=4, pin_memory=False, drop_last=True,
    )
    val_loader = DataLoader(
        val_ds, batch_size=train_config.batch_size, shuffle=False,
        num_workers=2, pin_memory=False, drop_last=True,
    )

    # --- Model ---
    model = LLM(model_config).to(device)
    print(f"Model: {model.count_parameters():,} parameters")

    # --- Optimizer ---
    decay_params = [p for p in model.parameters() if p.dim() >= 2]
    no_decay_params = [p for p in model.parameters() if p.dim() < 2]
    optimizer = torch.optim.AdamW([
        {"params": decay_params, "weight_decay": train_config.weight_decay},
        {"params": no_decay_params, "weight_decay": 0.0},
    ], lr=train_config.learning_rate, betas=(train_config.beta1, train_config.beta2))

    # --- Auto-resume from checkpoint ---
    start_step = 0
    metrics_history = {"losses": [], "val_losses": []}

    latest_ckpt = find_latest_checkpoint(train_config.checkpoint_dir)
    if latest_ckpt:
        print(f"Resuming from checkpoint: {latest_ckpt}")
        ckpt = torch.load(latest_ckpt, map_location=device, weights_only=False)
        model.load_state_dict(ckpt["model_state_dict"])
        optimizer.load_state_dict(ckpt["optimizer_state_dict"])
        start_step = ckpt["step"] + 1
        if "metrics" in ckpt:
            metrics_history = ckpt["metrics"]
        print(f"Resumed at step {start_step}")
    else:
        print("Starting fresh training")

    # --- Results directory ---
    results_dir = "results"
    os.makedirs(results_dir, exist_ok=True)
    log_path = os.path.join(results_dir, "training_log.jsonl")

    # If resuming, don't overwrite existing log — append mode
    print(f"Logging to: {log_path}")

    # --- Training loop ---
    train_iter = iter(train_loader)
    tokens_processed = 0
    t0 = time.time()
    cfg = train_config

    eff_batch = cfg.batch_size * cfg.gradient_accumulation_steps
    tok_per_step = eff_batch * model_config.context_length

    print(f"\nTraining: steps {start_step} -> {cfg.max_steps}")
    print(f"  Effective batch: {eff_batch} | Tokens/step: {tok_per_step:,}")
    print(f"  Eval every {cfg.eval_interval} steps | Save every {cfg.save_interval} steps")
    print(f"  Keep last {args.keep_checkpoints} checkpoints")
    print()

    for step in range(start_step, cfg.max_steps):
        # Check for graceful shutdown
        if _shutdown_requested:
            ckpt_path = os.path.join(cfg.checkpoint_dir, f"step_{step}.pt")
            save_checkpoint(ckpt_path, step, model, optimizer, model_config, cfg, metrics_history)
            print(f"\n[!] Saved checkpoint at step {step}. Re-run to resume.")
            return

        # --- Learning rate ---
        lr = get_lr(step, cfg.warmup_steps, cfg.max_steps, cfg.learning_rate, cfg.min_lr)
        for pg in optimizer.param_groups:
            pg["lr"] = lr

        # --- Forward/backward with gradient accumulation ---
        optimizer.zero_grad()
        accum_loss = 0.0

        for _ in range(cfg.gradient_accumulation_steps):
            try:
                x, y = next(train_iter)
            except StopIteration:
                train_iter = iter(train_loader)
                x, y = next(train_iter)

            x, y = x.to(device), y.to(device)
            tokens_processed += x.numel()

            with torch.autocast(device_type="cuda", dtype=dtype):
                logits = model(x)
                loss = torch.nn.functional.cross_entropy(
                    logits.view(-1, logits.size(-1)), y.view(-1)
                )
                loss = loss / cfg.gradient_accumulation_steps
            loss.backward()
            accum_loss += loss.item()

        torch.nn.utils.clip_grad_norm_(model.parameters(), cfg.grad_clip)
        optimizer.step()

        # --- Metrics ---
        elapsed = time.time() - t0
        tok_s = tokens_processed / elapsed if elapsed > 0 else 0
        metrics_history["losses"].append((step, round(accum_loss, 4)))

        # --- Logging ---
        log_entry = {
            "step": step,
            "loss": round(accum_loss, 4),
            "lr": lr,
            "tok_s": round(tok_s),
            "timestamp": time.strftime("%Y-%m-%dT%H:%M:%S"),
        }

        if step % cfg.log_interval == 0:
            print(f"step {step:6d} | loss {accum_loss:.4f} | lr {lr:.2e} | {tok_s:,.0f} tok/s")

        # --- Evaluation ---
        if step > 0 and step % cfg.eval_interval == 0:
            val_loss = evaluate(model, val_loader, device, dtype)
            metrics_history["val_losses"].append((step, round(val_loss, 4)))
            log_entry["val_loss"] = round(val_loss, 4)
            print(f"  >>> eval loss: {val_loss:.4f}")

        # Write log entry
        with open(log_path, "a") as f:
            f.write(json.dumps(log_entry) + "\n")

        # --- Checkpointing ---
        if step > 0 and step % cfg.save_interval == 0:
            ckpt_path = os.path.join(cfg.checkpoint_dir, f"step_{step}.pt")
            save_checkpoint(ckpt_path, step, model, optimizer, model_config, cfg, metrics_history)
            cleanup_old_checkpoints(cfg.checkpoint_dir, args.keep_checkpoints)
            print(f"  >>> checkpoint: {ckpt_path}")

    # --- Final save ---
    final_path = os.path.join(cfg.checkpoint_dir, "final.pt")
    save_checkpoint(final_path, cfg.max_steps, model, optimizer, model_config, cfg, metrics_history)
    print(f"\nTraining complete! Final checkpoint: {final_path}")

    total_time = time.time() - t0
    print(f"Total: {total_time:.1f}s | {tokens_processed / total_time:,.0f} tok/s avg")

    # --- Final evaluation ---
    print("\nRunning final evaluation...")
    val_loss = evaluate(model, val_loader, device, dtype)
    print(f"Final validation loss: {val_loss:.4f}")

    # Save final summary
    summary = {
        "final_step": cfg.max_steps,
        "final_train_loss": metrics_history["losses"][-1][1] if metrics_history["losses"] else None,
        "final_val_loss": round(val_loss, 4),
        "total_time_seconds": round(total_time, 1),
        "avg_tokens_per_sec": round(tokens_processed / total_time),
        "model_params": model.count_parameters(),
        "checkpoint": final_path,
    }
    summary_path = os.path.join(results_dir, "training_summary.json")
    with open(summary_path, "w") as f:
        json.dump(summary, f, indent=2)
    print(f"Summary saved: {summary_path}")


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------

def main():
    parser = argparse.ArgumentParser(description="Unattended ailmo training")
    parser.add_argument("--max-steps", type=int, default=5000)
    parser.add_argument("--lr", type=float, default=3e-4)
    parser.add_argument("--batch-size", type=int, default=32)
    parser.add_argument("--grad-accum", type=int, default=4)
    parser.add_argument("--data-dir", type=str, default="data")
    parser.add_argument("--checkpoint-dir", type=str, default="checkpoints")
    parser.add_argument("--keep-checkpoints", type=int, default=5,
                        help="Number of rolling checkpoints to keep")
    parser.add_argument("--eval-interval", type=int, default=250)
    parser.add_argument("--save-interval", type=int, default=500)
    parser.add_argument("--log-interval", type=int, default=10)
    args = parser.parse_args()

    print("=" * 60)
    print("  ailmo — Unattended Training")
    print("=" * 60)
    train(args)


if __name__ == "__main__":
    main()