`vp.gpt.train`

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

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

vp.gpt.train draws a model from seed and trains it on the text for steps steps. The model is the one of lib/gpt/gpt.viv: nanoGPT's layout cut down to a size a CPU trains in seconds, with the modded-nanogpt speedrun's choices (RMSNorm without weights, ReLU-squared in the MLP, the output map starting at zero). Each token is one character; the vocabulary is the sorted set of the text's characters, kept on the model as model.vocab. Token ids work too: pass an integer array and give GptConfig(vocab=...) if some ids never appear.

A step takes batch windows of context characters (window b of step t starts at (t * batch + b) mod (len(text) - context)), computes the mean cross-entropy of each next character and its gradient with respect to every matrix, and updates the model: Muon (momentum, then the nearest orthogonal matrix by Newton-Schulz) on the matrices inside the blocks, Adam on the two embeddings and the output map. The step size of both falls in a straight line from its start to a tenth of it. Nothing of this runs in Python: the whole loop is compiled Viv, and the gradient is the backward pass the compiler writes from sequence_loss (value_and_grad over the Gpt record).

The model comes back as Python objects holding NumPy arrays, named as the Viv records are:

m = r.model
m.tok, m.pos, m.unembed            # (vocab, width), (context, width), (width, vocab)
m.blocks[0].attn.heads[1].Wq       # (width, width / heads)
m.blocks[0].attn.Wo                # (width, width)
m.blocks[0].mlp.up, m.blocks[0].mlp.down
for name, w in m.parameters(): ...

vp.gpt.generate and vp.gpt.loss read these arrays again on every call, so a matrix changed in Python is the one they use.

Signature

vp.gpt.train(text, *, config=None, steps=30, seed=7, fast=False)

Parameters

text : str | ndarray[int], required: The training text (a character-level vocabulary is made from it), or token ids.
config : GptConfig, default None: The sizes (vp.gpt.GptConfig); None: GptConfig() with vocab from the text.
steps : int, default 30: Training steps; the step size falls in a straight line to a tenth of its start.
seed : int, default 7: The seed of the initial weights (Viv's rand(), Park-Miller MINSTD).
fast : bool, default False: The fast build (-O3, fused multiply-adds, the platform BLAS for matrix products): about twice as fast at 200,000 parameters; results may differ in the last bits.

Returns

A frozen GptTraining. Iterating it gives the fields below in order; it also has params, tests, log and check().

model : vp.gpt.Model: The trained model: model.tok, model.pos, model.blocks[l].attn.heads[h].Wq, model.blocks[l].attn.Wo, model.blocks[l].mlp.up, model.unembed, with model.config and model.vocab.
losses : ndarray[float64]: The batch loss (mean next-token cross-entropy, nats) of every step.

Raises

ValueError: Sizes that do not fit (width not a multiple of heads, a text no longer than the context), characters or token ids outside the model's vocabulary.
TypeError: An argument of the wrong type, or an unknown keyword.
vipyr.VivBuildError: An ImportError: no C compiler, or the program's C does not build.

Notes

Sizes (vp.gpt.GptConfig, the GptConfig record of lib/gpt/gpt.viv): context 8, width 16, heads 2, layers 2, hidden 64, batch 4 by default. A block is

X \leftarrow X + [softmax_{causal} (\frac{X ^ W _{q}^{h} ( X ^ W _{k}^{h} ) ^{T}}{d}) X^W_{v}^{h}]_{h} W_{o}, X \leftarrow X + relu (X^U)^{2} D

with $X^$ the RMS-normalized stream, $[\cdot]_{h}$ the heads side by side, and $d$ = width / heads. Muon updates a matrix $W$ of shape $r \times c$ by $W \leftarrow W - η max (1, r / c) NS_{5} (M)$ , $M$ the momentum (0.95) of the gradient and $NS_{5}$ five Newton-Schulz steps with the speedrun's coefficients (3.4445, -4.7750, 2.0315); Adam uses betas 0.9 and 0.99. Both start at a step size of 0.05.

References

Andrej Karpathy, nanoGPT (2022) https://github.com/karpathy/nanoGPT
Keller Jordan et al., modded-nanogpt: the NanoGPT speedrun (2024) https://github.com/KellerJordan/modded-nanogpt
Keller Jordan, Muon: an optimizer for hidden layers of neural networks (2024) https://kellerjordan.github.io/posts/muon/

Examples

import vipyr as vp
r = vp.gpt.train("hello world! " * 5, steps=30, seed=7)
print(r.model.size, r.losses[0], r.losses[-1])   # the first loss is log 9: nine characters
print(r.model.blocks[0].attn.Wo.shape)
assert r.check()

6560 2.1972245773362196 0.19911753110205457
(16, 16)

Each example ran as written when these pages were generated, and the output shown is what it printed.

Checked against

Reference: Against the NumPy program study/compare/ml/tiny_gpt.py (hand-written backward pass): its 30 losses, the first six bit for bit and all within 1.2 × 10⁻¹⁵, its trained matrices within 3.6 × 10⁻¹⁵ relative, and its sample to the character. Against a PyTorch 2.2.2 reimplementation (python/tests/gpt_reference.py: autograd, torch.optim.Adam, float64): 50 losses within 1.7 × 10⁻¹⁴ relative at 6,560 and 32,768 parameters. The gradient against central differences of the loss (lib/gpt/gpt.viv, run by r.check()). python/tests/test_gpt.py.

Source

Viv: lib/gpt/api.viv
Viv: lib/gpt/gpt.viv
Viv: lib/nn.viv
Viv: lib/optim.viv
Viv: lib/polar.viv
Viv: lib/random.viv
The defs it reaches and the calls between them: map of the code

vp.gpt.train