Update supernova/train.py

supernova/train.py

@@ -1,159 +1,415 @@
-import argparse
-import json
-import math
-import os
-import time
-from typing import Optional
-
-import torch
-import torch.nn as nn
-from torch.utils.data import DataLoader
-from transformers import get_cosine_schedule_with_warmup
-
-from .config import ModelConfig
-from .model import SupernovaModel
-from .tokenizer import load_gpt2_tokenizer
-from .data import load_sources_from_yaml, TokenChunkDataset
-
-
-def compute_grad_norm(model: nn.Module) -> float:
-    total = 0.0
-    for p in model.parameters():
-        if p.grad is not None:
-            param_norm = p.grad.data.float().norm(2).item()
-            total += param_norm * param_norm
-    return math.sqrt(total)
-
-
-def train(
-    config_path: str,
-    data_config_path: str,
-    seq_len: int = 1024,
-    batch_size: int = 16,
-    grad_accum: int = 8,
-    lr: float = 3e-4,
-    warmup_steps: int = 2000,
-    max_steps: int = 100_000,
-    save_every: int = 10_000,
-    out_dir: str = "checkpoints",
-    seed: int = 42,
-):
-    torch.manual_seed(seed)
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
-    cfg = ModelConfig.from_json_file(config_path)
-    # Assert exact parameter budget from formula
-    cfg.assert_exact_params(expected=25_000_000)
-
-    tok = load_gpt2_tokenizer()
-    assert tok.vocab_size == cfg.vocab_size, (
-        f"Tokenizer vocab size ({tok.vocab_size}) != config ({cfg.vocab_size})"
-    )
-
-    model = SupernovaModel(cfg).to(device)
-
-    # Double-check exact parameter count by instantiating
-    total_params = sum(p.numel() for p in model.parameters())
-    assert total_params == 25_000_000, f"Model has {total_params} params, expected 25,000,000"
-
-    sources = load_sources_from_yaml(data_config_path)
-    ds = TokenChunkDataset(tok, sources, seq_len=seq_len, eos_token_id=tok.eos_token_id)
-    dl = DataLoader(ds, batch_size=batch_size, shuffle=False, num_workers=0)
-
-    optimizer = torch.optim.AdamW(
-        model.parameters(), lr=lr, betas=(0.9, 0.95), weight_decay=0.1
-    )
-
-    # We use a token-based schedule; max_steps is optimizer steps, not micro-steps
-    scheduler = get_cosine_schedule_with_warmup(
-        optimizer,
-        num_warmup_steps=warmup_steps,
-        num_training_steps=max_steps,
-    )
-
-    model.train()
-    os.makedirs(out_dir, exist_ok=True)
-
-    step = 0
-    micro = 0
-    running_loss = 0.0
-    t0 = time.time()
-
-    while step < max_steps:
-        for batch in dl:
-            x, y = batch
-            x = x.to(device)
-            y = y.to(device)
-
-            logits, loss = model(x, y)
-            loss = loss / grad_accum
-            loss.backward()
-
-            micro += 1
-            running_loss += loss.item()
-
-            if micro % grad_accum == 0:
-                # Optional clip: leave off by default for pure monitoring
-                # torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)
-                optimizer.step()
-                optimizer.zero_grad(set_to_none=True)
-                scheduler.step()
-
-                step += 1
-                if step % 50 == 0:
-                    grad_norm = compute_grad_norm(model)
-                    avg_loss = running_loss * grad_accum / 50.0
-                    running_loss = 0.0
-                    elapsed = time.time() - t0
-                    lr_now = scheduler.get_last_lr()[0]
-                    print(f"step={step} loss={avg_loss:.4f} grad_norm={grad_norm:.2f} lr={lr_now:.6f} elapsed={elapsed:.1f}s")
-                    t0 = time.time()
-
-                if save_every and step % save_every == 0:
-                    ckpt_path = os.path.join(out_dir, f"supernova_step{step}.pt")
-                    torch.save({
-                        "model_state_dict": model.state_dict(),
-                        "config": cfg.__dict__,
-                        "step": step,
-                    }, ckpt_path)
-
-                if step >= max_steps:
-                    break
-
-    # final save
-    ckpt_path = os.path.join(out_dir, f"supernova_final.pt")
-    torch.save({
-        "model_state_dict": model.state_dict(),
-        "config": cfg.__dict__,
-        "step": step,
-    }, ckpt_path)
-
-
-if __name__ == "__main__":
-    ap = argparse.ArgumentParser()
-    ap.add_argument("--config", required=True)
-    ap.add_argument("--data-config", required=True)
-    ap.add_argument("--seq-len", type=int, default=1024)
-    ap.add_argument("--batch-size", type=int, default=16)
-    ap.add_argument("--grad-accum", type=int, default=8)
-    ap.add_argument("--lr", type=float, default=3e-4)
-    ap.add_argument("--warmup-steps", type=int, default=2000)
-    ap.add_argument("--max-steps", type=int, default=100000)
-    ap.add_argument("--save-every", type=int, default=10000)
-    ap.add_argument("--out-dir", type=str, default="checkpoints")
-    ap.add_argument("--seed", type=int, default=42)
-    args = ap.parse_args()
-
-    train(
-        config_path=args.config,
-        data_config_path=args.data_config,
-        seq_len=args.seq_len,
-        batch_size=args.batch_size,
-        grad_accum=args.grad_accum,
-        lr=args.lr,
-        warmup_steps=args.warmup_steps,
-        max_steps=args.max_steps,
-        save_every=args.save_every,
-        out_dir=args.out_dir,
-        seed=args.seed,
-    )
+# train.py (improved)
+import argparse
+import json
+import math
+import os
+import time
+from typing import Optional, Dict, Any
+
+import torch
+import torch.nn as nn
+from torch.utils.data import DataLoader, DistributedSampler
+from torch.utils.tensorboard import SummaryWriter
+from transformers import get_cosine_schedule_with_warmup
+
+from .config import ModelConfig
+from .model import SupernovaModel
+from .tokenizer import load_gpt2_tokenizer
+from .data import load_sources_from_yaml, TokenChunkDataset
+
+# -----------------------
+# Utilities
+# -----------------------
+def compute_grad_norm(model: nn.Module) -> float:
+    total = 0.0
+    for p in model.parameters():
+        if p.grad is not None:
+            param_norm = p.grad.data.float().norm(2).item()
+            total += param_norm * param_norm
+    return math.sqrt(total)
+
+def atomic_save(obj: Dict[str, Any], path: str):
+    tmp = path + ".tmp"
+    torch.save(obj, tmp)
+    os.replace(tmp, path)
+
+class EMA:
+    """Simple exponential moving average of model params (maintains shadow copy)."""
+    def __init__(self, model: nn.Module, decay: float = 0.9999):
+        self.decay = decay
+        self.shadow = {}
+        for name, p in model.named_parameters():
+            if p.requires_grad:
+                self.shadow[name] = p.data.clone()
+
+    def update(self, model: nn.Module):
+        for name, p in model.named_parameters():
+            if p.requires_grad:
+                self.shadow[name].mul_(self.decay).add_(p.data, alpha=1.0 - self.decay)
+
+    def store(self, model: nn.Module):
+        self.backup = {n: p.data.clone() for n, p in model.named_parameters() if p.requires_grad}
+
+    def copy_to(self, model: nn.Module):
+        for name, p in model.named_parameters():
+            if p.requires_grad:
+                p.data.copy_(self.shadow[name])
+
+    def restore(self, model: nn.Module):
+        for name, p in model.named_parameters():
+            if p.requires_grad:
+                p.data.copy_(self.backup[name])
+        del self.backup
+
+# -----------------------
+# Training loop
+# -----------------------
+def train(
+    config_path: str,
+    data_config_path: str,
+    seq_len: int = 1024,
+    batch_size: int = 16,
+    grad_accum: int = 8,
+    lr: float = 3e-4,
+    warmup_steps: int = 2000,
+    max_steps: int = 100_000,
+    save_every: int = 10_000,
+    out_dir: str = "checkpoints",
+    seed: int = 42,
+    validate_every: int = 1000,
+    val_steps: int = 100,
+    clip_grad_norm: Optional[float] = 1.0,
+    use_ema: bool = True,
+    ema_decay: float = 0.9999,
+    resume_from: Optional[str] = None,
+    use_tensorboard: bool = True,
+    ddp: bool = False,
+    local_rank: int = 0,
+    num_workers: int = 4,
+    pin_memory: bool = True,
+    compile_model: bool = False,
+):
+    # reproducibility
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    import random
+    random.seed(seed)
+    # performance flags
+    torch.backends.cudnn.benchmark = True
+
+    # device / distributed
+    if ddp:
+        torch.distributed.init_process_group(backend="nccl")
+        device = torch.device(f"cuda:{local_rank}")
+        torch.cuda.set_device(device)
+    else:
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    # config & tokenizer
+    cfg = ModelConfig.from_json_file(config_path)
+    cfg.assert_exact_params(expected=25_000_000)
+    tok = load_gpt2_tokenizer()
+    assert tok.vocab_size == cfg.vocab_size, "Tokenizer vocab size mismatch."
+
+    model = SupernovaModel(cfg)
+    # optional: enable gradient checkpointing for memory saving if model supports it
+    if hasattr(model, "gradient_checkpointing_enable"):
+        try:
+            model.gradient_checkpointing_enable()
+        except Exception:
+            pass
+
+    model.to(device)
+
+    # double-check params
+    total_params = sum(p.numel() for p in model.parameters())
+    assert total_params == 25_000_000, f"Model has {total_params} params, expected 25,000,000"
+
+    # optional compile (PyTorch 2.0)
+    if compile_model:
+        try:
+            model = torch.compile(model)
+        except Exception as e:
+            print("torch.compile not available/failed:", e)
+
+    # DDP wrap
+    if ddp:
+        model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[local_rank], find_unused_parameters=False)
+
+    # dataset and dataloader
+    sources = load_sources_from_yaml(data_config_path)
+    # TODO: improve TokenChunkDataset to perform token-packing (pack multiple short examples into one sequence)
+    ds = TokenChunkDataset(tok, sources, seq_len=seq_len, eos_token_id=tok.eos_token_id)
+
+    sampler = DistributedSampler(ds) if ddp else None
+    dl = DataLoader(
+        ds,
+        batch_size=batch_size,
+        shuffle=(sampler is None),
+        sampler=sampler,
+        num_workers=num_workers,
+        pin_memory=pin_memory,
+        prefetch_factor=2,
+        drop_last=True,
+    )
+
+    # optimizer with simple parameter grouping example to avoid weight decay on norms/bias
+    def param_groups(model):
+        decay, no_decay = [], []
+        for n, p in model.named_parameters():
+            if not p.requires_grad:
+                continue
+            if any(nd in n for nd in ["bias", "ln", "layernorm", "LayerNorm", "norm"]):
+                no_decay.append(p)
+            else:
+                decay.append(p)
+        return [
+            {"params": decay, "weight_decay": 0.1},
+            {"params": no_decay, "weight_decay": 0.0},
+        ]
+
+    optimizer = torch.optim.AdamW(param_groups(model), lr=lr, betas=(0.9, 0.95), eps=1e-8)
+
+    # scheduler
+    scheduler = get_cosine_schedule_with_warmup(optimizer, num_warmup_steps=warmup_steps, num_training_steps=max_steps)
+
+    # AMP scaler
+    scaler = torch.cuda.amp.GradScaler(enabled=(device.type == "cuda"))
+
+    # EMA
+    ema = EMA(model if not ddp else model.module, decay=ema_decay) if use_ema else None
+
+    # logging + checkpoint dir
+    os.makedirs(out_dir, exist_ok=True)
+    writer = SummaryWriter(log_dir=os.path.join(out_dir, "runs")) if use_tensorboard and (not ddp or local_rank == 0) else None
+
+    # validation dataset (simple split: user should provide a separate validation YAML ideally)
+    # TODO: Implement a proper validation dataset pipeline. For now, we use a small random subset of training data.
+    val_ds = None
+    val_dl = None
+
+    # resume
+    start_step = 0
+    best_val_loss = float("inf")
+    if resume_from and os.path.exists(resume_from):
+        ckpt = torch.load(resume_from, map_location=device)
+        model_state = ckpt["model_state_dict"]
+        # if ddp, load into module
+        target = model.module if ddp else model
+        target.load_state_dict(model_state)
+        optimizer.load_state_dict(ckpt.get("optimizer_state_dict", {}))
+        scheduler_state = ckpt.get("scheduler_state_dict", None)
+        if scheduler_state:
+            scheduler.load_state_dict(scheduler_state)
+        if "scaler_state_dict" in ckpt and scaler is not None:
+            scaler.load_state_dict(ckpt["scaler_state_dict"])
+        start_step = ckpt.get("step", 0)
+        best_val_loss = ckpt.get("best_val_loss", best_val_loss)
+        print(f"Resumed from {resume_from} at step {start_step}")
+
+    model.train()
+    step = start_step
+    micro = 0
+    running_loss = 0.0
+    t0 = time.time()
+    no_improve_steps = 0
+    early_stop_patience = 10_000  # you can tune this
+
+    # training loop
+    while step < max_steps:
+        if sampler is not None:
+            sampler.set_epoch(step)  # shuffle differently per epoch for DDP
+
+        for batch in dl:
+            x, y = batch
+            x = x.to(device, non_blocking=True)
+            y = y.to(device, non_blocking=True)
+
+            with torch.cuda.amp.autocast(enabled=(device.type == "cuda")):
+                logits, loss = model(x, y)
+                loss = loss / grad_accum
+
+            scaler.scale(loss).backward()
+            micro += 1
+            running_loss += loss.item()
+
+            if micro % grad_accum == 0:
+                # gradient clipping
+                if clip_grad_norm is not None:
+                    scaler.unscale_(optimizer)
+                    torch.nn.utils.clip_grad_norm_(model.parameters(), clip_grad_norm)
+
+                scaler.step(optimizer)
+                scaler.update()
+                optimizer.zero_grad(set_to_none=True)
+                scheduler.step()
+
+                if ema:
+                    ema.update(model if not ddp else model.module)
+
+                step += 1
+
+                # logging
+                if step % 50 == 0 and (not ddp or local_rank == 0):
+                    grad_norm = compute_grad_norm(model if not ddp else model.module)
+                    avg_loss = running_loss * grad_accum / 50.0
+                    running_loss = 0.0
+                    elapsed = time.time() - t0
+                    lr_now = scheduler.get_last_lr()[0]
+                    print(f"step={step} loss={avg_loss:.6f} grad_norm={grad_norm:.3f} lr={lr_now:.6f} elapsed={elapsed:.1f}s")
+                    if writer:
+                        writer.add_scalar("train/loss", avg_loss, step)
+                        writer.add_scalar("train/grad_norm", grad_norm, step)
+                        writer.add_scalar("train/lr", lr_now, step)
+                    t0 = time.time()
+
+                # periodic validation
+                if validate_every and step % validate_every == 0:
+                    if val_dl is None:
+                        # quick in-memory val split: take first N batches (user should replace with real val)
+                        # NOTE: for production, create a dedicated validation dataset.
+                        val_ds = TokenChunkDataset(tok, sources[: max(1, len(sources) // 20)], seq_len=seq_len, eos_token_id=tok.eos_token_id)
+                        val_dl = DataLoader(val_ds, batch_size=batch_size, shuffle=False, num_workers=0, pin_memory=True, drop_last=False)
+
+                    model.eval()
+                    # optionally swap in EMA weights for evaluation
+                    if ema:
+                        ema.store(model if not ddp else model.module)
+                        ema.copy_to(model if not ddp else model.module)
+
+                    val_losses = []
+                    with torch.no_grad():
+                        for i, (vx, vy) in enumerate(val_dl):
+                            if i >= val_steps:
+                                break
+                            vx = vx.to(device)
+                            vy = vy.to(device)
+                            with torch.cuda.amp.autocast(enabled=(device.type == "cuda")):
+                                _, vloss = model(vx, vy)
+                            val_losses.append(float(vloss.detach().cpu().item()))
+                    mean_val = float(sum(val_losses) / max(1, len(val_losses)))
+                    if writer and (not ddp or local_rank == 0):
+                        writer.add_scalar("val/loss", mean_val, step)
+                    print(f"[eval] step={step} val_loss={mean_val:.6f}")
+
+                    # restore weights
+                    if ema:
+                        ema.restore(model if not ddp else model.module)
+                    model.train()
+
+                    # early stop / best model saving
+                    if mean_val < best_val_loss:
+                        best_val_loss = mean_val
+                        no_improve_steps = 0
+                        best_path = os.path.join(out_dir, f"supernova_best_step{step}.pt")
+                        ckpt = {
+                            "model_state_dict": (model.module.state_dict() if ddp else model.state_dict()),
+                            "optimizer_state_dict": optimizer.state_dict(),
+                            "scheduler_state_dict": scheduler.state_dict(),
+                            "scaler_state_dict": (scaler.state_dict() if scaler else None),
+                            "step": step,
+                            "best_val_loss": best_val_loss,
+                            "config": cfg.__dict__,
+                        }
+                        if not ddp or local_rank == 0:
+                            atomic_save(ckpt, best_path)
+                            print(f"Saved best checkpoint to {best_path}")
+                    else:
+                        no_improve_steps += validate_every
+                        if no_improve_steps >= early_stop_patience:
+                            print("Early stopping triggered.")
+                            step = max_steps
+                            break
+
+                # periodic checkpointing
+                if save_every and step % save_every == 0 and (not ddp or local_rank == 0):
+                    ckpt_path = os.path.join(out_dir, f"supernova_step{step}.pt")
+                    ckpt = {
+                        "model_state_dict": (model.module.state_dict() if ddp else model.state_dict()),
+                        "optimizer_state_dict": optimizer.state_dict(),
+                        "scheduler_state_dict": scheduler.state_dict(),
+                        "scaler_state_dict": (scaler.state_dict() if scaler else None),
+                        "step": step,
+                        "best_val_loss": best_val_loss,
+                        "config": cfg.__dict__,
+                    }
+                    atomic_save(ckpt, ckpt_path)
+                    print(f"Saved checkpoint {ckpt_path}")
+
+                if step >= max_steps:
+                    break
+
+        if step >= max_steps:
+            break
+
+    # final save
+    if not ddp or local_rank == 0:
+        ckpt_path = os.path.join(out_dir, f"supernova_final_step{step}.pt")
+        ckpt = {
+            "model_state_dict": (model.module.state_dict() if ddp else model.state_dict()),
+            "optimizer_state_dict": optimizer.state_dict(),
+            "scheduler_state_dict": scheduler.state_dict(),
+            "scaler_state_dict": (scaler.state_dict() if scaler else None),
+            "step": step,
+            "best_val_loss": best_val_loss,
+            "config": cfg.__dict__,
+        }
+        atomic_save(ckpt, ckpt_path)
+        print(f"Saved final checkpoint to {ckpt_path}")
+
+    if writer:
+        writer.close()
+
+
+if __name__ == "__main__":
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--config", required=True)
+    ap.add_argument("--data-config", required=True)
+    ap.add_argument("--seq-len", type=int, default=1024)
+    ap.add_argument("--batch-size", type=int, default=16)
+    ap.add_argument("--grad-accum", type=int, default=8)
+    ap.add_argument("--lr", type=float, default=3e-4)
+    ap.add_argument("--warmup-steps", type=int, default=2000)
+    ap.add_argument("--max-steps", type=int, default=100000)
+    ap.add_argument("--save-every", type=int, default=10000)
+    ap.add_argument("--out-dir", type=str, default="checkpoints")
+    ap.add_argument("--seed", type=int, default=42)
+    ap.add_argument("--validate-every", type=int, default=1000)
+    ap.add_argument("--val-steps", type=int, default=100)
+    ap.add_argument("--clip-grad-norm", type=float, default=1.0)
+    ap.add_argument("--resume-from", type=str, default=None)
+    ap.add_argument("--use-ema", action="store_true")
+    ap.add_argument("--ema-decay", type=float, default=0.9999)
+    ap.add_argument("--no-tensorboard", dest="use_tensorboard", action="store_false")
+    ap.add_argument("--ddp", action="store_true", help="enable DistributedDataParallel; use with torchrun")
+    ap.add_argument("--local-rank", type=int, default=0)
+    ap.add_argument("--num-workers", type=int, default=4)
+    ap.add_argument("--pin-memory", type=bool, default=True)
+    ap.add_argument("--compile", dest="compile_model", action="store_true")
+    args = ap.parse_args()
+
+    train(
+        config_path=args.config,
+        data_config_path=args.data_config,
+        seq_len=args.seq_len,
+        batch_size=args.batch_size,
+        grad_accum=args.grad_accum,
+        lr=args.lr,
+        warmup_steps=args.warmup_steps,
+        max_steps=args.max_steps,
+        save_every=args.save_every,
+        out_dir=args.out_dir,
+        seed=args.seed,
+        validate_every=args.validate_every,
+        val_steps=args.val_steps,
+        clip_grad_norm=args.clip_grad_norm,
+        use_ema=args.use_ema,
+        ema_decay=args.ema_decay,
+        resume_from=args.resume_from,
+        use_tensorboard=args.use_tensorboard,
+        ddp=args.ddp,
+        local_rank=args.local_rank,
+        num_workers=args.num_workers,
+        pin_memory=args.pin_memory,
+        compile_model=args.compile_model,
+    )