train_gpu.py

"""
train_gpu.py — AngstromE1-Nano GPU Training
Optimized for 2x NVIDIA T4 (16GB VRAM each)

Features:
  - Mixed precision (FP16/BF16)
  - Gradient checkpointing
  - DataParallel across 2 GPUs
  - Cosine LR schedule with warmup
  - Gradient clipping
  - Periodic checkpointing
  - Wandb logging (optional)

Usage:
  python train_gpu.py                         # default config
  python train_gpu.py --config large          # larger model
  python train_gpu.py --resume checkpoint.pt  # resume training
"""
import sys; sys.path.insert(0, '.')
import os
import math
import time
import json
import argparse
from pathlib import Path

import torch
import torch.nn as nn
from torch.cuda.amp import GradScaler, autocast
from torch.nn.parallel import DistributedDataParallel as DDP

from angstrom_nano import AngstromNanoConfig, AngstromNanoForCausalLM
from angstrom_nano.tokenizer import AngstromNanoTokenizer


# ═══════════════════════════════════════════════════════════════════
# Configs
# ═══════════════════════════════════════════════════════════════════

CONFIGS = {
    "small": {
        "vocab_size": 8192, "hidden_size": 256, "intermediate_size": 1024,
        "num_hidden_layers": 8, "num_attention_heads": 8, "num_key_value_heads": 2,
        "head_dim": 32, "num_local_experts": 4, "num_experts_per_tok": 2,
        "max_position_embeddings": 2048, "sliding_window": 512,
        "scoring_func": "sigmoid", "use_qk_norm": True, "use_routing_bias": True,
        "tie_word_embeddings": True,
    },
    "medium": {
        "vocab_size": 16384, "hidden_size": 512, "intermediate_size": 2048,
        "num_hidden_layers": 12, "num_attention_heads": 16, "num_key_value_heads": 4,
        "head_dim": 32, "num_local_experts": 8, "num_experts_per_tok": 2,
        "max_position_embeddings": 4096, "sliding_window": 1024,
        "scoring_func": "sigmoid", "use_qk_norm": True, "use_routing_bias": True,
        "tie_word_embeddings": True,
    },
    "large": {
        "vocab_size": 32000, "hidden_size": 1024, "intermediate_size": 4096,
        "num_hidden_layers": 24, "num_attention_heads": 16, "num_key_value_heads": 4,
        "head_dim": 64, "num_local_experts": 8, "num_experts_per_tok": 2,
        "max_position_embeddings": 4096, "sliding_window": 1024,
        "scoring_func": "sigmoid", "use_qk_norm": True, "use_routing_bias": True,
        "tie_word_embeddings": True,
    },
}


# ═══════════════════════════════════════════════════════════════════
# Dataset
# ═══════════════════════════════════════════════════════════════════

class TextDataset(torch.utils.data.Dataset):
    """Memory-mapped token dataset for large corpora."""

    def __init__(self, token_ids: torch.Tensor, seq_len: int):
        self.token_ids = token_ids
        self.seq_len = seq_len
        self.n_samples = len(token_ids) - seq_len - 1

    def __len__(self):
        return self.n_samples

    def __getitem__(self, idx):
        x = self.token_ids[idx : idx + self.seq_len]
        y = self.token_ids[idx + 1 : idx + self.seq_len + 1]
        return x, y


# ═══════════════════════════════════════════════════════════════════
# Training
# ═══════════════════════════════════════════════════════════════════

def setup_device():
    """Setup multi-GPU or single GPU."""
    if not torch.cuda.is_available():
        print("WARNING: No GPU found, using CPU (will be slow)")
        return torch.device("cpu"), 1

    n_gpus = torch.cuda.device_count()
    device = torch.device("cuda:0")
    print(f"Using {n_gpus} GPU(s):")
    for i in range(n_gpus):
        props = torch.cuda.get_device_properties(i)
        print(f"  GPU {i}: {props.name} ({props.total_mem / 1e9:.1f} GB)")
    return device, n_gpus


def get_lr(step, warmup_steps, max_steps, base_lr, min_lr):
    """Cosine learning rate schedule with warmup."""
    if step < warmup_steps:
        return base_lr * step / max(1, warmup_steps)
    progress = (step - warmup_steps) / max(1, max_steps - warmup_steps)
    return min_lr + 0.5 * (base_lr - min_lr) * (1.0 + math.cos(math.pi * progress))


def count_params(model):
    """Count trainable parameters."""
    return sum(p.numel() for p in model.parameters() if p.requires_grad)


