train.py

#!/usr/bin/env python3
"""
v3GPT — production training script (built for an Azure A100, runs anywhere with CUDA).

Trains the 100M Hinglish/Indian-internet model FROM SCRATCH on the packed token
binaries. Self-contained on purpose: the model is the same architecture validated
in notebook 01 (RoPE + GQA + SwiGLU + RMSNorm), inlined here so Azure only needs
ONE file + the data.

Usage (on the GPU box, in a dir containing tokenizer.json, train.bin, val.bin):
    pip install torch numpy tokenizers
    python train.py                 # fresh run (or auto-resumes if out/ckpt.pt exists)

Key techniques:
  - bf16 mixed precision  (fp32 master weights + bf16 compute  -> full quality, ~2x speed)
  - gradient accumulation (~0.5M-token effective batch -> stable gradients)
  - resume-safe checkpoints (survives spot preemption / crashes)
  - warmup -> cosine LR, grad clipping
  - val-loss early stopping (stop when a small model saturates / plateaus)
"""

import math
import glob
import os
import time
from contextlib import nullcontext
from dataclasses import dataclass, field
from typing import Optional

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


# ============================================================================
# CONFIG  —  edit here
# ============================================================================
@dataclass
class Config:
    # --- model: 100M default.  For 200M (Chinchilla-optimal for 3.8B tokens):
    #     d_model=1024, n_layers=16, n_heads=16, n_kv_heads=4, d_ff=2816 ---
    vocab_size: int = 32000
    max_seq_len: int = 1024
    d_model: int = 768
    n_layers: int = 12
    n_heads: int = 12
    n_kv_heads: int = 3
    d_ff: int = 2048
    dropout: float = 0.0          # over-fed corpus -> no overfit -> no dropout
    rope_theta: float = 10000.0
    norm_eps: float = 1e-5
    tie_embeddings: bool = True

    # --- training schedule ---
    epochs: float = 3.0           # ceiling; early-stop usually halts before this
    batch_size: int = 32          # micro-batch (sequences/forward). LOWER if OOM, RAISE if memory to spare (faster)
    grad_accum: int = 16          # micro-batches per optimizer step. eff batch = 32*1024*16 ≈ 0.5M tokens
    lr: float = 6e-4
    min_lr: float = 6e-5
    warmup_iters: int = 300
    weight_decay: float = 0.1
    beta1: float = 0.9
    beta2: float = 0.95
    grad_clip: float = 1.0

    # --- eval / checkpoint / early stop ---
    eval_interval: int = 250
    eval_iters: int = 100
    ckpt_interval: int = 500
    early_stop_patience: int = 8  # stop if val loss hasn't improved for this many evals

    # --- io ---
    data_dir: str = "."
    out_dir: str = "out"
    resume_dir: str = ""          # external ckpt dir to warm-start from (Kaggle dataset mount); "" = none
    warm_start: bool = False      # load WEIGHTS ONLY from the resume ckpt -> fresh optimizer + LR, iter 0 (SFT)
    seed: int = 1337
    compile_model: bool = True    # torch.compile ~1.5-2x on CUDA; set False if it errors
    max_steps: int = 0            # >0 caps total steps (smoke test); 0 = use epochs


# ============================================================================
# MODEL  (same architecture as notebook 01, inlined for a single-file Azure run)
# ============================================================================
class RMSNorm(nn.Module):
    def __init__(self, dim, eps=1e-5):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))

    def forward(self, x):
        normed = x.float() * torch.rsqrt(x.float().pow(2).mean(-1, keepdim=True) + self.eps)
        return normed.type_as(x) * self.weight


def precompute_rope_cache(head_dim, max_seq_len, theta):
    inv_freq = 1.0 / (theta ** (torch.arange(0, head_dim, 2).float() / head_dim))
    angles = torch.outer(torch.arange(max_seq_len).float(), inv_freq)
    angles = torch.cat((angles, angles), dim=-1)
    return angles.cos(), angles.sin()


def rotate_half(x):
    x1, x2 = x.chunk(2, dim=-1)
    return torch.cat((-x2, x1), dim=-1)


def apply_rotary_emb(x, cos, sin):
    cos = cos[None, :, None, :].to(x.dtype)
    sin = sin[None, :, None, :].to(x.dtype)
    return (x * cos) + (rotate_half(x) * sin)


def repeat_kv(x, n_rep):
    if n_rep == 1:
        return x
    b, n_kv, s, hd = x.shape
    return x[:, :, None, :, :].expand(b, n_kv, n_rep, s, hd).reshape(b, n_kv * n_rep, s, hd)


