train_local.py

#!/usr/bin/env python3
"""Local GPU training — tiny model for proof of concept.

Downloads TinyStories, tokenizes, trains a ~16M param Hybrid Liquid-Dense model.
Designed for GTX 1650 (4GB VRAM). Produces coherent English in ~1-2 hours.
"""

import argparse
import math
import os
import sys
import time
from pathlib import Path

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

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

from hybrid_liquid_dense.config import HybridLiquidDenseConfig
from hybrid_liquid_dense.model import HybridLiquidDenseModel


# ================================================================
# Data: Download + Tokenize TinyStories
# ================================================================


def prepare_data(data_path: str, max_tokens: int = 100_000_000):
    """Download TinyStories and tokenize to uint16 binary file."""
    if Path(data_path).exists():
        data = np.memmap(data_path, dtype=np.uint16, mode="r")
        print(f"Data exists: {len(data)/1e6:.1f}M tokens at {data_path}")
        return

    print("Downloading TinyStories (streaming)...")
    import tiktoken
    from datasets import load_dataset

    enc = tiktoken.get_encoding("gpt2")
    eot = enc.eot_token

    ds = load_dataset("roneneldan/TinyStories", split="train", streaming=True)

    tokens = []
    t0 = time.time()
    for i, ex in enumerate(ds):
        text = ex.get("text", "")
        if text:
            tokens.extend(enc.encode(text, allowed_special={"<|endoftext|>"}))
            tokens.append(eot)

        if len(tokens) >= max_tokens:
            break

        if (i + 1) % 10000 == 0:
            elapsed = time.time() - t0
            print(f"  {len(tokens)/1e6:.1f}M tokens from {i+1} docs ({elapsed:.0f}s)")

    tokens = tokens[:max_tokens]
    arr = np.array(tokens, dtype=np.uint16)

    Path(data_path).parent.mkdir(parents=True, exist_ok=True)
    arr.tofile(data_path)
    print(f"Saved {len(arr)/1e6:.1f}M tokens to {data_path} ({elapsed:.0f}s)")


# ================================================================
# Training
# ================================================================


def cosine_lr(step, warmup, total, peak, minimum=0.0):
    if step < warmup:
        return peak * step / max(warmup, 1)
    if step >= total:
        return minimum
    progress = (step - warmup) / max(total - warmup, 1)
    return minimum + 0.5 * (peak - minimum) * (1.0 + math.cos(math.pi * progress))


def train(args):
    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:.1f}GB")

    # ---- Small config for local training ----
    config = HybridLiquidDenseConfig(
        vocab_size=50_257,
        d_model=args.d_model,
        d_ff=args.d_ff,
        n_layers=args.n_layers,
        max_seq_len=args.seq_len,
        conv_kernel=4,
        scan_backend="pytorch",  # No Triton on GTX 1650
    )

    model = HybridLiquidDenseModel(config).to(device)
    n_params = model.num_parameters()
    print(f"Model: {n_params/1e6:.1f}M parameters")
    print(f"Config: d={config.d_model}, ff={config.d_ff}, layers={config.n_layers}, seq={config.max_seq_len}")

    # ---- Data ----
    data = np.memmap(args.data_path, dtype=np.uint16, mode="r")
    n_tokens = len(data)
    stride = args.seq_len + 1
    n_samples = n_tokens // stride
    print(f"Data: {n_tokens/1e6:.1f}M tokens, {n_samples:,} samples")

    # ---- Optimizer ----
    optimizer = torch.optim.AdamW(
        model.parameters(),
        lr=args.peak_lr,
        betas=(0.9, 0.95),
        weight_decay=0.1,
        eps=1e-8,
    )

    total_steps = int(args.total_tokens / (args.batch_size * args.seq_len))
    print(f"Training: {args.total_tokens/1e6:.0f}M tokens, {total_steps:,} steps")
    print(f"Batch: {args.batch_size}, seq: {args.seq_len}")
    print()

    # ---- Train ----
    model.train()
    rng = np.random.RandomState(42)
    global_step = 0
    t0 = time.time()
    total_loss = 0.0
    total_tok = 0
    best_loss = float("inf")

    for step in range(total_steps):
        # Random sample from data
        starts = rng.randint(0, n_tokens - stride, size=args.batch_size)
        batch = np.stack([data[s : s + stride].astype(np.int64) for s in starts])
        batch = torch.from_numpy(batch).to(device)
        input_ids = batch[:, :-1]
        targets = batch[:, 1:]

        # LR schedule
        lr = cosine_lr(step, args.warmup_steps, total_steps, args.peak_lr, args.min_lr)
        for pg in optimizer.param_groups:
            pg["lr"] = lr

        # Forward
        with torch.amp.autocast("cuda", dtype=torch.float16, enabled=device.type == "cuda"):
            logits, _ = model(input_ids)
            loss = F.cross_entropy(
                logits.reshape(-1, config.vocab_size), targets.reshape(-1)
            )

        # Backward
        optimizer.zero_grad(set_to_none=True)
        loss.backward()
        grad_norm = nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
        optimizer.step()

        n_tok = targets.numel()
        total_loss += loss.item() * n_tok
        total_tok += n_tok
        global_step += 1

        # Log
        if step % args.log_interval == 0 or step == total_steps - 1:
            avg = total_loss / total_tok
            ppl = math.exp(min(avg, 20))
            elapsed = time.time() - t0
            tok_s = total_tok / max(elapsed, 1)
            eta = (total_steps - step) / max(step / elapsed, 1e-9)
            print(
                f"step {step:>6d}/{total_steps} | loss {loss.item():.3f} | "
                f"avg {avg:.3f} | ppl {ppl:.1f} | lr {lr:.1e} | "
                f"gnorm {grad_norm:.2f} | {tok_s:.0f} tok/s | "
                f"ETA {eta/60:.0f}m"
            )

        # Save checkpoint periodically
        if step > 0 and step % args.save_interval == 0:
            _save(model, config, step, args.checkpoint_dir)

        # Generate sample every N steps
        if step > 0 and step % args.sample_interval == 0:
            _generate_sample(model, config, device)

    elapsed = time.time() - t0
    avg_loss = total_loss / total_tok
    print(f"\nDone: {total_steps:,} steps, {elapsed/60:.1f}min, {total_tok/elapsed:.0f} tok/s")
    print(f"Final loss: {avg_loss:.3f}, ppl: {math.exp(min(avg_loss, 20)):.1f}")

    # Final save
    _save(model, config, total_steps, args.checkpoint_dir)
    return model, config


