src/train.py

"""
Base pretraining script for Retriever500M.

Memory optimizations for 8GB VRAM (RTX 4070 Laptop):
  - bf16 mixed precision (autocast)
  - Gradient checkpointing (recompute activations during backward)
  - 8-bit AdamW optimizer (bitsandbytes) — halves optimizer state memory
  - Gradient accumulation (effective batch size > micro batch size)
  - Short sequence length (512 tokens) for base training
  - Tied embeddings (shared input/output weight)
  - Flash Attention via torch SDPA

Usage:
  python src/train.py [--steps N] [--seq_len N] [--batch_size N] [--grad_accum N]
"""

import argparse
import json
import os
import sys
import time
from dataclasses import asdict

import numpy as np
import torch
import torch.nn.functional as F
from tqdm import tqdm

# Add src to path
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from model import ModelConfig, Retriever500M
from tokenizers import Tokenizer

# ─── Paths ───────────────────────────────────────────────────────────────────
PROJECT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
DATA_DIR = os.path.join(PROJECT_DIR, "data")
TOKENIZER_DIR = os.path.join(PROJECT_DIR, "tokenizer")
CHECKPOINT_DIR = os.path.join(PROJECT_DIR, "checkpoints")
LOGS_DIR = os.path.join(PROJECT_DIR, "logs")

CORPUS_PATH = os.path.join(DATA_DIR, "corpus.txt")
CURATED_CORPUS_PATH = os.path.join(DATA_DIR, "corpus_curated.txt")
TOKENIZER_PATH = os.path.join(TOKENIZER_DIR, "tokenizer.json")

# ─── Data Loading ────────────────────────────────────────────────────────────

def load_and_tokenize(corpus_path: str, tokenizer: Tokenizer) -> np.ndarray:
    """Load corpus, tokenize everything, return a flat numpy array of token IDs."""
    print(f"Loading corpus from {corpus_path}...")
    with open(corpus_path, "r", encoding="utf-8") as f:
        text = f.read()

    print(f"Corpus size: {len(text) / 1e6:.1f} MB")

    # Tokenize in chunks to avoid memory issues
    chunk_size = 1_000_000  # 1MB chunks
    all_tokens = []

    print("Tokenizing corpus...")
    for i in tqdm(range(0, len(text), chunk_size)):
        chunk = text[i : i + chunk_size]
        encoded = tokenizer.encode(chunk)
        all_tokens.extend(encoded.ids)

    tokens = np.array(all_tokens, dtype=np.int32)
    print(f"Total tokens: {len(tokens):,}")
    return tokens