class GroupedQueryAttention(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.n_heads, self.n_kv_heads = cfg.n_heads, cfg.n_kv_heads
        self.head_dim = cfg.d_model // cfg.n_heads
        self.n_rep = cfg.n_heads // cfg.n_kv_heads
        self.dropout_p = cfg.dropout
        self.wq = nn.Linear(cfg.d_model, self.n_heads * self.head_dim, bias=False)
        self.wk = nn.Linear(cfg.d_model, self.n_kv_heads * self.head_dim, bias=False)
        self.wv = nn.Linear(cfg.d_model, self.n_kv_heads * self.head_dim, bias=False)
        self.wo = nn.Linear(self.n_heads * self.head_dim, cfg.d_model, bias=False)

    def forward(self, x, cos, sin):
        B, T, _ = x.shape
        q = self.wq(x).view(B, T, self.n_heads, self.head_dim)
        k = self.wk(x).view(B, T, self.n_kv_heads, self.head_dim)
        v = self.wv(x).view(B, T, self.n_kv_heads, self.head_dim)
        q, k = apply_rotary_emb(q, cos, sin), apply_rotary_emb(k, cos, sin)
        q, k, v = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)
        k, v = repeat_kv(k, self.n_rep), repeat_kv(v, self.n_rep)
        out = F.scaled_dot_product_attention(
            q, k, v, is_causal=True, dropout_p=self.dropout_p if self.training else 0.0)
        return self.wo(out.transpose(1, 2).contiguous().view(B, T, -1))


class SwiGLU(nn.Module):
    def __init__(self, d_model, d_ff):
        super().__init__()
        self.w_gate = nn.Linear(d_model, d_ff, bias=False)
        self.w_up = nn.Linear(d_model, d_ff, bias=False)
        self.w_down = nn.Linear(d_ff, d_model, bias=False)

    def forward(self, x):
        return self.w_down(F.silu(self.w_gate(x)) * self.w_up(x))