def _save(model, config, step, checkpoint_dir):
    Path(checkpoint_dir).mkdir(parents=True, exist_ok=True)
    path = Path(checkpoint_dir) / f"step_{step}.pt"
    torch.save({"model": model.state_dict(), "config": vars(config), "step": step}, path)
    print(f"  Saved: {path} ({path.stat().st_size/1e6:.1f}MB)")


def _generate_sample(model, config, device):
    """Generate a quick sample to see progress."""
    import tiktoken
    enc = tiktoken.get_encoding("gpt2")

    model.eval()
    prompt = "Once upon a time"
    ids = enc.encode(prompt)
    input_ids = torch.tensor([ids], dtype=torch.long, device=device)

    with torch.no_grad():
        generated = model.generate(input_ids, max_new_tokens=60, temperature=0.8, top_k=40)

    text = enc.decode(generated[0].tolist())
    print(f"  Sample: \"{prompt}{text}\"")
    print()
    model.train()


# ================================================================
# CLI
# ================================================================


def main():
    p = argparse.ArgumentParser(description="Local training — proof of concept")

    # Data
    p.add_argument("--data_path", default="data/tinystories.bin")
    p.add_argument("--max_download_tokens", type=float, default=50e6,
                    help="How many tokens to download (default 50M)")
    p.add_argument("--total_tokens", type=float, default=50e6,
                    help="How many tokens to train on (default 50M)")

    # Model (tiny defaults for GTX 1650 — ~8M params, fits in 4GB)
    p.add_argument("--d_model", type=int, default=192)
    p.add_argument("--d_ff", type=int, default=576)
    p.add_argument("--n_layers", type=int, default=4)
    p.add_argument("--seq_len", type=int, default=256)

    # Training
    p.add_argument("--batch_size", type=int, default=16)
    p.add_argument("--peak_lr", type=float, default=1e-3)
    p.add_argument("--min_lr", type=float, default=1e-4)
    p.add_argument("--warmup_steps", type=int, default=300)
    p.add_argument("--grad_clip", type=float, default=1.0)

    # Logging
    p.add_argument("--log_interval", type=int, default=100)
    p.add_argument("--save_interval", type=int, default=5000)
    p.add_argument("--sample_interval", type=int, default=2000)
    p.add_argument("--checkpoint_dir", default="checkpoints_local")

    args = p.parse_args()

    # Step 1: Prepare data
    prepare_data(args.data_path, int(args.max_download_tokens))

    # Step 2: Train
    model, config = train(args)

    # Step 3: Interactive chat
    print("\n" + "=" * 50)
    print("Training complete! Starting chat...")
    print("=" * 50 + "\n")

    from chat import interactive_chat
    interactive_chat(
        checkpoint=f"{args.checkpoint_dir}/step_{int(args.total_tokens / (args.batch_size * args.seq_len))}.pt",
        device="cuda" if torch.cuda.is_available() else "cpu",
    )


if __name__ == "__main__":
    main()