def save_checkpoint(model, optimizer, scaler, step, loss, config_dict, out_dir):
    """Save training checkpoint."""
    out_dir = Path(out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)

    # Save model weights
    sd = model.module.state_dict() if hasattr(model, 'module') else model.state_dict()
    if "lm_head.weight" not in sd:
        sd["lm_head.weight"] = sd["model.embed_tokens.weight"]

    from safetensors.torch import save_file
    weights_path = out_dir / f"checkpoint-{step}.safetensors"
    save_file({k: v.contiguous().cpu() for k, v in sd.items()}, str(weights_path))

    # Save training state
    state = {
        "step": step,
        "loss": loss,
        "config": config_dict,
        "optimizer": optimizer.state_dict(),
        "scaler": scaler.state_dict() if scaler else None,
    }
    torch.save(state, str(out_dir / f"checkpoint-{step}.pt"))

    # Save config
    (out_dir / "config.json").write_text(json.dumps(config_dict, indent=2))

    print(f"  Saved checkpoint at step {step}")


def train(args):
    """Main training loop."""
    torch.manual_seed(42)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(42)

    # ── Setup ──────────────────────────────────────────────────────
    device, n_gpus = setup_device()
    use_amp = device.type == "cuda" and torch.cuda.is_available()

    # ── Tokenizer ──────────────────────────────────────────────────
    tok_path = Path("checkpoints/tokenizer.json")
    if tok_path.exists():
        tok = AngstromNanoTokenizer.from_bpe_file(str(tok_path))
        print(f"Loaded tokenizer: {len(tok)} vocab")
    else:
        print("Training new tokenizer...")
        tok = AngstromNanoTokenizer.train_bpe(
            [str(args.data_path)], vocab_size=args.vocab_size,
            out_path=str(tok_path),
        )
        print(f"Trained tokenizer: {len(tok)} vocab")

    # ── Load and tokenize data ─────────────────────────────────────
    print(f"\nLoading data from {args.data_path}...")
    text = Path(args.data_path).read_text(encoding="utf-8")
    print(f"  Raw: {len(text):,} chars ({len(text)/1e6:.1f} MB)")

    ids = torch.tensor(tok.encode(text, add_bos=True, add_eos=True), dtype=torch.long)
    print(f"  Tokens: {len(ids):,} ({len(ids)/1e6:.1f}M)")

    dataset = TextDataset(ids, args.seq_len)
    print(f"  Samples: {len(dataset):,}")

    dataloader = torch.utils.data.DataLoader(
        dataset, batch_size=args.batch_size, shuffle=True,
        num_workers=2, pin_memory=True, drop_last=True,
    )

    # ── Model ──────────────────────────────────────────────────────
    config_dict = CONFIGS[args.config]
    config_dict["vocab_size"] = len(tok)
    cfg = AngstromNanoConfig(**config_dict)
    model = AngstromNanoForCausalLM(cfg)

    n_params = count_params(model)
    print(f"\nModel: {n_params:,} params ({n_params * 4 / 1e6:.1f} MB FP32)")

    model = model.to(device)

    # Multi-GPU
    if n_gpus > 1:
        model = nn.DataParallel(model, device_ids=list(range(n_gpus)))
        print(f"  Wrapped in DataParallel across {n_gpus} GPUs")

    # ── Optimizer ──────────────────────────────────────────────────
    optimizer = torch.optim.AdamW(
        model.parameters(), lr=args.lr,
        weight_decay=0.1, betas=(0.9, 0.95),
    )
    scaler = GradScaler(enabled=use_amp)

    # ── Resume ─────────────────────────────────────────────────────
    start_step = 0
    if args.resume and Path(args.resume).exists():
        print(f"\nResuming from {args.resume}...")
        ckpt = torch.load(args.resume, map_location=device)
        if hasattr(model, 'module'):
            model.module.load_state_dict(ckpt["model"])
        else:
            model.load_state_dict(ckpt["model"])
        optimizer.load_state_dict(ckpt["optimizer"])
        start_step = ckpt["step"]
        print(f"  Resumed at step {start_step}")

    # ── Training ───────────────────────────────────────────────────
    max_steps = args.steps
    warmup_steps = args.warmup_steps
    grad_clip = args.grad_clip
    log_every = args.log_every
    save_every = args.save_every

    print(f"\nTraining for {max_steps} steps, seq_len={args.seq_len}, batch={args.batch_size}")
    print(f"  LR: {args.lr} → {args.min_lr}, warmup: {warmup_steps} steps")
    print(f"  Gradient clipping: {grad_clip}")
    print(f"  Mixed precision: {use_amp}")
    print(f"  Checkpoint every: {save_every} steps")
    print()

    model.train()
    t0 = time.time()
    running_loss = 0.0
    running_steps = 0

    for step in range(start_step + 1, max_steps + 1):
        # Get batch
        try:
            x, y = next(dataloader_iter)
        except (StopIteration, NameError):
            dataloader_iter = iter(dataloader)
            x, y = next(dataloader_iter)

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

        # Forward pass with AMP
        with autocast(enabled=use_amp, dtype=torch.float16):
            out = model(x, labels=y, output_router_logits=True)
            loss = out["loss"]
            aux_loss = out["aux_loss"]

        # Backward pass
        optimizer.zero_grad()
        scaler.scale(loss).backward()

        # Gradient clipping
        if grad_clip > 0:
            scaler.unscale_(optimizer)
            nn.utils.clip_grad_norm_(model.parameters(), grad_clip)

        scaler.step(optimizer)
        scaler.update()

        # Update LR
        lr = get_lr(step, warmup_steps, max_steps, args.lr, args.min_lr)
        for param_group in optimizer.param_groups:
            param_group["lr"] = lr

        # Logging
        running_loss += loss.item()
        running_steps += 1

        if step % log_every == 0 or step == 1:
            avg_loss = running_loss / running_steps
            ppl = math.exp(min(avg_loss, 20))  # cap to avoid overflow
            elapsed = time.time() - t0
            tokens_per_sec = (args.batch_size * args.seq_len * running_steps) / elapsed
            gpu_mem = torch.cuda.memory_allocated(0) / 1e9 if device.type == "cuda" else 0

            print(f"  step {step:>6d}/{max_steps}  loss={avg_loss:.4f}  ppl={ppl:.2f}  "
                  f"aux={aux_loss.item():.6f}  lr={lr:.1e}  "
                  f"tok/s={tokens_per_sec:.0f}  gpu={gpu_mem:.1f}GB  "
                  f"{elapsed:.0f}s")
            running_loss = 0.0
            running_steps = 0

        # Save checkpoint
        if step % save_every == 0:
            save_checkpoint(model, optimizer, scaler, step, avg_loss,
                          config_dict, args.output_dir)

    # ── Final save ─────────────────────────────────────────────────
    print("\nTraining complete!")
    save_checkpoint(model, optimizer, scaler, max_steps, avg_loss,
                    config_dict, args.output_dir)

    # Save final model as the main model file
    final_path = Path(args.output_dir) / "model_final.safetensors"
    sd = model.module.state_dict() if hasattr(model, 'module') else model.state_dict()
    if "lm_head.weight" not in sd:
        sd["lm_head.weight"] = sd["model.embed_tokens.weight"]
    from safetensors.torch import save_file
    save_file({k: v.contiguous().cpu() for k, v in sd.items()}, str(final_path))
    print(f"Saved final model: {final_path}")

    total_time = time.time() - t0
    print(f"Total time: {total_time/3600:.1f} hours")