def get_batch(
    tokens: np.ndarray,
    batch_size: int,
    seq_len: int,
    device: torch.device,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Sample a random batch of sequences from the token array.

    Returns (input_ids, targets) where targets are shifted by 1.
    """
    # Random starting indices
    max_start = len(tokens) - seq_len - 1
    indices = np.random.randint(0, max_start, size=batch_size)

    # Gather sequences
    input_ids = np.stack([tokens[i : i + seq_len] for i in indices])
    targets = np.stack([tokens[i + 1 : i + seq_len + 1] for i in indices])

    input_ids = torch.from_numpy(input_ids).long().to(device)
    targets = torch.from_numpy(targets).long().to(device)

    return input_ids, targets


# ─── Training ────────────────────────────────────────────────────────────────

def setup_optimizer(model: Retriever500M, lr: float, use_8bit: bool = True):
    """Set up optimizer — 8-bit AdamW if available, else standard AdamW."""
    # Separate embedding params (no weight decay) from rest
    decay_params = []
    no_decay_params = []
    for name, param in model.named_parameters():
        if not param.requires_grad:
            continue
        if "embedding" in name or "norm" in name or "weight" in name and ".weight" not in name:
            no_decay_params.append(param)
        else:
            decay_params.append(param)

    param_groups = [
        {"params": decay_params, "weight_decay": 0.1},
        {"params": no_decay_params, "weight_decay": 0.0},
    ]

    if use_8bit:
        try:
            import bitsandbytes as bnb
            optimizer = bnb.optim.AdamW8bit(
                param_groups, lr=lr, betas=(0.9, 0.95), eps=1e-8,
            )
            print("Using 8-bit AdamW (bitsandbytes)")
            return optimizer
        except Exception as e:
            print(f"8-bit optimizer unavailable ({e}), falling back to AdamW")

    optimizer = torch.optim.AdamW(
        param_groups, lr=lr, betas=(0.9, 0.95), eps=1e-8,
    )
    print("Using standard AdamW")
    return optimizer


def resume_from_checkpoint(
    model: Retriever500M,
    optimizer,
    resume_path: str,
    device: torch.device,
) -> tuple[int, float]:
    """Load model + optimizer state from a checkpoint.

    Returns (start_step, best_loss) so the training loop can continue.
    """
    print(f"Resuming from {resume_path}")
    ckpt = torch.load(resume_path, map_location=device, weights_only=False)

    model.load_state_dict(ckpt["model_state_dict"])
    print(f"  Loaded model weights (step {ckpt.get('step', '?')})")

    if "optimizer_state_dict" in ckpt:
        try:
            optimizer.load_state_dict(ckpt["optimizer_state_dict"])
            print("  Loaded optimizer state")
        except Exception as e:
            print(f"  Could not load optimizer state ({e}); starting fresh optimizer")

    start_step = int(ckpt.get("step", 0))
    best_loss = float(ckpt.get("loss", float("inf")))
    print(f"  Resuming at step {start_step} (best_loss={best_loss:.4f})")
    return start_step, best_loss


def get_lr(step: int, warmup_steps: int, max_steps: int, max_lr: float, min_lr: float) -> float:
    """Cosine learning rate schedule with linear warmup."""
    if step < warmup_steps:
        return max_lr * (step + 1) / warmup_steps
    if step > max_steps:
        return min_lr
    decay_ratio = (step - warmup_steps) / (max_steps - warmup_steps)
    coeff = 0.5 * (1.0 + np.cos(np.pi * decay_ratio))
    return min_lr + coeff * (max_lr - min_lr)


def train(args):
    # ─── Setup ───────────────────────────────────────────────────────────────
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Device: {device}")
    if device.type == "cuda":
        print(f"GPU: {torch.cuda.get_device_name(0)}")
        print(f"VRAM: {torch.cuda.get_device_properties(0).total_memory / 1e9:.2f} GB")

    os.makedirs(CHECKPOINT_DIR, exist_ok=True)
    os.makedirs(LOGS_DIR, exist_ok=True)

    # ─── Tokenizer ───────────────────────────────────────────────────────────
    print("Loading tokenizer...")
    tokenizer = Tokenizer.from_file(TOKENIZER_PATH)
    vocab_size = tokenizer.get_vocab_size()
    print(f"Vocab size: {vocab_size}")

    # ─── Data ────────────────────────────────────────────────────────────────
    if args.corpus == "curated":
        corpus_path = CURATED_CORPUS_PATH
        if not os.path.exists(corpus_path):
            raise FileNotFoundError(f"Curated corpus not found: {corpus_path}. Run src/curate.py first.")
        print(f"Using CURATED corpus: {corpus_path}")
    elif args.corpus == "default":
        corpus_path = CORPUS_PATH
    else:
        corpus_path = args.corpus
    tokens = load_and_tokenize(corpus_path, tokenizer)

    # ─── Model ───────────────────────────────────────────────────────────────
    config = ModelConfig(
        vocab_size=vocab_size,
        d_model=1_280,
        n_layers=23,
        n_heads=20,
        d_ff=3_456,
        max_seq_len=args.seq_len,
        dropout=0.0,
        tie_embeddings=True,
    )

    model = Retriever500M(config).to(device)
    total_params = model.count_parameters()
    print(f"Model parameters: {total_params:,} ({total_params / 1e6:.1f}M)")

    # ─── Optimizer ───────────────────────────────────────────────────────────
    optimizer = setup_optimizer(model, args.lr, use_8bit=args.use_8bit_adam)

    # ─── Resume from checkpoint ──────────────────────────────────────────────
    start_step = 0
    best_loss = float("inf")
    prev_log_steps = []
    if args.resume:
        resume_path = args.resume_path or os.path.join(CHECKPOINT_DIR, "latest.pt")
        if not os.path.exists(resume_path):
            raise FileNotFoundError(f"Cannot resume: {resume_path} does not exist")
        start_step, best_loss = resume_from_checkpoint(model, optimizer, resume_path, device)
        accum_loss = best_loss  # continue EMA from saved loss
        # Load previous log entries so we append rather than overwrite history
        prev_log_path = os.path.join(LOGS_DIR, "training_log.json")
        if os.path.exists(prev_log_path):
            try:
                with open(prev_log_path, "r") as f:
                    prev_log = json.load(f)
                prev_log_steps = prev_log.get("steps", [])
                print(f"  Loaded {len(prev_log_steps)} previous log entries")
            except Exception as e:
                print(f"  Could not load previous log ({e})")
    else:
        accum_loss = 0.0

    # ─── Training loop ───────────────────────────────────────────────────────
    effective_batch = args.batch_size * args.grad_accum
    max_steps_total = start_step + args.steps  # for LR schedule continuity
    print(f"\nTraining configuration:")
    print(f"  Micro batch size:  {args.batch_size}")
    print(f"  Gradient accum:    {args.grad_accum}")
    print(f"  Effective batch:   {effective_batch}")
    print(f"  Sequence length:   {args.seq_len}")
    print(f"  Learning rate:     {args.lr}")
    print(f"  Steps this run:    {args.steps}")
    print(f"  Start step:        {start_step}")
    print(f"  Target step:       {max_steps_total}")
    print(f"  Warmup steps:      {args.warmup}")
    print(f"  Grad checkpointing: {args.grad_checkpoint}")
    print()

    # Training log
    log = {
        "config": asdict(config),
        "train_args": vars(args),
        "total_params": total_params,
        "steps": list(prev_log_steps),  # carry over previous entries
    }

    model.train()
    step = start_step
    start_time = time.time()

    pbar = tqdm(range(start_step, max_steps_total), desc="Training", initial=start_step, total=max_steps_total)
    for step in pbar:
        # Learning rate schedule (uses absolute step for continuity)
        lr = get_lr(step, args.warmup, max_steps_total, args.lr, args.lr * 0.1)
        for pg in optimizer.param_groups:
            pg["lr"] = lr

        optimizer.zero_grad(set_to_none=True)

        # Gradient accumulation
        total_loss = 0.0
        for micro_step in range(args.grad_accum):
            input_ids, targets = get_batch(tokens, args.batch_size, args.seq_len, device)

            with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
                out = model(input_ids, targets=targets, use_checkpoint=args.grad_checkpoint)
                loss = out["loss"] / args.grad_accum

            loss.backward()
            total_loss += loss.item()

        # Gradient clipping
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)

        # Optimizer step
        optimizer.step()

        avg_loss = total_loss  # already divided by grad_accum
        accum_loss = accum_loss * 0.95 + avg_loss * 0.05  # EMA

        # Logging
        if step % args.log_every == 0 or step == max_steps_total - 1:
            elapsed = time.time() - start_time
            steps_this_run = step - start_step + 1
            steps_per_sec = steps_this_run / elapsed
            vram_used = torch.cuda.max_memory_allocated() / 1e9 if device.type == "cuda" else 0

            log_entry = {
                "step": step,
                "loss": avg_loss,
                "ema_loss": accum_loss,
                "lr": lr,
                "elapsed_s": elapsed,
                "steps_per_sec": steps_per_sec,
                "vram_gb": vram_used,
            }
            log["steps"].append(log_entry)

            pbar.set_postfix({
                "loss": f"{avg_loss:.4f}",
                "ema": f"{accum_loss:.4f}",
                "lr": f"{lr:.2e}",
                "vram": f"{vram_used:.1f}G",
            })

        # Save checkpoint
        if (step + 1) % args.save_every == 0 or step == max_steps_total - 1:
            ckpt_path = os.path.join(CHECKPOINT_DIR, f"model_step_{step + 1}.pt")
            torch.save({
                "model_state_dict": model.state_dict(),
                "optimizer_state_dict": optimizer.state_dict(),
                "config": asdict(config),
                "step": step + 1,
                "loss": accum_loss,
            }, ckpt_path)
            print(f"\n  Saved checkpoint: {ckpt_path}")

            # Also save latest
            latest_path = os.path.join(CHECKPOINT_DIR, "latest.pt")
            torch.save({
                "model_state_dict": model.state_dict(),
                "config": asdict(config),
                "step": step + 1,
                "loss": accum_loss,
            }, latest_path)

            if accum_loss < best_loss:
                best_loss = accum_loss
                best_path = os.path.join(CHECKPOINT_DIR, "best.pt")
                torch.save({
                    "model_state_dict": model.state_dict(),
                    "config": asdict(config),
                    "step": step + 1,
                    "loss": accum_loss,
                }, best_path)

        # Reset peak memory stats periodically
        if step % 50 == 0 and device.type == "cuda":
            torch.cuda.reset_peak_memory_stats()

    # ─── Save training log ───────────────────────────────────────────────────
    log_path = os.path.join(LOGS_DIR, "training_log.json")
    with open(log_path, "w") as f:
        json.dump(log, f, indent=2)
    print(f"\nTraining log saved to {log_path}")

    total_time = time.time() - start_time
    print(f"\nTraining complete!")
    print(f"  Total time:     {total_time:.1f}s ({total_time/60:.1f} min)")
    print(f"  Final EMA loss: {accum_loss:.4f}")
    print(f"  Best loss:      {best_loss:.4f}")
    print(f"  Steps/sec:      {args.steps / total_time:.2f}")

    return model, log


def main():
    parser = argparse.ArgumentParser(description="Train Retriever500M base model")
    parser.add_argument("--steps", type=int, default=2000, help="Total training steps")
    parser.add_argument("--batch_size", type=int, default=4, help="Micro batch size")
    parser.add_argument("--grad_accum", type=int, default=8, help="Gradient accumulation steps")
    parser.add_argument("--seq_len", type=int, default=512, help="Sequence length")
    parser.add_argument("--lr", type=float, default=3e-4, help="Peak learning rate")
    parser.add_argument("--warmup", type=int, default=100, help="Warmup steps")
    parser.add_argument("--save_every", type=int, default=500, help="Save checkpoint every N steps")
    parser.add_argument("--log_every", type=int, default=10, help="Log every N steps")
    parser.add_argument("--grad_checkpoint", action="store_true", default=True, help="Use gradient checkpointing")
    parser.add_argument("--no_grad_checkpoint", dest="grad_checkpoint", action="store_false")
    parser.add_argument("--use_8bit_adam", action="store_true", default=True, help="Use 8-bit AdamW")
    parser.add_argument("--no_8bit_adam", dest="use_8bit_adam", action="store_false")
    parser.add_argument("--resume", action="store_true", help="Resume training from latest checkpoint")
    parser.add_argument("--resume_path", type=str, default=None, help="Specific checkpoint to resume from (default: checkpoints/latest.pt)")
    parser.add_argument("--corpus", type=str, default="default", help="Corpus to use: 'default', 'curated', or a path")
    args = parser.parse_args()

    train(args)


if __name__ == "__main__":
    main()