train.py

import torch

from config.model_config import MODEL_CONFIG
from model.lexora import LexoraGPT
from training.dataloader import create_dataloaders
from training.evaluate import estimate_loss
from training.loss import calculate_loss
from training.utils import get_device


# -----------------------------
# Device
# -----------------------------

device = get_device()

print("Using device:", device)


# -----------------------------
# Load dataset
# -----------------------------

with open(
    "data/raw/input.txt",
    "r",
    encoding="utf-8"
) as f:
    text = f.read()


tokenizer, train_loader, val_loader = create_dataloaders(
    text=text,
    block_size=MODEL_CONFIG["context_length"],
    batch_size=MODEL_CONFIG["batch_size"]
)


# -----------------------------
# Create model
# -----------------------------

model = LexoraGPT(
    vocab_size=tokenizer.vocab_size,
    embedding_dim=MODEL_CONFIG["embedding_dim"],
    num_heads=MODEL_CONFIG["num_heads"],
    num_layers=MODEL_CONFIG["num_layers"],
    context_length=MODEL_CONFIG["context_length"],
    dropout=MODEL_CONFIG["dropout"]
)

model = model.to(device)


print(
    "Parameters:",
    sum(
        p.numel()
        for p in model.parameters()
        if p.requires_grad
    )
)


# -----------------------------
# Optimizer
# -----------------------------

optimizer = torch.optim.AdamW(
    model.parameters(),
    lr=MODEL_CONFIG["learning_rate"],
    weight_decay=MODEL_CONFIG["weight_decay"]
)


# -----------------------------
# Training
# -----------------------------

best_val_loss = float("inf")

train_iterator = iter(train_loader)

for step in range(
    MODEL_CONFIG["max_steps"]
):

    try:
        x, y = next(train_iterator)

    except StopIteration:
        train_iterator = iter(train_loader)
        x, y = next(train_iterator)

    x = x.to(device)
    y = y.to(device)

    # Forward pass
    logits = model(x)

    # Calculate loss
    loss = calculate_loss(
        logits,
        y
    )

    # Backpropagation
    optimizer.zero_grad(
        set_to_none=True
    )

    loss.backward()

    optimizer.step()

    # Evaluation
    
    if (
        step % MODEL_CONFIG["eval_interval"] == 0
        or step == MODEL_CONFIG["max_steps"] - 1
    ):

        losses = estimate_loss(
            model,
            train_loader,
            val_loader,
            device,
            MODEL_CONFIG["eval_batches"]
        )

        print(
            f"Step {step:5d} | "
            f"Train loss: {losses['train']:.4f} | "
            f"Val loss: {losses['val']:.4f}"
        )

        # Save best model
        if losses["val"] < best_val_loss:

            best_val_loss = losses["val"]

            torch.save(
                {
                    "step": step,
                    "model_state_dict": model.state_dict(),
                    "optimizer_state_dict": optimizer.state_dict(),
                    "train_loss": losses["train"],
                    "val_loss": losses["val"],
                    "config": MODEL_CONFIG,
                },
                "checkpoints/best_model.pt"
            )

            print(
                f"✓ Saved new best model "
                f"(val loss: {best_val_loss:.4f})"
            )