train.py

"""
Training module for 63.8M SLM on BookCorpusOpen
Can be called directly inside a notebook cell:
    from config import TrainConfig
    from train import train
    model, tok = train(train_cfg=TrainConfig())
"""

import os
import time
import math
import json
import numpy as np
import matplotlib.pyplot as plt
import torch
from torch.utils.data import DataLoader

from config import GPTConfig, TrainConfig
from model import GPT
from dataset import PackedDataset, build_and_cache_dataset, get_tokenizer
from checkpoint import save_checkpoint, load_checkpoint, mirror_checkpoint


def get_lr(step: int, warmup_steps: int, max_steps: int, peak_lr: float, min_lr: float) -> float:
    if step < warmup_steps:
        return peak_lr * (0.01 + 0.99 * (step / max(1, warmup_steps)))
    if step > max_steps:
        return min_lr
    decay_ratio = (step - warmup_steps) / (max_steps - warmup_steps)
    coeff = 0.5 * (1.0 + math.cos(math.pi * decay_ratio))
    return min_lr + coeff * (peak_lr - min_lr)


@torch.no_grad()
def evaluate_loss(model, val_loader, eval_iters: int, device: str):
    model.eval()
    losses = []
    for i, (x, y) in enumerate(val_loader):
        if i >= eval_iters:
            break
        x, y = x.to(device, non_blocking=True), y.to(device, non_blocking=True)
        with torch.autocast(device_type="cuda" if "cuda" in device else "cpu", dtype=torch.float16):
            _, loss = model(x, targets=y)
        losses.append(loss.item())
    model.train()
    return float(np.mean(losses)) if losses else float("nan")


def save_training_plots(history: dict, output_dir: str):
    os.makedirs(output_dir, exist_ok=True)
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(13, 4.5))

    ax1.plot(history["steps"], history["train_loss"], label="Train Loss", color="#2563eb", alpha=0.8)
    if history["eval_steps"]:
        ax1.plot(history["eval_steps"], history["val_loss"], label="Val Loss", color="#dc2626", marker="o", markersize=4)
    ax1.set_title("Training Loss Progression")
    ax1.set_xlabel("Steps")
    ax1.set_ylabel("Loss")
    ax1.grid(True, alpha=0.3)
    ax1.legend()

    ax2.plot(history["steps"], history["lr"], label="Learning Rate", color="#16a34a")
    ax2.set_title("Cosine LR Schedule")
    ax2.set_xlabel("Steps")
    ax2.set_ylabel("LR")
    ax2.grid(True, alpha=0.3)
    ax2.legend()

    plt.tight_layout()
    plt.savefig(os.path.join(output_dir, "training_curves.png"), dpi=150)
    plt.close()