class Block(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.attn_norm = RMSNorm(cfg.d_model, cfg.norm_eps)
        self.attn = GroupedQueryAttention(cfg)
        self.ffn_norm = RMSNorm(cfg.d_model, cfg.norm_eps)
        self.ffn = SwiGLU(cfg.d_model, cfg.d_ff)
        self.dropout = nn.Dropout(cfg.dropout)

    def forward(self, x, cos, sin):
        x = x + self.dropout(self.attn(self.attn_norm(x), cos, sin))
        x = x + self.dropout(self.ffn(self.ffn_norm(x)))
        return x


class Transformer(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.cfg = cfg
        self.tok_emb = nn.Embedding(cfg.vocab_size, cfg.d_model)
        self.blocks = nn.ModuleList(Block(cfg) for _ in range(cfg.n_layers))
        self.norm = RMSNorm(cfg.d_model, cfg.norm_eps)
        self.lm_head = nn.Linear(cfg.d_model, cfg.vocab_size, bias=False)
        if cfg.tie_embeddings:
            self.lm_head.weight = self.tok_emb.weight
        cos, sin = precompute_rope_cache(cfg.d_model // cfg.n_heads, cfg.max_seq_len, cfg.rope_theta)
        self.register_buffer("rope_cos", cos, persistent=False)
        self.register_buffer("rope_sin", sin, persistent=False)
        self.apply(self._init)
        scale = 0.02 / math.sqrt(2 * cfg.n_layers)          # GPT-2 residual scaling
        for name, p in self.named_parameters():
            if name.endswith(("wo.weight", "w_down.weight")):
                nn.init.normal_(p, 0.0, scale)

    def _init(self, m):
        if isinstance(m, nn.Linear):
            nn.init.normal_(m.weight, 0.0, 0.02)
            if m.bias is not None:
                nn.init.zeros_(m.bias)
        elif isinstance(m, nn.Embedding):
            nn.init.normal_(m.weight, 0.0, 0.02)

    def forward(self, idx, targets=None):
        T = idx.size(1)
        x = self.tok_emb(idx)
        cos, sin = self.rope_cos[:T], self.rope_sin[:T]
        for block in self.blocks:
            x = block(x, cos, sin)
        logits = self.lm_head(self.norm(x))
        loss = None
        if targets is not None:
            loss = F.cross_entropy(logits.reshape(-1, logits.size(-1)), targets.reshape(-1), ignore_index=-100)
        return logits, loss

    def num_params(self, non_embedding=True):
        n = sum(p.numel() for p in self.parameters())
        return n - self.tok_emb.weight.numel() if non_embedding else n


# ============================================================================
# TRAINING
# ============================================================================
def main(cfg=None):
    cfg = cfg or Config()
    os.makedirs(cfg.out_dir, exist_ok=True)
    torch.manual_seed(cfg.seed)
    torch.set_float32_matmul_precision("high")              # enable TF32 for any fp32 matmuls

    device = "cuda" if torch.cuda.is_available() else ("mps" if torch.backends.mps.is_available() else "cpu")
    dev_type = "cuda" if device == "cuda" else "cpu"
    # Kaggle GPUs (P100/T4) have no bf16 -> fp16 autocast + GradScaler
    ctx = torch.autocast(device_type=dev_type, dtype=torch.float16) if dev_type == "cuda" else nullcontext()
    scaler = torch.cuda.amp.GradScaler(enabled=(dev_type == "cuda"))
    print(f"device={device} | fp16 autocast={'on' if dev_type=='cuda' else 'off (cpu/mps fallback)'}")

    # --- data (memory-mapped token streams) ---
    # Resolve data_dir robustly: if it has no train.bin, find it under a mounted Kaggle input dataset.
    _dd = cfg.data_dir
    if not os.path.exists(os.path.join(_dd, "train.bin")):
        _hits = sorted(glob.glob("/kaggle/input/**/train.bin", recursive=True))
        if _hits:
            _dd = os.path.dirname(_hits[0]); print(f"data_dir resolved via glob -> {_dd}")
    train_data = np.memmap(os.path.join(_dd, "train.bin"), dtype=np.uint16, mode="r")
    val_data = np.memmap(os.path.join(_dd, "val.bin"), dtype=np.uint16, mode="r")
    # optional SFT loss mask (uint8: 1 = train on this target token). Absent => loss on all tokens.
    def _opt_mask(name):
        p = os.path.join(_dd, name)
        return np.memmap(p, dtype=np.uint8, mode="r") if os.path.exists(p) else None
    train_mask, val_mask = _opt_mask("train_mask.bin"), _opt_mask("val_mask.bin")
    if train_mask is not None:
        print("loss-masking ON (train_mask.bin found) -> loss only on response tokens")
    tokens_per_step = cfg.batch_size * cfg.max_seq_len * cfg.grad_accum
    max_iters = int(cfg.epochs * (len(train_data) // tokens_per_step))
    print(f"train={len(train_data):,} tok | val={len(val_data):,} tok | "
          f"{tokens_per_step:,} tok/step | {max_iters:,} steps (~{cfg.epochs} epochs)")

    def get_batch(split):
        data = train_data if split == "train" else val_data
        msk = train_mask if split == "train" else val_mask
        ix = torch.randint(len(data) - cfg.max_seq_len, (cfg.batch_size,))
        x = torch.stack([torch.from_numpy(data[i:i + cfg.max_seq_len].astype(np.int64)) for i in ix])
        y = torch.stack([torch.from_numpy(data[i + 1:i + 1 + cfg.max_seq_len].astype(np.int64)) for i in ix])
        if msk is not None:                                   # SFT: loss only where mask==1 (response tokens)
            m = torch.stack([torch.from_numpy(msk[i + 1:i + 1 + cfg.max_seq_len].astype(np.bool_)) for i in ix])
            y = y.masked_fill(~m, -100)                       # -100 = ignore_index in cross_entropy
        if device == "cuda":
            return x.pin_memory().to(device, non_blocking=True), y.pin_memory().to(device, non_blocking=True)
        return x.to(device), y.to(device)

    # --- model + optimizer ---
    model = Transformer(cfg).to(device)
    print(f"params: {model.num_params(False):,} total (~{model.num_params(False)/1e6:.0f}M)")
    decay = [p for p in model.parameters() if p.dim() >= 2]
    nodecay = [p for p in model.parameters() if p.dim() < 2]
    optimizer = torch.optim.AdamW(
        [{"params": decay, "weight_decay": cfg.weight_decay},
         {"params": nodecay, "weight_decay": 0.0}],
        lr=cfg.lr, betas=(cfg.beta1, cfg.beta2),
        fused=(device == "cuda"))

    # --- resume / warm-start ---
    # Own out_dir ckpt ALWAYS wins (full resume of an interrupted run); else an external ckpt
    # (resume_dir or any mounted /kaggle/input/**/ckpt.pt). warm_start => load weights only, iter 0.
    iter_num, best_val, patience = 0, float("inf"), 0
    _own = os.path.join(cfg.out_dir, "ckpt.pt")
    _cands = [_own]
    _cands += ([os.path.join(cfg.resume_dir, "ckpt.pt")] if cfg.resume_dir else [])
    _cands += sorted(glob.glob("/kaggle/input/**/ckpt.pt", recursive=True))
    ckpt_path = next((p for p in _cands if os.path.exists(p)), None)
    if ckpt_path:
        ck = torch.load(ckpt_path, map_location=device)
        model.load_state_dict(ck["model"])
        if cfg.warm_start and ckpt_path != _own:
            print(f"WARM-START from {ckpt_path}: weights only -> fresh optimizer + LR schedule (iter 0)")
        else:
            optimizer.load_state_dict(ck["optimizer"])
            iter_num, best_val = ck["iter_num"], ck["best_val"]
            print(f"RESUMED from {ckpt_path} at iter {iter_num} (best val {best_val:.4f})")

    raw_model = model
    if cfg.compile_model and device == "cuda":
        model = torch.compile(model)
        print("torch.compile: on")

    def get_lr(it):
        if it < cfg.warmup_iters:
            return cfg.lr * (it + 1) / cfg.warmup_iters
        if it > max_iters:
            return cfg.min_lr
        r = (it - cfg.warmup_iters) / max(1, max_iters - cfg.warmup_iters)
        return cfg.min_lr + 0.5 * (1 + math.cos(math.pi * r)) * (cfg.lr - cfg.min_lr)

    @torch.no_grad()
    def estimate_loss():
        model.eval()
        out = {}
        for split in ("train", "val"):
            losses = torch.zeros(cfg.eval_iters)
            for k in range(cfg.eval_iters):
                X, Y = get_batch(split)
                with ctx:
                    _, loss = model(X, Y)
                losses[k] = loss.item()
            out[split] = losses.mean().item()
        model.train()
        return out

    def save(name):
        torch.save({"model": raw_model.state_dict(), "optimizer": optimizer.state_dict(),
                    "config": vars(cfg), "iter_num": iter_num, "best_val": best_val},
                   os.path.join(cfg.out_dir, name))

    # --- the loop ---
    log = open(os.path.join(cfg.out_dir, "log.csv"), "a")
    if iter_num == 0:
        log.write("iter,train_loss,val_loss,lr,elapsed_s\n")
    model.train()
    x, y = get_batch("train")               # prefetch first batch
    t0 = time.time()
    while iter_num <= max_iters:
        for g in optimizer.param_groups:
            g["lr"] = get_lr(iter_num)

        if iter_num % cfg.eval_interval == 0:
            L = estimate_loss()
            el = time.time() - t0
            print(f"iter {iter_num:6d}/{max_iters} | train {L['train']:.4f} | val {L['val']:.4f} "
                  f"| lr {get_lr(iter_num):.2e} | {el/60:.1f}min", flush=True)
            log.write(f"{iter_num},{L['train']:.4f},{L['val']:.4f},{get_lr(iter_num):.3e},{el:.0f}\n"); log.flush()
            if L["val"] < best_val:
                best_val, patience = L["val"], 0
                save("best.pt")                 # keep the best-val model separately
            elif iter_num > cfg.warmup_iters:
                patience += 1
                if patience >= cfg.early_stop_patience:
                    print(f"EARLY STOP: val loss plateaued for {patience} evals (best {best_val:.4f})")
                    break
        if iter_num > 0 and iter_num % cfg.ckpt_interval == 0:
            save("ckpt.pt")                     # latest, for resume

        # forward/backward with gradient accumulation (fp16 AMP -> GradScaler)
        for _ in range(cfg.grad_accum):
            with ctx:
                _, loss = model(x, y)
                loss = loss / cfg.grad_accum
            x, y = get_batch("train")           # fetch next batch while GPU works on this one
            scaler.scale(loss).backward()
        scaler.unscale_(optimizer)
        torch.nn.utils.clip_grad_norm_(model.parameters(), cfg.grad_clip)
        scaler.step(optimizer)
        scaler.update()
        optimizer.zero_grad(set_to_none=True)
        iter_num += 1
        if cfg.max_steps and iter_num >= cfg.max_steps:
            print(f"max_steps {cfg.max_steps} reached"); break

    save("ckpt.pt")
    log.close()
    print(f"DONE at iter {iter_num} | best val loss {best_val:.4f} | "
          f"final model -> {cfg.out_dir}/best.pt")


if __name__ == "__main__":
    main()