# ═══════════════════════════════════════════════════════════════════
# CLI
# ═══════════════════════════════════════════════════════════════════

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="AngstromE1-Nano GPU Training")
    parser.add_argument("--config", default="medium", choices=["small", "medium", "large"],
                        help="Model config (default: medium)")
    parser.add_argument("--data-path", default="data/corpus.txt",
                        help="Path to training corpus")
    parser.add_argument("--output-dir", default="checkpoints",
                        help="Output directory for checkpoints")
    parser.add_argument("--vocab-size", type=int, default=16384,
                        help="BPE vocab size (if training tokenizer)")
    parser.add_argument("--seq-len", type=int, default=512,
                        help="Sequence length (default: 512)")
    parser.add_argument("--batch-size", type=int, default=4,
                        help="Batch size per GPU (default: 4)")
    parser.add_argument("--steps", type=int, default=50000,
                        help="Total training steps (default: 50000)")
    parser.add_argument("--lr", type=float, default=3e-3,
                        help="Peak learning rate (default: 3e-3)")
    parser.add_argument("--min-lr", type=float, default=3e-4,
                        help="Min learning rate (default: 3e-4)")
    parser.add_argument("--warmup-steps", type=int, default=500,
                        help="Warmup steps (default: 500)")
    parser.add_argument("--grad-clip", type=float, default=1.0,
                        help="Gradient clipping (default: 1.0)")
    parser.add_argument("--log-every", type=int, default=100,
                        help="Log every N steps (default: 100)")
    parser.add_argument("--save-every", type=int, default=5000,
                        help="Save checkpoint every N steps (default: 5000)")
    parser.add_argument("--resume", type=str, default=None,
                        help="Resume from checkpoint path")
    args = parser.parse_args()

    train(args)