Language models (vp.gpt)

A small GPT in the shape of nanoGPT and the modded-nanogpt speedrun, trained on a CPU in seconds.

• token and position embeddings, RMSNorm, multi-head causal attention and a ReLU-squared MLP
• training by Muon on the block matrices and Adam on the rest, with gradients from Viv's value_and_grad
• train returns the model as named arrays (model.blocks[0].attn.Wo), and generate and loss take it back, edits included

Checked against a NumPy version, a PyTorch reimplementation and central differences.

Part of the distribution vipyr; it runs from a checkout of the repository. Planned on PyPI as pip install vipyr.

Functions

vp.gpt.attention

The attention weights of every head of a GPT for one prompt.

vp.gpt.generate

Continue a prompt with a trained GPT, one sampled token at a time.

vp.gpt.loss

The mean next-token cross-entropy of a GPT on a text.

vp.gpt.train

Train a small GPT on a text, character by character, and get the model back as named arrays.

Examples in the gallery

• A small GPT trained on one paragraph