Python and Viv · A small GPT
The same GPT in PyTorch, against Viv
What to notice
- The Python column is the model of tiny_gpt in PyTorch:
nn.Modulesubclasses,nn.Linearandnn.Embeddinglayers,F.scaled_dot_product_attention(..., is_causal=True),F.cross_entropy,torch.optim.Adamwith aLambdaLRschedule, and Muon as an 11-linetorch.optim.Optimizersubclass. There is no backward code:loss.backward()fills everyp.grad, which is why its loss and backward regions hold only a comment. It is the reference for what the Viv program would look like with autograd, arrays of modules and batched tensors. It needs torch (run with 2.2.2, in float64 on one thread). - The Viv column imports what the torch program imports:
#! use lib/nnforattention(one causal head, asF.scaled_dot_product_attention(..., is_causal=True)),softmax_rows,rmsnorm,relu2,cross_entropyandembedding(the row lookup ofnn.Embedding);#! use lib/optimforadam_updateandmuon_update, which update one matrix and its memories in place with the settings in anAdamSettingsorMuonSettingsrecord; and#! use lib/randomforrandn_matrixandchoice. It is the programlib/gpt/story_gpt.vivruns withlib/gpt/gpt.vivcopied in, and prints the same bytes as the written-out Viv of the NumPy pair (tiny_gpt): all 30 losses and every trained weight to 17 digits, and the sample, in the interpreter, in C and in wasm. Its model is the same tree as the torch modules, as records:Gpt { tok, pos, blocks: [Block], unembed },Block { attn, mlp },Attention { heads: [Head], Wo },Head { Wq, Wk, Wv }; its gradient isvalue_and_grad(sequence_loss, model, ...), aGpt; andoptimizer_stepisfor_leaves(muon_update, model.blocks, g.blocks, m.blocks, lr_scale, MUON_OPT)and threeadam_updatecalls, where torch callso.step()on each optimizer. - The outputs are the same text. Printed with 17 digits, the 30 losses differ from Viv's by at most 6.7e-16. The weights are copied from the same MINSTD draws (
init_fromfills each layer in the orderlib/gptlays them out, transposed wherenn.Linearstores(out, in)), and the samples use the same draws. Muon's Newton-Schulz step returns the matrix unchanged when its norm is zero, aslib/polardoes: at the first step the output map is zero, so every block matrix has a zero gradient, and dividing by the norm would give NaN. - torch 2.2 has no
F.rms_normand no Muon; both are written out here, in 2 and 11 lines (plus the 9 ofnewton_schulz). Viv imports both from its libraries, so the optimizer region is shorter in Viv (8 lines against 20). The torch program writesMinstd(12 lines) only to draw the weights and the samples Viv draws; Viv's random draws come fromlib/random.
Measured
Lines are code lines: blank lines, comments and docstrings are not counted; characters are those lines' characters. A lookup is a read of a size from the configuration record (cfg.<field>, c.<field>); torch's self.q(x) is not counted, and neither is Viv's h.Wq. study/compare/measure.py counts them.
| region | Python lines | Viv lines | Viv / Python | Python lookups | Viv lookups |
|---|---|---|---|---|---|
| imports | 8 | 3 | 0.4 | 0 | 0 |
| config | 15 | 15 | 1.0 | 0 | 0 |
| random draws | 12 | 0 | 0.0 | 0 | 0 |
| parameters and init | 8 | 25 | 3.1 | 0 | 19 |
| embedding | 0 | 0 | 0 | 0 | |
| attention | 14 | 9 | 0.6 | 6 | 0 |
| block | 11 | 5 | 0.5 | 4 | 0 |
| forward | 13 | 8 | 0.6 | 7 | 0 |
| loss | 0 | 3 | 0 | 0 | |
| backward | 0 | 0 | 0 | 0 | |
| optimizer | 20 | 8 | 0.4 | 0 | 0 |
| training loop | 21 | 18 | 0.9 | 4 | 4 |
| sampling | 7 | 10 | 1.4 | 1 | 0 |
| text as tokens | 4 | 4 | 1.0 | 0 | 0 |
| run | 8 | 12 | 1.5 | 0 | 0 |
| checks | 4 | 12 | 3.0 | 0 | 0 |
| all | 145 | 132 | 0.9 | 22 | 23 |
| characters of code | Python | Viv | Viv / Python |
|---|---|---|---|
| imports | 188 | 189 | 1.0 |
| config | 307 | 362 | 1.2 |
| random draws | 496 | 0 | 0.0 |
| parameters and init | 318 | 937 | 2.9 |
| embedding | 0 | 0 | |
| attention | 552 | 231 | 0.4 |
| block | 385 | 155 | 0.4 |
| forward | 440 | 215 | 0.5 |
| loss | 0 | 86 | |
| backward | 0 | 0 | |
| optimizer | 714 | 400 | 0.6 |
| training loop | 996 | 627 | 0.6 |
| sampling | 227 | 287 | 1.3 |
| text as tokens | 163 | 159 | 1.0 |
| run | 333 | 398 | 1.2 |
| checks | 176 | 469 | 2.7 |
| all | 5295 | 4515 | 0.9 |
| the forward and backward defs | Python | Viv |
|---|---|---|
| defs | 5 | 4 |
| parameters, all of them together | 5 | 9 |
| output parameters (arrays the def itself writes results into) | 0 | 0 |
| most in one def | 1 (forward) | 3 (sequence_loss) |
Where the Viv is still longer
The code follows the house style: a step of one statement is one line (step mlp: X = X + relu2(rmsnorm(X) @ b.mlp.up) @ b.mlp.down), an array of records is built by a comprehension of constructor calls ([Block(c) for l to c.layers]), and no line is wider than 100 columns, so a record that does not fit puts one field per line.
- parameters and init (25 against 8): the torch layers make and size their own weights in their constructors (
nn.Linear(cfg.width, size), in the attention, block and forward regions); the Viv records only declare the matrices, and the constructor defsHead,BlockandGptbuild them (model = Gpt(cfg)). Counting the parameters is a def handed tofor_leaves(count_into, into a one-entry array), where torch writessum(p.numel() for p in model.parameters()): a closure cannot assign the running count it captured, andleaves(model)is not built. - run (12 against 8): the configuration is one line,
GptConfig { vocab: len(vocab) }, with the other sizes the record's field defaults; the seed is set beforeGpt(cfg)draws, the parameter count takes two lines, and the losses are printed by a loop after training, where torch prints insidetrain. - loss (3 against 0):
value_and_graddifferentiates a named def, so the loss issequence_loss, where torch callsF.cross_entropyinside the training loop. - sampling (10 against 7): the context is
take the last len(m.pos) elements of out, where torch writesout[-n:](a Viv slice does not clamp a negative start). - text as tokens (4 against 4):
unique(TEXT)forsorted(set(TEXT)), a map comprehension forstoi, andencodeanddecodebound to lambdas whose bodies are comprehensions, as the Python binds them. - checks (12 against 4): each check is a named
testblock, where the Python writes one bareassert. - training loop (18 against 21):
batchcuts the windows with two comprehensions (the second stacks the windows into a batch-by-context matrix), as torch stacks a list of slices.value_and_gradreturns the loss and the gradient as two names, so its step is braced, to store the loss. - block (5 against 11), attention (9 against 14) and forward (8 against 13) are shorter than torch's: the block is the paper's two lines,
X = X + self_attention(b.attn, rmsnorm(X))and the MLP's, one step each, and lib/nn'sattentionreturns only the head's output.
Python tiny_gpt_torch.py · 204 lines
Viv tiny_gpt_torch.viv · 194 lines
py2viv refuses tiny_gpt_torch.py
1 imports
4import math
5from dataclasses import dataclass
6
7import numpy as np
8import torch
9import torch.nn as nn
10import torch.nn.functional as F
11
12torch.set_default_dtype(torch.float64)
13torch.set_num_threads(1)
2# lib/gpt/gpt and lib/gpt/story_gpt in one file. The layers (lib/nn), the
3# optimizers (lib/optim) and the random draws (lib/random) are imported, as
4# lib/gpt imports them and the PyTorch program imports torch.
5from nn import attention, cross_entropy, embedding, relu2, rmsnorm, softmax_rows
6from optim import adam_update, AdamSettings, muon_update, MuonSettings
7from random import choice, randn_matrix
py2viv refuses this program:
tiny_gpt_torch.py:8:1: not numeric Python: `import torch`
import torch
^^^^^^^^^^^^
py2viv does not translate PyTorch model code; write it in NumPy.
tiny_gpt_torch.py:9:1: not numeric Python: `import torch.nn`
import torch.nn as nn
^^^^^^^^^^^^^^^^^^^^^
py2viv does not translate PyTorch model code; write it in NumPy.
tiny_gpt_torch.py:10:1: not numeric Python: `import torch.nn.functional`
import torch.nn.functional as F
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
py2viv does not translate PyTorch model code; write it in NumPy.
tiny_gpt_torch.py:12:1: not numeric Python: a use of `torch` (from `torch`, refused at its import)
torch.set_default_dtype(torch.float64)
^^^^^
py2viv does not translate PyTorch model code; write it in NumPy.
tiny_gpt_torch.py:13:1: not numeric Python: a use of `torch` (from `torch`, refused at its import)
torch.set_num_threads(1)
^^^^^
py2viv does not translate PyTorch model code; write it in NumPy.
tiny_gpt_torch.py:46:56: not numeric Python: range of a float
u = [(self.rand(), self.rand()) for _ in range(rows * cols)]
^^^^^^^^^^^
Python raises TypeError: range takes ints.
tiny_gpt_torch.py:47:90: not numeric Python: iterating over a value of unknown type
z = [math.sqrt(-2.0 * math.log(u1)) * math.cos(2.0 * math.pi * u2) for u1, u2 in u]
^
Fix the refusal that produced it.
tiny_gpt_torch.py:48:16: not numeric Python: a use of `torch` (from `torch`, refused at its import)
return torch.tensor(z).reshape(rows, cols)
^^^^^
py2viv does not translate PyTorch model code; write it in NumPy.
tiny_gpt_torch.py:51:24: not numeric Python: `np.searchsorted`
return min(int(np.searchsorted(np.cumsum(p), self.rand(), side="right")), len(p) - 1)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
Not in the NumPy mapping yet (tools/py2viv/SUBSET.md lists what is).
tiny_gpt_torch.py:54:2: not numeric Python: a decorator
@torch.no_grad()
^^^^^^^^^^^^^^^
Drop the decorator.
tiny_gpt_torch.py:68:12: not numeric Python: a class that inherits from `nn.Module`
class Head(nn.Module):
^^^^^^^^^
Viv records have no inheritance: put the fields and methods in one class, or write functions.
tiny_gpt_torch.py:78:17: not numeric Python: a class that inherits from `nn.Module`
class Attention(nn.Module):
^^^^^^^^^
Viv records have no inheritance: put the fields and methods in one class, or write functions.
tiny_gpt_torch.py:89:16: not numeric Python: a use of `torch` (from `torch`, refused at its import)
return x / torch.sqrt(x.pow(2).mean(-1, keepdim=True) + 1e-8)
^^^^^
py2viv does not translate PyTorch model code; write it in NumPy.
tiny_gpt_torch.py:92:13: not numeric Python: a class that inherits from `nn.Module`
class Block(nn.Module):
^^^^^^^^^
Viv records have no inheritance: put the fields and methods in one class, or write functions.
tiny_gpt_torch.py:104:11: not numeric Python: a class that inherits from `nn.Module`
class GPT(nn.Module):
^^^^^^^^^
Viv records have no inheritance: put the fields and methods in one class, or write functions.
tiny_gpt_torch.py:128:30: not numeric Python: the attribute `.T` of a value of unknown type (it comes from a construct refused elsewhere)
return newton_schulz(g.T, steps, a, b, c).T
^^^
Fix the refusal that produced it.
tiny_gpt_torch.py:129:12: not numeric Python: the method `.norm` of a value of unknown type (it comes from a construct refused elsewhere)
norm = g.norm()
^^^^^^^^
Fix the refusal that produced it.
tiny_gpt_torch.py:132:17: not numeric Python: the attribute `.T` of a value of unknown type (it comes from a construct refused elsewhere)
s = x @ x.T
^^^
Fix the refusal that produced it.
tiny_gpt_torch.py:137:12: not numeric Python: a class that inherits from `torch.optim.Optimizer`
class Muon(torch.optim.Optimizer):
^^^^^^^^^^^^^^^^^^^^^
Viv records have no inheritance: put the fields and methods in one class, or write functions.
tiny_gpt_torch.py:152:29: not numeric Python: the attribute `.context` of a value of unknown type (it comes from a construct refused elsewhere)
windows = len(tokens) - cfg.context
^^^^^^^^^^^
Fix the refusal that produced it.
tiny_gpt_torch.py:153:63: not numeric Python: the attribute `.batch` of a value of unknown type (it comes from a construct refused elsewhere)
starts = [(step * cfg.batch + b) % windows for b in range(cfg.batch)]
^^^^^^^^^
Fix the refusal that produced it.
tiny_gpt_torch.py:154:12: not numeric Python: a use of `torch` (from `torch`, refused at its import)
return torch.stack([tokens[s + shift : s + shift + cfg.context] for s in starts])
^^^^^
py2viv does not translate PyTorch model code; write it in NumPy.
tiny_gpt_torch.py:158:34: not numeric Python: the method `.named_parameters` of a value of unknown type (it comes from a construct refused elsewhere)
matrices = [p for name, p in model.named_parameters() if name.startswith("blocks.")]
^^^^^^^^^^^^^^^^^^^^^^^^
Fix the refusal that produced it.
tiny_gpt_torch.py:159:32: not numeric Python: the method `.named_parameters` of a value of unknown type (it comes from a construct refused elsewhere)
others = [p for name, p in model.named_parameters() if not name.startswith("blocks.")]
^^^^^^^^^^^^^^^^^^^^^^^^
Fix the refusal that produced it.
tiny_gpt_torch.py:160:53: not numeric Python: a use of `torch` (from `torch`, refused at its import)
opts = [Muon(matrices, MUON_LR, MUON_MOMENTUM), torch.optim.Adam(others, ADAM_LR, betas=ADAM_BETAS)]
^^^^^
py2viv does not translate PyTorch model code; write it in NumPy.
tiny_gpt_torch.py:161:110: not numeric Python: iterating over a value of unknown type
schedules = [torch.optim.lr_scheduler.LambdaLR(o, lambda s: 1.0 - (1.0 - LR_FLOOR) * s / steps) for o in opts]
^^^^
Fix the refusal that produced it.
tiny_gpt_torch.py:164:34: not numeric Python: the attribute `.cfg` of a value of unknown type (it comes from a construct refused elsewhere)
x, y = get_batch(tokens, model.cfg, step, 0), get_batch(tokens, model.cfg, step, 1)
^^^^^^^^^
Fix the refusal that produced it.
tiny_gpt_torch.py:165:18: not numeric Python: a call of `model`
logits = model(x)
^^^^^^^^
Not a def in this file and not in the numeric subset.
tiny_gpt_torch.py:166:16: not numeric Python: a use of `F` (from `torch.nn.functional`, refused at its import)
loss = F.cross_entropy(logits.view(-1, logits.shape[-1]), y.view(-1))
^
py2viv does not translate PyTorch model code; write it in NumPy.
tiny_gpt_torch.py:167:9: not numeric Python: the method `.zero_grad` of a value of unknown type (it comes from a construct refused elsewhere)
model.zero_grad()
^^^^^^^^^^^^^^^^^
Fix the refusal that produced it.
tiny_gpt_torch.py:168:9: not numeric Python: the method `.backward` of a value of unknown type (it comes from a construct refused elsewhere)
loss.backward()
^^^^^^^^^^^^^^^
Fix the refusal that produced it.
tiny_gpt_torch.py:169:18: not numeric Python: iterating over a value of unknown type
for o in opts + schedules:
^^^^^^^^^^^^^^^^
Loop over `range(len(a))` and index.
tiny_gpt_torch.py:171:23: not numeric Python: the method `.item` of a value of unknown type (it comes from a construct refused elsewhere)
losses.append(loss.item())
^^^^^^^^^^^
Fix the refusal that produced it.
tiny_gpt_torch.py:172:40: not numeric Python: the method `.item` of a value of unknown type (it comes from a construct refused elsewhere)
print(f"step {step + 1}: loss {loss.item():.6f}")
^^^^^^^^^^^
Fix the refusal that produced it.
tiny_gpt_torch.py:176:2: not numeric Python: a decorator
@torch.no_grad()
^^^^^^^^^^^^^^^
Drop the decorator.
tiny_gpt_torch.py:188:48: not numeric Python: iterating over a value of unknown type
decode = lambda ids: "".join(vocab[i] for i in ids)
^^^
Fix the refusal that produced it.
tiny_gpt_torch.py:194:32: not numeric Python: the method `.parameters`, which an object of class `GPT` does not have
count = sum(p.numel() for p in model.parameters())
^^^^^^^^^^^^^^^^^^
Python raises AttributeError.
tiny_gpt_torch.py:195:10: not numeric Python: printing a value of unknown type (it comes from a construct refused elsewhere)
print(f"{count} parameters, {len(vocab)} characters")
^^^^^
Fix the refusal that produced it.
tiny_gpt_torch.py:196:23: not numeric Python: a use of `torch` (from `torch`, refused at its import)
losses = train(model, torch.tensor(encode(TEXT)), STEPS)
^^^^^
py2viv does not translate PyTorch model code; write it in NumPy.
tiny_gpt_torch.py:198:7: not numeric Python: printing a value of unknown type (it comes from a construct refused elsewhere)
print(written)
^^^^^^^
Fix the refusal that produced it.
40 constructs outside the numeric subset (tools/py2viv/SUBSET.md) 2 config
16@dataclass
17class Config:
18 vocab: int
19 context: int = 8
20 width: int = 16
21 heads: int = 2
22 layers: int = 2
23 hidden: int = 64
24 batch: int = 4
25
26
27MUON_LR, MUON_MOMENTUM = 0.05, 0.95
28ADAM_LR, ADAM_BETAS = 0.05, (0.9, 0.99)
29LR_FLOOR = 0.1
30TEXT, PROMPT = "hello world! " * 5, "hello"
31STEPS, SAMPLES = 30, 24
32INIT_SEED, SAMPLE_SEED = 7, 11
10record GptConfig {
11 vocab: Num,
12 context: Num = 8,
13 width: Num = 16,
14 heads: Num = 2,
15 layers: Num = 2,
16 hidden: Num = 64,
17 batch: Num = 4,
18}
19
20MUON_OPT = MuonSettings { lr: 0.05, momentum: 0.95 }
21ADAM_OPT = AdamSettings { lr: 0.05, beta1: 0.9, beta2: 0.99, eps: 1.0e-8 }
22# the step size falls in a straight line to this fraction of its start
23LR_FLOOR = 0.1
24TEXT, PROMPT = "hello world! " * 5, "hello"
25STEPS, SAMPLES = 30, 24
26INIT_SEED, SAMPLE_SEED = 7, 11
(refused)
3 random draws
35class Minstd:
36 """Viv's rand(), so the weights and the samples are the ones the Viv program draws."""
37
38 def __init__(self, seed):
39 self.state = math.floor(abs(seed)) % 2147483646 + 1
40
41 def rand(self):
42 self.state = 16807 * self.state % 2147483647
43 return self.state / 2147483647
44
45 def randn(self, rows, cols):
46 u = [(self.rand(), self.rand()) for _ in range(rows * cols)]
47 z = [math.sqrt(-2.0 * math.log(u1)) * math.cos(2.0 * math.pi * u2) for u1, u2 in u]
48 return torch.tensor(z).reshape(rows, cols)
49
50 def choice(self, p):
51 return min(int(np.searchsorted(np.cumsum(p), self.rand(), side="right")), len(p) - 1)
29# randn_matrix and choice: lib/random
(refused)
4 parameters and init
54@torch.no_grad()
55def init_from(model, rng):
56 """N(0, 1) embeddings, N(0, 1/fan_in) matrices, a zero output map, drawn in the Viv program's order."""
57 for m in model.modules():
58 if isinstance(m, nn.Embedding):
59 m.weight.copy_(rng.randn(*m.weight.shape))
60 elif isinstance(m, nn.Linear) and m is not model.unembed:
61 m.weight.copy_(rng.randn(m.in_features, m.out_features).T / math.sqrt(m.in_features))
62 nn.init.zeros_(model.unembed.weight)
32record Head { Wq: Arr, Wk: Arr, Wv: Arr }
33record Attention { heads: [Head], Wo: Arr }
34record Mlp { up: Arr, down: Arr }
35record Block { attn: Attention, mlp: Mlp }
36record Gpt { tok: Arr, pos: Arr, blocks: [Block], unembed: Arr }
37
38def init(rows, cols) {
39 return (1.0 / sqrt(rows)) * randn_matrix(rows, cols)
40}
41
42def Head(c: GptConfig) -> Head {
43 hd = c.width / c.heads
44 return Head { Wq: init(c.width, hd), Wk: init(c.width, hd), Wv: init(c.width, hd) }
45}
46
47def Block(c: GptConfig) -> Block {
48 attn = Attention { heads: [Head(c) for h to c.heads], Wo: init(c.width, c.width) }
49 mlp = Mlp { up: init(c.width, c.hidden), down: init(c.hidden, c.width) }
50 return Block { attn: attn, mlp: mlp }
51}
52
53def Gpt(c: GptConfig) -> Gpt {
54 tok, pos = randn_matrix(c.vocab, c.width), randn_matrix(c.context, c.width)
55 blocks = [Block(c) for l to c.layers]
56 return Gpt { tok: tok, pos: pos, blocks: blocks, unembed: zeros(c.width, c.vocab) }
57}
58
59def count_into(w: Arr, n: Arr) {
60 n[0] += numel(w)
61}
(refused)
5 embedding
65# nn.Embedding: model.tok and model.pos below
64# embedding: lib/nn, in gpt_forward below
(refused)
6 attention
68class Head(nn.Module):
69 def __init__(self, cfg):
70 super().__init__()
71 size = cfg.width // cfg.heads
72 self.q, self.k, self.v = (nn.Linear(cfg.width, size, bias=False) for _ in range(3))
73
74 def forward(self, x):
75 return F.scaled_dot_product_attention(self.q(x), self.k(x), self.v(x), is_causal=True)
76
77
78class Attention(nn.Module):
79 def __init__(self, cfg):
80 super().__init__()
81 self.heads = nn.ModuleList(Head(cfg) for _ in range(cfg.heads))
82 self.out = nn.Linear(cfg.width, cfg.width, bias=False)
83
84 def forward(self, x):
85 return self.out(torch.cat([h(x) for h in self.heads], dim=-1))
67# attention (one causal head) and softmax_rows: lib/nn
68
69def self_attention(a: Attention, N: Arr) {
70 hd = shape(N)[-1] / len(a.heads)
71 Z = zeros_like(N)
72 for i, h in enumerate(a.heads) {
73 step head: Z[:, :, i * hd:(i + 1) * hd] = attention(N, h.Wq, h.Wk, h.Wv)
74 }
75 step project: Y = Z @ a.Wo
76 return Y
77}
(refused)
7 block
88def rmsnorm(x):
89 return x / torch.sqrt(x.pow(2).mean(-1, keepdim=True) + 1e-8)
90
91
92class Block(nn.Module):
93 def __init__(self, cfg):
94 super().__init__()
95 self.attn = Attention(cfg)
96 self.up = nn.Linear(cfg.width, cfg.hidden, bias=False)
97 self.down = nn.Linear(cfg.hidden, cfg.width, bias=False)
98
99 def forward(self, x):
100 x = x + self.attn(rmsnorm(x))
101 return x + self.down(F.relu(self.up(rmsnorm(x))).square())
80# rmsnorm and relu2: lib/nn
81
82def block(b: Block, X: Arr) {
83 step attention: X = X + self_attention(b.attn, rmsnorm(X))
84 step mlp: X = X + relu2(rmsnorm(X) @ b.mlp.up) @ b.mlp.down
85 return X
86}
(refused)
8 forward
104class GPT(nn.Module):
105 def __init__(self, cfg):
106 super().__init__()
107 self.cfg = cfg
108 self.tok = nn.Embedding(cfg.vocab, cfg.width)
109 self.pos = nn.Embedding(cfg.context, cfg.width)
110 self.blocks = nn.ModuleList(Block(cfg) for _ in range(cfg.layers))
111 self.unembed = nn.Linear(cfg.width, cfg.vocab, bias=False)
112
113 def forward(self, idx):
114 x = self.tok(idx) + self.pos(torch.arange(idx.shape[1]))
115 for block in self.blocks:
116 x = block(x)
117 return self.unembed(rmsnorm(x))
89def gpt_forward(m: Gpt, tokens: Arr) {
90 step embed: X = embedding(m.tok, tokens) + m.pos[0:shape(tokens)[-1], :]
91 for b in m.blocks {
92 step block: X = block(b, X)
93 }
94 step logits: logits = rmsnorm(X) @ m.unembed
95 return logits
96}
(refused)
9 loss
120# F.cross_entropy, in the training loop
99# cross_entropy: lib/nn
100
101def sequence_loss(m: Gpt, x: Arr, y: Arr) {
102 return cross_entropy(gpt_forward(m, x), y)
103}
(refused)
10 backward
123# autograd: loss.backward() in the training loop fills every p.grad
106# nothing: `train` calls `value_and_grad(sequence_loss, model, ...)`, and the
107# compiler writes the backward pass from the forward pass above
(refused)
11 optimizer
126def newton_schulz(g, steps=5, a=3.4445, b=-4.775, c=2.0315):
127 if g.shape[0] > g.shape[1]:
128 return newton_schulz(g.T, steps, a, b, c).T
129 norm = g.norm()
130 x = g / norm if norm > 0 else g
131 for _ in range(steps):
132 s = x @ x.T
133 x = a * x + b * (s @ x) + c * ((s @ s) @ x)
134 return x
135
136
137class Muon(torch.optim.Optimizer):
138 def __init__(self, params, lr, momentum):
139 super().__init__(params, dict(lr=lr, momentum=momentum))
140
141 @torch.no_grad()
142 def step(self):
143 for group in self.param_groups:
144 for p in group["params"]:
145 buf = self.state[p].setdefault("buf", torch.zeros_like(p))
146 buf.mul_(group["momentum"]).add_(p.grad)
147 cols, rows = p.shape # a Linear weight is stored (out, in)
148 p.sub_(group["lr"] * math.sqrt(max(1.0, rows / cols)) * newton_schulz(buf.T).T)
110# adam_update and muon_update: lib/optim
111
112def optimizer_step(model: Gpt, g: Gpt, m: Gpt, v: Gpt, taken, lr_scale) {
113 step muon: for_leaves(muon_update, model.blocks, g.blocks, m.blocks, lr_scale, MUON_OPT)
114 step adam {
115 adam_update(model.tok, g.tok, m.tok, v.tok, taken, lr_scale, ADAM_OPT)
116 adam_update(model.pos, g.pos, m.pos, v.pos, taken, lr_scale, ADAM_OPT)
117 adam_update(model.unembed, g.unembed, m.unembed, v.unembed, taken, lr_scale, ADAM_OPT)
118 }
119}
(refused)
12 training loop
151def get_batch(tokens, cfg, step, shift):
152 windows = len(tokens) - cfg.context
153 starts = [(step * cfg.batch + b) % windows for b in range(cfg.batch)]
154 return torch.stack([tokens[s + shift : s + shift + cfg.context] for s in starts])
155
156
157def train(model, tokens, steps):
158 matrices = [p for name, p in model.named_parameters() if name.startswith("blocks.")]
159 others = [p for name, p in model.named_parameters() if not name.startswith("blocks.")]
160 opts = [Muon(matrices, MUON_LR, MUON_MOMENTUM), torch.optim.Adam(others, ADAM_LR, betas=ADAM_BETAS)]
161 schedules = [torch.optim.lr_scheduler.LambdaLR(o, lambda s: 1.0 - (1.0 - LR_FLOOR) * s / steps) for o in opts]
162 losses = []
163 for step in range(steps):
164 x, y = get_batch(tokens, model.cfg, step, 0), get_batch(tokens, model.cfg, step, 1)
165 logits = model(x)
166 loss = F.cross_entropy(logits.view(-1, logits.shape[-1]), y.view(-1))
167 model.zero_grad()
168 loss.backward()
169 for o in opts + schedules:
170 o.step()
171 losses.append(loss.item())
172 print(f"step {step + 1}: loss {loss.item():.6f}")
173 return losses
122def batch(tokens: Arr, c: GptConfig, t, shift) {
123 starts = [mod(t * c.batch + b, len(tokens) - c.context) + shift for b to c.batch]
124 return [tokens[s:s + c.context] for s in starts]
125}
126
127def train(model: Gpt, c: GptConfig, tokens: Arr, steps) {
128 losses = zeros(steps)
129 m, v = zeros_like(model), zeros_like(model)
130 for t to steps {
131 step batch: xs, ys = batch(tokens, c, t, 0), batch(tokens, c, t, 1)
132 step forward_backward {
133 loss, grad = value_and_grad(sequence_loss, model, xs, ys)
134 losses[t] = loss
135 }
136 step measure: gnorm = sqrt(sum(grad * grad))
137 step update: optimizer_step(model, grad, m, v, t + 1, 1.0 - (1.0 - LR_FLOOR) * t / steps)
138 }
139 return losses
140}
(refused)
13 sampling
176@torch.no_grad()
177def generate(model, prompt, count, rng):
178 out = list(prompt)
179 for _ in range(count):
180 logits = model(torch.tensor([out[-model.cfg.context :]]))
181 out.append(rng.choice(F.softmax(logits[0, -1], dim=-1).numpy()))
182 return out
143def sample(m: Gpt, prompt: Arr, count, seed) {
144 srand(seed)
145 out = prompt
146 for k to count {
147 step context: ctx = take the last len(m.pos) elements of out
148 step probs: p = softmax_rows(gpt_forward(m, reshape(ctx, 1, len(ctx)))[0, len(ctx) - 1, :])
149 step draw: out = append(out, choice(p))
150 }
151 return out
152}
(refused)
14 text as tokens
185vocab = sorted(set(TEXT))
186stoi = {c: i for i, c in enumerate(vocab)}
187encode = lambda text: [stoi[c] for c in text]
188decode = lambda ids: "".join(vocab[i] for i in ids)
155vocab = unique(TEXT)
156stoi = {c: i for i, c in enumerate(vocab)}
157encode = fn(text) -> [stoi[c] for c in text]
158decode = fn(ids) -> join("", [vocab[i] for i in ids])
(refused)
15 run
191cfg = Config(vocab=len(vocab))
192model = GPT(cfg)
193init_from(model, Minstd(INIT_SEED))
194count = sum(p.numel() for p in model.parameters())
195print(f"{count} parameters, {len(vocab)} characters")
196losses = train(model, torch.tensor(encode(TEXT)), STEPS)
197written = decode(generate(model, encode(PROMPT), SAMPLES, Minstd(SAMPLE_SEED)))
198print(written)
161cfg = GptConfig { vocab: len(vocab) }
162srand(INIT_SEED)
163model = Gpt(cfg)
164count = [0.0]
165for_leaves(count_into, model, count)
166# story_gpt prints the four numbers bare; here every loss is printed, to six
167# decimals, so the Python run can be compared with it line by line
168print(f"{count[0]:d} parameters, {len(vocab):d} characters")
169step train: losses = train(model, cfg, encode(TEXT), STEPS)
170for t to STEPS {
171 print(f"step {t + 1:d}: loss {losses[t]:.6f}")
172}
173step generate: written = decode(sample(model, encode(PROMPT), SAMPLES, SAMPLE_SEED))
174print(written)
(refused)
16 checks
201assert count == 6560 and len(vocab) == 9
202assert abs(losses[0] - math.log(9.0)) <= 1e-15
203assert abs(losses[-1] - 0.199118) <= 1e-6
204assert written == "hello world! hello world! hel"
177test "the model has 6,560 parameters, and 9 characters make its vocabulary" {
178 assert_true(count[0] == 6560 and len(vocab) == 9)
179}
180
181test "the first loss is log 9 exactly, because the output map starts at zero" {
182 assert_near(losses[0], log(9.0), 1.0e-15)
183}
184
185test "thirty steps take the loss from 2.197 to 0.1991" {
186 assert_near(losses[STEPS - 1], 0.199118, 1.0e-6)
187}
188
189test "from the prompt hello, the model writes the line on" {
190 assert_true(written == "hello world! hello world! hel")
191}
(refused)
What they print
identical — the two programs print the same text.
Dimmed lines are left out of the comparison. Each Viv test block prints one PASS line; the Python checks the same facts with assert, which prints nothing.
$ python3 tiny_gpt_torch.py
6560 parameters, 9 charactersstep 1: loss 2.197225step 2: loss 2.102097step 3: loss 1.879860step 4: loss 1.575851step 5: loss 1.150941step 6: loss 1.198318step 7: loss 0.840376step 8: loss 0.673002step 9: loss 0.686226step 10: loss 0.541481step 11: loss 0.403199step 12: loss 0.305271step 13: loss 0.563067step 14: loss 0.306272step 15: loss 0.314534step 16: loss 0.195234step 17: loss 0.539175step 18: loss 0.290064step 19: loss 0.191980step 20: loss 0.334400step 21: loss 0.255869step 22: loss 0.229073step 23: loss 0.151354step 24: loss 0.341850step 25: loss 0.196461step 26: loss 0.141672step 27: loss 0.266627step 28: loss 0.179300step 29: loss 0.143552step 30: loss 0.199118hello world! hello world! hel
$ viv run tiny_gpt_torch.viv
6560 parameters, 9 charactersstep 1: loss 2.197225step 2: loss 2.102097step 3: loss 1.879860step 4: loss 1.575851step 5: loss 1.150941step 6: loss 1.198318step 7: loss 0.840376step 8: loss 0.673002step 9: loss 0.686226step 10: loss 0.541481step 11: loss 0.403199step 12: loss 0.305271step 13: loss 0.563067step 14: loss 0.306272step 15: loss 0.314534step 16: loss 0.195234step 17: loss 0.539175step 18: loss 0.290064step 19: loss 0.191980step 20: loss 0.334400step 21: loss 0.255869step 22: loss 0.229073step 23: loss 0.151354step 24: loss 0.341850step 25: loss 0.196461step 26: loss 0.141672step 27: loss 0.266627step 28: loss 0.179300step 29: loss 0.143552step 30: loss 0.199118hello world! hello world! helPASS the model has 6,560 parameters, and 9 characters make its vocabularyPASS the first loss is log 9 exactly, because the output map starts at zeroPASS thirty steps take the loss from 2.197 to 0.1991PASS from the prompt hello, the model writes the line on