train.py

from pathlib import Path
import logging
import torch
import torch.nn as nn
import torch.nn.functional as F
from time import time
from torch.utils.data import DataLoader
from transformers import get_cosine_schedule_with_warmup
from types import SimpleNamespace

from model import Model
from dataset import TextDataset, ResponseDataset

logger = logging.getLogger(__name__)

assert torch.cuda.is_available(), "CUDA GPU required for training."
device = "cuda"


def val(model, val_data, criterion):

    model.eval()
    total_loss = torch.tensor(0.0, device=device)
    for x, y in val_data:
        x, y = x.to(device), y.to(device)
        with torch.amp.autocast(device_type="cuda", dtype=torch.bfloat16):
            logits = model(x)
            total_loss += criterion(
                logits.view(-1, logits.size(-1)), y.view(-1)
            ).detach()
    model.train()
    return total_loss.item() / len(val_data)


def train(config=None, **kwargs):

    config_defaults = {
        "checkpoint_dir": "checkpoints",
        "context": 1024,
        "dataloader_workers": 4,
        "dataset_dir": "datasets",
        "dataset_file": "wikitext.npy",
        "dataset_type": "text",
        "epochs": 1,
        "grad_norm": 1,
        "grad_steps": 8,
        "microbatch_size": 4,
        "model_file": None,
        "logging_rate": 20,
        "lr": 6e-4,
        "pad_token": 50257,
        "val_dataset_file": "wikitext_val.npy",
        "vocab_size": 50260,  # Default GPT 2 tokenization + user + assistant + pad
        "warmup_batches": 200,
        "weight_decay": 0.1,
    }
    if config is None:
        config = kwargs
    config = SimpleNamespace(**(config_defaults | config))

    Path(config.checkpoint_dir).mkdir(parents=True, exist_ok=True)

    DatasetClass = TextDataset if config.dataset_type == "text" else ResponseDataset
    dataset = DatasetClass(
        Path(config.dataset_dir) / config.dataset_file, seq_len=config.context
    )
    dataloader = DataLoader(
        dataset,
        batch_size=config.microbatch_size,
        drop_last=True,
        shuffle=True,
        num_workers=config.dataloader_workers,
        pin_memory=True,
        persistent_workers=True,
    )
    val_dataset = DatasetClass(
        Path(config.dataset_dir) / config.val_dataset_file, seq_len=config.context
    )
    val_dataloader = DataLoader(
        val_dataset,
        batch_size=config.microbatch_size,
        drop_last=True,
        shuffle=True,
        num_workers=config.dataloader_workers,
        pin_memory=True,
        persistent_workers=True,
    )
    logger.info("Loaded dataset!")

    raw_model = Model(
        vocab_size=config.vocab_size,
        max_context=config.context,
    ).to(device)
    if config.model_file:
        state_dict = torch.load(Path(config.checkpoint_dir) / config.model_file)[
            "model_state"
        ]
        raw_model.load_state_dict(state_dict)
        logger.info("Got model from checkpoint!")
    model = torch.compile(raw_model)
    logger.info("Loaded model!")

    param_dict = {pn: p for pn, p in model.named_parameters() if p.requires_grad}
    decay_params = [p for n, p in param_dict.items() if p.dim() >= 2]
    nodecay_params = [p for n, p in param_dict.items() if p.dim() < 2]
    optim_groups = [
        {"params": decay_params, "weight_decay": config.weight_decay},
        {"params": nodecay_params, "weight_decay": 0.0},
    ]
    optimizer = torch.optim.AdamW(
        optim_groups, lr=config.lr, betas=(0.9, 0.95), eps=1e-8
    )
    criterion = nn.CrossEntropyLoss(ignore_index=config.pad_token)

    scheduler = get_cosine_schedule_with_warmup(
        optimizer=optimizer,
        num_warmup_steps=config.warmup_batches,
        num_training_steps=config.epochs * (len(dataloader) // config.grad_steps),
    )

    model.train()

    starting_time = time()
    for epoch in range(config.epochs):
        optimizer.zero_grad(set_to_none=True)
        accum_loss = torch.tensor(0.0, device=device)
        for step, (x, y) in enumerate(dataloader):
            x, y = x.to(device, non_blocking=True), y.to(device, non_blocking=True)
            with torch.amp.autocast(device_type="cuda", dtype=torch.bfloat16):
                logits = model(x)
                loss = criterion(logits.view(-1, logits.size(-1)), y.view(-1))
                loss = loss / config.grad_steps
            accum_loss += loss.detach()
            loss.backward()

            if step % config.grad_steps == config.grad_steps - 1:
                norm = nn.utils.clip_grad_norm_(
                    model.parameters(), max_norm=config.grad_norm
                )
                optimizer.step()
                optimizer.zero_grad(set_to_none=True)
                scheduler.step()

            if (
                step % (config.logging_rate * config.grad_steps)
                == config.logging_rate * config.grad_steps - 1
            ):
                logger.info(
                    f"Step: {step + 1} | Loss per token: {accum_loss.item() / config.logging_rate} | Avg time / step: {((time() - starting_time) / (step + 1)):.2f}"
                )
                accum_loss.zero_()

        validation_loss = val(model, val_dataloader, criterion)
        logger.info(f"Validation loss for {epoch + 1}: {validation_loss:.2f}")

        checkpoint = {
            "epoch": epoch + 1,
            "model_state": raw_model.state_dict(),
            "scheduler_state": scheduler.state_dict(),
            "optimizer_state": optimizer.state_dict(),
            "config": vars(config),
        }
        torch.save(checkpoint, Path(config.checkpoint_dir) / f"{epoch}.pt")
        logger.info(f"Saved checkpoint for epoch {epoch + 1}")


if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO, filename="train.log", filemode="w")
    rtx4050_config = {
        "dataloader_workers": 2,
        "epochs": 5,
        "grad_steps": 8,
        "microbatch_size": 4,
        "logging_rate": 20,
        "warmup_batches": 200,
    }
    l4_config = {
        "dataloader_workers": 4,
        "epochs": 5,
        "grad_steps": 2,
        "microbatch_size": 16,
        "logging_rate": 50,
        "warmup_batches": 2000,
    }
    rtx4050_ft = rtx4050_config | {
        "epochs": 1,
        "lr": 1e-4,
        "model_file": "pretrained.pt",
        "dataset_type": "response",
        "dataset_file": "ultrachat.npy",
        "val_dataset_file": "ultrachat_val.npy"
    }
    train(config=rtx4050_ft)