train.py

import math
import model as m
import sentencepiece as spm
import numpy as np
import argparse
import torch
import torch.nn.functional as F

CONTEXT_LENGTH = 1024
DIM = 480
NUM_HEADS = 8
EXPANSION = 4
NUM_LAYERS = 6

tokenizer = spm.SentencePieceProcessor(model_file="tokenizer.model")
vocab_size = tokenizer.get_piece_size()

def pretokenize_dataset(dataset_name):
    all_ids = []
    with open(f"{dataset_name}.txt", "r", encoding="utf-8") as f:
        for line in f:
            text = line.strip()
            if text:
                all_ids.extend(tokenizer.encode(text))
                all_ids.append(tokenizer.eos_id())
                
    arr = np.array(all_ids, dtype=np.uint16) #binary file of uint16 ints
    arr.tofile(f"{dataset_name}_tokenized.bin")
    print(f"Saved {len(arr):,} tokens")
    return arr

def load_dataset(dataset_name, tokenize):
    if tokenize:
        arr = pretokenize_dataset(dataset_name)
    else:
        arr = np.fromfile(f"{dataset_name}_tokenized.bin", dtype=np.uint16)
    return torch.from_numpy(arr.astype(np.int64))

def get_batch(data, batch_size, device):
    starts = torch.randint(0, len(data) - CONTEXT_LENGTH - 1, (batch_size,))
    x = torch.stack([data[i   : i + CONTEXT_LENGTH] for i in starts]).to(device)
    y = torch.stack([data[i+1 : i + CONTEXT_LENGTH + 1] for i in starts]).to(device)
    return x, y

@torch.no_grad()
def estimate_val_loss(model, val_data, batch_size, device, dtype, eval_iters=50):
    model.eval()
    losses = []
    for _ in range(eval_iters):
        x, y = get_batch(val_data, batch_size, device)
        with torch.autocast(device_type=device, dtype=dtype):
            logits = model(x)
            loss = F.cross_entropy(logits.view(-1, vocab_size), y.view(-1))
        losses.append(loss.item())
    model.train()
    return sum(losses) / len(losses)

def get_lr(step, learning_rate, warmup_iters, total_iters):
    if step < warmup_iters:
        return learning_rate * step / warmup_iters
    progress = (step - warmup_iters) / (total_iters - warmup_iters)
    return learning_rate * (0.1 + 0.9 * 0.5 * (1 + math.cos(math.pi * progress)))

def main(args):
    torch.manual_seed(args.seed)
    device = "cuda" if torch.cuda.is_available() else "cpu"
    dtype = {"float32": torch.float32, "bfloat16": torch.bfloat16, "float16": torch.float16}[args.dtype]
    print(f"Using device: {device} | dtype: {args.dtype}")

    data = load_dataset(args.dataset_name, args.tokenize)
    split = int(0.9 * len(data))
    train_data, val_data = data[:split], data[split:]
    print(f"Train tokens: {len(train_data):,} | Val tokens: {len(val_data):,}")

    model = m.GPT(
        vocab_size=vocab_size,
        dim=DIM,
        num_heads=NUM_HEADS,
        expansion=EXPANSION,
        num_layers=NUM_LAYERS,
        context_length=CONTEXT_LENGTH,
        dropout=args.dropout,
    ).to(device=device, dtype=dtype)
    print(f"Parameters: {sum(p.numel() for p in model.parameters()):,}")

    model = torch.compile(model)
    optimizer = torch.optim.AdamW(model.parameters(), lr=args.learning_rate, weight_decay=args.weight_decay)

    model.train()
    optimizer.zero_grad()
    for i in range(args.iters):
        lr = args.learning_rate if args.no_decay_lr else get_lr(i, args.learning_rate, args.warmup_iters, args.iters)
        for param_group in optimizer.param_groups:
            param_group["lr"] = lr

        for micro_step in range(args.grad_accum_steps):
            x, y = get_batch(train_data, args.batch_size, device)
            with torch.autocast(device_type=device, dtype=dtype):
                logits = model(x)
                loss = F.cross_entropy(logits.view(-1, vocab_size), y.view(-1))
                loss = loss / args.grad_accum_steps  # scale loss to average over accumulation steps
            loss.backward()

        torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
        optimizer.step()
        optimizer.zero_grad()

        if i % 50 == 0:
            val_loss = estimate_val_loss(model, val_data, args.batch_size, device, dtype)
            print(f"iter {i:6d} | train loss {loss.item() * args.grad_accum_steps:.4f} | val loss {val_loss:.4f} | lr {lr:.2e}")

        if i % 5000 == 0 and i > 0:
            torch.save(model.state_dict(), f"{args.model_dir}_step{i}.pt")
            print(f"Checkpoint saved at step {i}")

    torch.save(model.state_dict(), f"{args.model_dir}.pt")
    print(f"Model saved to {args.model_dir}.pt")

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--model_dir", type=str, default="model")
    parser.add_argument("--tokenize", action="store_true")
    parser.add_argument("--dataset_name", type=str, default="corpus")
    parser.add_argument("--iters", type=int, default=40000)
    parser.add_argument("--batch_size", type=int, default=12)
    parser.add_argument("--learning_rate", type=float, default=1e-3)
    parser.add_argument("--weight_decay", type=float, default=0.1)
    parser.add_argument("--dropout", type=float, default=0.1)
    parser.add_argument("--grad_clip", type=float, default=1.0)
    parser.add_argument("--grad_accum_steps", type=int, default=2)
    parser.add_argument("--no_decay_lr", action="store_true")
    parser.add_argument("--warmup_iters", type=int, default=400)
    parser.add_argument("--dtype", type=str, default="bfloat16", choices=["float32", "bfloat16", "float16"])
    parser.add_argument("--seed", type=int, default=42)
    args = parser.parse_args()
    main(args)