train.py

import torch
from torch.utils.data import DataLoader
from data.dataset import TinyStoriesDataset
from model.bigram import Bigram
from model.transformer import Transformer, TINY_CONFIG, FULL_CONFIG

if __name__ == "__main__":
    device = torch.device("cuda")

    training_data = TinyStoriesDataset("data/train.bin", context_len=512)
    valid_data = TinyStoriesDataset("data/val.bin", context_len=512)

    model = Transformer(**FULL_CONFIG) #GPT2 encoder vocab size is 50257
    model.to(device)

    n_params = sum(p.numel() for p in model.parameters())
    print(f"Model parameters: {n_params:,}")

    loss_function = torch.nn.CrossEntropyLoss()

    optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4)

    training_loader = DataLoader(
        training_data,
        batch_size=4,
        shuffle=True,
        num_workers=0,
    )

    valid_loader = DataLoader(
        valid_data,
        batch_size=4,
        shuffle=True,
        num_workers=0,
    )

    epoch = 30

    for e in range(epoch):
        for i, data in enumerate(training_loader):
            x, y = data
            x = x.to(device)
            y = y.to(device)

            outputs = model(x)

            optimizer.zero_grad()

            loss = loss_function(outputs.view(-1, 50257), y.view(-1))
            loss.backward()

            optimizer.step()

            if i % 100 == 0:
                print(f"Epoch {e}, Step {i}, Loss: {loss.item():.4f}")
    
    torch.save({
        'model_state_dict': model.state_dict(),
        'optimizer_state_dict': optimizer.state_dict(),
        'loss': loss.item(),
    }, 'checkpoints/transformer_full.pt')
    print("Model saved!")