def train(train_cfg: TrainConfig = None, model_cfg: GPTConfig = None, resume_from: str = None):
    """
    Main training function callable from Jupyter / Kaggle / Colab notebooks.
    """
    if train_cfg is None:
        train_cfg = TrainConfig()
    if model_cfg is None:
        model_cfg = GPTConfig()
    if resume_from is not None:
        train_cfg.resume_from = resume_from

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

    os.makedirs(train_cfg.checkpoint_dir, exist_ok=True)
    if train_cfg.drive_checkpoint_dir:
        os.makedirs(train_cfg.drive_checkpoint_dir, exist_ok=True)

    # 1. Dataset setup
    train_path, val_path = build_and_cache_dataset(
        cache_dir=train_cfg.data_dir,
        max_books=train_cfg.max_books,
        dataset_name=train_cfg.dataset_name,
        corpus_dir=train_cfg.corpus_dir,
        target_tokens=train_cfg.target_tokens,
    )
    train_dataset = PackedDataset(train_path, context_length=model_cfg.context_length)
    val_dataset = PackedDataset(val_path, context_length=model_cfg.context_length)

    train_loader = DataLoader(train_dataset, batch_size=train_cfg.batch_size, shuffle=True, num_workers=2, pin_memory=True)
    val_loader = DataLoader(val_dataset, batch_size=train_cfg.batch_size, shuffle=False, num_workers=2, pin_memory=True)

    # 2. Model setup
    model = GPT(model_cfg).to(device)
    print(f"[*] Total Model Parameters: {model.get_num_params():,} (~{model.get_num_params()/1e6:.2f}M)")

    # 3. Optimizer setup with weight decay splitting
    optimizer = model.configure_optimizers(
        weight_decay=train_cfg.weight_decay,
        learning_rate=train_cfg.peak_lr,
        betas=(train_cfg.beta1, train_cfg.beta2),
        device_type=device
    )
    scaler = torch.amp.GradScaler("cuda", enabled=(device == "cuda" and train_cfg.use_fp16))
    tokenizer = get_tokenizer()

    start_step = 0
    best_val_loss = float("inf")

    # 4. Resume from Checkpoint if specified
    if train_cfg.resume_from is not None:
        start_step, best_val_loss = load_checkpoint(
            train_cfg.resume_from, model, optimizer=optimizer, scaler=scaler, map_location=device
        )
        if best_val_loss is None:
            best_val_loss = float("inf")

    # Save model config
    with open(os.path.join(train_cfg.checkpoint_dir, "model_config.json"), "w") as f:
        json.dump(model_cfg.__dict__, f, indent=2)

    history = {"steps": [], "train_loss": [], "eval_steps": [], "val_loss": [], "lr": []}
    step = start_step
    start_time = time.time()
    train_iter = iter(train_loader)

    print(f"[*] Training from step {start_step + 1} to {train_cfg.max_steps}...")
    model.train()

    while step < train_cfg.max_steps:
        optimizer.zero_grad(set_to_none=True)
        accum_loss = 0.0

        lr = get_lr(step, train_cfg.warmup_steps, train_cfg.max_steps, train_cfg.peak_lr, train_cfg.min_lr)
        for param_group in optimizer.param_groups:
            param_group["lr"] = lr

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

            x, y = x.to(device, non_blocking=True), y.to(device, non_blocking=True)

            device_type = "cuda" if "cuda" in device else "cpu"
            with torch.autocast(device_type=device_type, dtype=torch.float16, enabled=train_cfg.use_fp16):
                _, loss = model(x, targets=y)
                loss = loss / train_cfg.grad_accum

            accum_loss += loss.item()

            if train_cfg.use_fp16 and device == "cuda":
                scaler.scale(loss).backward()
            else:
                loss.backward()

        if train_cfg.use_fp16 and device == "cuda":
            scaler.unscale_(optimizer)
            torch.nn.utils.clip_grad_norm_(model.parameters(), train_cfg.grad_clip)
            scaler.step(optimizer)
            scaler.update()
        else:
            torch.nn.utils.clip_grad_norm_(model.parameters(), train_cfg.grad_clip)
            optimizer.step()

        step += 1

        # Logging
        if step % 100 == 0:
            elapsed = time.time() - start_time
            print(f"step {step:6d} | train_loss {accum_loss:.4f} | lr {lr:.2e} | {elapsed:.1f}s")
            history["steps"].append(step)
            history["train_loss"].append(accum_loss)
            history["lr"].append(lr)

        # Evaluation & Checkpointing
        if step % train_cfg.eval_interval == 0 or step == train_cfg.max_steps:
            val_loss = evaluate_loss(model, val_loader, train_cfg.eval_iters, device)
            history["eval_steps"].append(step)
            history["val_loss"].append(val_loss)
            print(f"step {step:6d} | val_loss {val_loss:.4f}")

            # Save periodic checkpoint
            ckpt_path = os.path.join(train_cfg.checkpoint_dir, f"ckpt_step{step}.pt")
            save_checkpoint(ckpt_path, model, optimizer, scaler, step=step, val_loss=val_loss, config=model_cfg)
            print(f"Saved checkpoint at step {step} -> {ckpt_path}")
            mirrored_ckpt = mirror_checkpoint(ckpt_path, train_cfg.drive_checkpoint_dir)
            if mirrored_ckpt:
                print(f"Mirrored checkpoint -> {mirrored_ckpt}")

            # Save best model
            if val_loss < best_val_loss:
                best_val_loss = val_loss
                best_path = os.path.join(train_cfg.checkpoint_dir, "best_model.pt")
                save_checkpoint(best_path, model, optimizer, scaler, step=step, val_loss=val_loss, config=model_cfg)
                mirrored_best = mirror_checkpoint(best_path, train_cfg.drive_checkpoint_dir)
                if mirrored_best:
                    print(f"Mirrored best checkpoint -> {mirrored_best}")

                weights_path = os.path.join(train_cfg.checkpoint_dir, "slm_weights.pt")
                torch.save(model.state_dict(), weights_path)
                mirrored_weights = mirror_checkpoint(weights_path, train_cfg.drive_checkpoint_dir)
                if mirrored_weights:
                    print(f"Mirrored weights -> {mirrored_weights}")

            save_training_plots(history, train_cfg.checkpoint_dir)

        # Sample generation
        if step % train_cfg.sample_interval == 0:
            model.eval()
            prompt_tokens = tokenizer.encode(train_cfg.sample_prompt)
            idx = torch.tensor(prompt_tokens, dtype=torch.long, device=device).unsqueeze(0)
            sample_idx = model.generate(idx, max_new_tokens=60, temperature=0.8, top_k=40)
            print("-" * 60)
            print(f"[Sample @ step {step}]: " + tokenizer.decode(sample_idx[0].tolist()))
            print("-" * 60)
            model.train()

    print("Training complete.")
    return model, tokenizer


if __name__ == "__main__":
    train()