This site is version 0.0.1. Most of the code on this site is drafted by an agent and needs review. I apologize for the AI-speak, we are working on rewriting it!
viv-lang.org

Python and Viv · A small GPT

A small GPT in NumPy and in Viv

What to notice

Measured

Lines are code lines: blank lines, comments and docstrings are not counted; characters are those lines' characters. A lookup is one place the program finds a weight or a size: in Python params["q"] or cfg.width; in Viv a weight read from its record by name (h.Wq, b.mlp.up, m.tok), or a read of a size from the configuration record (c.width). study/compare/measure.py counts them, and check.sh fails when these tables are out of date.

regionPython linesViv linesViv / PythonPython lookupsViv lookupsViv in phrases linesViv in phrases lookups
imports300.00000
config23150.700150
random draws15191.300190
parameters and init14251.80192619
embedding1430.210232
attention41290.7154284
block35160.5102162
forward1580.55181
loss1190.80090
backward600.01000
optimizer27381.40123912
training loop15181.244214
sampling6101.711101
text as tokens640.70070
run7121.700210
checks4123.000120
all2422180.9464523445
characters of codePythonVivViv / PythonViv in phrases
imports6200.00
config4233380.8312
random draws4972910.6291
parameters and init4669372.0937
embedding524990.299
attention15187060.5854
block12243390.3339
forward4931750.4175
loss3913590.9377
backward24000.00
optimizer90812611.41346
training loop5966271.1656
sampling1972901.5287
text as tokens1671591.0193
run3223981.2503
checks1764692.7469
all820464480.86838
the forward and backward defsPythonViv
defs1611
parameters, all of them together1921
output parameters (arrays the def itself writes results into)00
most in one def3 (rmsnorm_backward)4 (attention)

Where the extra Viv lines come from

Open the Viv in the workbench The atlas of the Viv py2viv refuses this Python (the reason is in its column).
Python tiny_gpt.py · 327 lines
Viv tiny_gpt.viv · 293 lines
py2viv refuses tiny_gpt.py
Viv in phrases tiny_gpt.phrased.viv · 311 lines
1 imports
4import math
5from dataclasses import dataclass
6
7import numpy as np
2# lib/gpt/gpt and the parts of lib/nn, lib/polar, lib/optim and lib/random it
3# calls, copied into one file (the program is lib/gpt/story_gpt.viv with its
4# imports expanded by hand, and the optimizers' settings written as
5# constants). The NumPy program writes its layers and optimizers out, and so
6# does this one; tiny_gpt_torch.viv imports them, as the PyTorch program does.
py2viv refuses this program:
tiny_gpt.py:52:26: not numeric Python: a call of `Minstd.rand`, which changes the object passed as `self`, inside an expression
                    u1, u2 = self.rand(), self.rand()
                             ^^^^^^^^^^^
  Call it on its own line, or as the whole right side of an assignment (`x = obj.m()`), so the translation can give the changed object back to its variable.
tiny_gpt.py:57: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.py:68:13: not numeric Python: an expression whose value is dropped
                yield prefix + name, p, self.grads.get(name)
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  Assign it, print it, or remove it.
tiny_gpt.py:70:13: not numeric Python: an expression whose value is dropped
                yield from child.named_parameters(f"{prefix}{name}.")
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  Assign it, print it, or remove it.
tiny_gpt.py:74:8: not numeric Python: `is` / `is not`
        if scale is None:
           ^^^^^^^^^^^^^
  Compare values with `==`; `None` is not a Viv value.
tiny_gpt.py:76:20: not numeric Python: the method `.randn` of a value of unknown type (it comes from a construct refused elsewhere)
        return scale * rng.randn(rows, cols)
                       ^^^^^^^^^^^^^^^^^^^^^
  Fix the refusal that produced it.
tiny_gpt.py:81:12: not numeric Python: the method `.reshape` of a value of unknown type (it comes from a construct refused elsewhere)
        return a.reshape(-1, a.shape[-1])
               ^^^^^^^^^^^^^^^^^^^^^^^^^^
  Fix the refusal that produced it.
tiny_gpt.py:84:17: not numeric Python: a class that inherits from `Module`
    class Embedding(Module):
                    ^^^^^^
  Viv records have no inheritance: put the fields and methods in one class, or write functions.
tiny_gpt.py:103:20: not numeric Python: the method `.max` of a value of unknown type (it comes from a construct refused elsewhere)
        e = np.exp(x - x.max(axis=-1, keepdims=True))
                       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  Fix the refusal that produced it.
tiny_gpt.py:104:16: not numeric Python: the method `.sum` of a value of unknown type (it comes from a construct refused elsewhere)
        return e / e.sum(axis=-1, keepdims=True)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  Fix the refusal that produced it.
tiny_gpt.py:107:12: not numeric Python: a class that inherits from `Module`
    class Head(Module):
               ^^^^^^
  Viv records have no inheritance: put the fields and methods in one class, or write functions.
tiny_gpt.py:136:17: not numeric Python: a class that inherits from `Module`
    class Attention(Module):
                    ^^^^^^
  Viv records have no inheritance: put the fields and methods in one class, or write functions.
tiny_gpt.py:154:17: not numeric Python: the method `.mean` of a value of unknown type (it comes from a construct refused elsewhere)
        r = np.sqrt((x * x).mean(axis=-1, keepdims=True) + EPS)
                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  Fix the refusal that produced it.
tiny_gpt.py:159:11: not numeric Python: the method `.sum` of a value of unknown type (it comes from a construct refused elsewhere)
        dot = (dy * x).sum(axis=-1, keepdims=True)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  Fix the refusal that produced it.
tiny_gpt.py:163:11: not numeric Python: a class that inherits from `Module`
    class MLP(Module):
              ^^^^^^
  Viv records have no inheritance: put the fields and methods in one class, or write functions.
tiny_gpt.py:182:13: not numeric Python: a class that inherits from `Module`
    class Block(Module):
                ^^^^^^
  Viv records have no inheritance: put the fields and methods in one class, or write functions.
tiny_gpt.py:200:11: not numeric Python: a class that inherits from `Module`
    class GPT(Module):
              ^^^^^^
  Viv records have no inheritance: put the fields and methods in one class, or write functions.
tiny_gpt.py:244: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.py:245:12: not numeric Python: np.linalg.norm of a value of unknown type
        norm = np.linalg.norm(g)
               ^^^^^^^^^^^^^^^^^
  Pass a number array.
tiny_gpt.py:248: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.py:262:27: not numeric Python: the method `.named_parameters`, which an object of class `GPT` does not have
            for name, p, g in model.named_parameters():
                              ^^^^^^^^^^^^^^^^^^^^^^^^
  Python raises AttributeError.
tiny_gpt.py:277:29: not numeric Python: the attribute `.context`
        windows = len(tokens) - cfg.context
                                ^^^^^^^^^^^
  Viv has no objects; use the NumPy function form, e.g. np.sum(a).
tiny_gpt.py:278:63: not numeric Python: the attribute `.batch`
        starts = [(step * cfg.batch + b) % windows for b in range(cfg.batch)]
                                                                  ^^^^^^^^^
  Viv has no objects; use the NumPy function form, e.g. np.sum(a).
tiny_gpt.py:279:75: not numeric Python: iterating over a value of unknown type
        return np.array([tokens[s + shift : s + shift + cfg.context] for s in starts])
                                                                              ^^^^^^
  Fix the refusal that produced it.
tiny_gpt.py:287:25: not numeric Python: unpacking a value of unknown type (lifted by R5)
            loss, dlogits = model.loss(model.forward(x), y)
                            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  Only several returned values, a vector (its rows for a matrix are not unpacked) and a shape unpack.
tiny_gpt.py:290:23: not numeric Python: appending a value of unknown type to a list (lifted by R17)
            losses.append(loss)
                          ^^^^
  A list py2viv grows holds numbers, or rows of numbers (a new list, a copy, or a tuple).
tiny_gpt.py:291:40: not numeric Python: printing a value of unknown type (it comes from a construct refused elsewhere)
            print(f"step {step + 1}: loss {loss:.6f}")
                                           ^^^^
  Fix the refusal that produced it.
tiny_gpt.py:298:47: not numeric Python: the attribute `.context`
            logits = model.forward(np.array([out[-model.cfg.context :]]))
                                                  ^^^^^^^^^^^^^^^^^
  Viv has no objects; use the NumPy function form, e.g. np.sum(a).
tiny_gpt.py:299:20: not numeric Python: appending a value of unknown type to a list (lifted by R17)
            out.append(rng.choice(softmax(logits[0, -1])))
                       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  A list py2viv grows holds numbers, or rows of numbers (a new list, a copy, or a tuple).
tiny_gpt.py:317:35: not numeric Python: the method `.named_parameters`, which an object of class `GPT` does not have
    count = sum(p.size for _, p, _ in model.named_parameters())
                                      ^^^^^^^^^^^^^^^^^^^^^^^^
  Python raises AttributeError.
tiny_gpt.py:318: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.
31 constructs outside the numeric subset (tools/py2viv/SUBSET.md)
2# lib/gpt/gpt and the parts of lib/nn, lib/polar, lib/optim and lib/random it
3# calls, copied into one file (the program is lib/gpt/story_gpt.viv with its
4# imports expanded by hand), with the phrases those modules declare and
5# use
2 config
10@dataclass
11class Config:
12 vocab: int
13 context: int = 8
14 width: int = 16
15 heads: int = 2
16 layers: int = 2
17 hidden: int = 64
18 batch: int = 4
19
20
21MUON_LR = 0.05
22MUON_MOMENTUM = 0.95
23ADAM_LR = 0.05
24ADAM_BETA1, ADAM_BETA2 = 0.9, 0.99
25ADAM_EPS = 1e-8
26LR_FLOOR = 0.1 # the step size falls in a straight line to this fraction of its start
27EPS = 1e-8
28MASKED = -1e30
29
30TEXT = "hello world! " * 5
31PROMPT = "hello"
32STEPS = 30
33SAMPLES = 24
34INIT_SEED = 7
35SAMPLE_SEED = 11
9record GptConfig {
10 vocab: Num,
11 context: Num = 8,
12 width: Num = 16,
13 heads: Num = 2,
14 layers: Num = 2,
15 hidden: Num = 64,
16 batch: Num = 4,
17}
18
19# the optimizers' constants
20MUON_LR, MUON_MOMENTUM = 0.05, 0.95
21ADAM_LR, ADAM_BETA1, ADAM_BETA2, ADAM_EPS = 0.05, 0.9, 0.99, 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)
8record GptConfig {
9 vocab: Num,
10 context: Num,
11 width: Num,
12 heads: Num,
13 layers: Num,
14 hidden: Num,
15 batch: Num,
16}
17
18# the optimizers' constants
19MUON_LR, MUON_MOMENTUM = 0.05, 0.95
20ADAM_LR, ADAM_BETA1, ADAM_BETA2, ADAM_EPS = 0.05, 0.9, 0.99, 1.0e-8
21# the step size falls in a straight line to this fraction of its start
22LR_FLOOR = 0.1
23TEXT, PROMPT = "hello world! " * 5, "hello"
24STEPS, SAMPLES = 30, 24
25INIT_SEED, SAMPLE_SEED = 7, 11
3 random draws
38class Minstd:
39 """Viv's rand(): Park-Miller MINSTD. Using it here makes both programs draw the same numbers."""
40
41 def __init__(self, seed):
42 self.state = math.floor(abs(seed)) % 2147483646 + 1
43
44 def rand(self):
45 self.state = 16807 * self.state % 2147483647
46 return self.state / 2147483647
47
48 def randn(self, rows, cols):
49 out = np.empty((rows, cols))
50 for i in range(rows):
51 for j in range(cols):
52 u1, u2 = self.rand(), self.rand()
53 out[i, j] = math.sqrt(-2.0 * math.log(u1)) * math.cos(2.0 * math.pi * u2)
54 return out
55
56 def choice(self, p):
57 return min(int(np.searchsorted(np.cumsum(p), self.rand(), side="right")), len(p) - 1)
29pi = 3.141592653589793
30
31def randn() {
32 u1, u2 = rand(), rand()
33 return sqrt(-2.0 * log(u1)) * cos(2.0 * pi * u2)
34}
35
36def randn_matrix(rows, cols) {
37 return [[randn() for j to cols] for i to rows]
38}
39
40def choice(p: Arr) {
41 u = rand()
42 cum = 0.0
43 for k to len(p) {
44 cum += p[k]
45 if u < cum {
46 return k
47 }
48 }
49 return len(p) - 1
50}
(refused)
28pi = 3.141592653589793
29
30def randn() {
31 u1, u2 = rand(), rand()
32 return sqrt(-2.0 * log(u1)) * cos(2.0 * pi * u2)
33}
34
35def randn_matrix(rows, cols) {
36 return [[randn() for j to cols] for i to rows]
37}
38
39def choice(p: Arr) {
40 u = rand()
41 cum = 0.0
42 for k to len(p) {
43 cum += p[k]
44 if u < cum {
45 return k
46 }
47 }
48 return len(p) - 1
49}
4 parameters and init
60class Module:
61 """A layer: its own arrays in self.params, their gradients in self.grads, its sublayers in self.children."""
62
63 def __init__(self):
64 self.params, self.grads, self.children = {}, {}, {}
65
66 def named_parameters(self, prefix=""):
67 for name, p in self.params.items():
68 yield prefix + name, p, self.grads.get(name)
69 for name, child in self.children.items():
70 yield from child.named_parameters(f"{prefix}{name}.")
71
72
73def init_matrix(rng, rows, cols, scale=None):
74 if scale is None:
75 scale = 1.0 / math.sqrt(rows)
76 return scale * rng.randn(rows, cols)
77
78
79def flat(a):
80 """(batch, time, features) -> (batch * time, features), for the weight gradients."""
81 return a.reshape(-1, a.shape[-1])
53record Head { Wq: Arr, Wk: Arr, Wv: Arr }
54record Attention { heads: [Head], Wo: Arr }
55record Mlp { up: Arr, down: Arr }
56record Block { attn: Attention, mlp: Mlp }
57record Gpt { tok: Arr, pos: Arr, blocks: [Block], unembed: Arr }
58
59def init(rows, cols) {
60 return (1.0 / sqrt(rows)) * randn_matrix(rows, cols)
61}
62
63def Head(c: GptConfig) -> Head {
64 hd = c.width / c.heads
65 return Head { Wq: init(c.width, hd), Wk: init(c.width, hd), Wv: init(c.width, hd) }
66}
67
68def Block(c: GptConfig) -> Block {
69 attn = Attention { heads: [Head(c) for h to c.heads], Wo: init(c.width, c.width) }
70 mlp = Mlp { up: init(c.width, c.hidden), down: init(c.hidden, c.width) }
71 return Block { attn: attn, mlp: mlp }
72}
73
74def Gpt(c: GptConfig) -> Gpt {
75 tok, pos = randn_matrix(c.vocab, c.width), randn_matrix(c.context, c.width)
76 blocks = [Block(c) for l to c.layers]
77 return Gpt { tok: tok, pos: pos, blocks: blocks, unembed: zeros(c.width, c.vocab) }
78}
79
80def count_into(w: Arr, n: Arr) {
81 n[0] += numel(w)
82}
(refused)
52record Head { Wq: Arr, Wk: Arr, Wv: Arr }
53record Attention { heads: [Head], Wo: Arr }
54record Mlp { up: Arr, down: Arr }
55record Block { attn: Attention, mlp: Mlp }
56record Gpt { tok: Arr, pos: Arr, blocks: [Block], unembed: Arr }
57
58def init(rows, cols) {
59 return (1.0 / sqrt(rows)) * randn_matrix(rows, cols)
60}
61
62def init_head(c: GptConfig) {
63 hd = c.width / c.heads
64 return Head { Wq: init(c.width, hd), Wk: init(c.width, hd), Wv: init(c.width, hd) }
65}
66
67def init_block(c: GptConfig) {
68 attn = Attention { heads: [init_head(c) for h to c.heads], Wo: init(c.width, c.width) }
69 mlp = Mlp { up: init(c.width, c.hidden), down: init(c.hidden, c.width) }
70 return Block { attn: attn, mlp: mlp }
71}
72
73def init_gpt(c: GptConfig) {
74 tok = randn_matrix(c.vocab, c.width)
75 pos = randn_matrix(c.context, c.width)
76 blocks = [init_block(c) for l to c.layers]
77 return Gpt { tok: tok, pos: pos, blocks: blocks, unembed: zeros(c.width, c.vocab) }
78}
79
80def count_into(w: Arr, n: Arr) {
81 n[0] += numel(w)
82}
5 embedding
84class Embedding(Module):
85 def __init__(self, cfg, rng):
86 super().__init__()
87 self.params["tok"] = init_matrix(rng, cfg.vocab, cfg.width, scale=1.0)
88 self.params["pos"] = init_matrix(rng, cfg.context, cfg.width, scale=1.0)
89
90 def forward(self, idx):
91 self.idx = idx
92 return self.params["tok"][idx] + self.params["pos"][: idx.shape[1]]
93
94 def backward(self, dx):
95 dtok = np.zeros_like(self.params["tok"])
96 np.add.at(dtok, self.idx, dx)
97 dpos = np.zeros_like(self.params["pos"])
98 dpos[: dx.shape[1]] = dx.sum(axis=0)
99 self.grads = {"tok": dtok, "pos": dpos}
85def embed(m: Gpt, tokens: Arr) {
86 return take(m.tok, tokens, axis=0) + m.pos[0:shape(tokens)[-1], :]
87}
(refused)
85def embed(m: Gpt, tokens: Arr) {
86 return take(m.tok, tokens, axis=0) + m.pos[0:shape(tokens)[-1], :]
87}
6 attention
102def softmax(x):
103 e = np.exp(x - x.max(axis=-1, keepdims=True))
104 return e / e.sum(axis=-1, keepdims=True)
105
106
107class Head(Module):
108 def __init__(self, cfg, rng):
109 super().__init__()
110 size = cfg.width // cfg.heads
111 for name in ("q", "k", "v"):
112 self.params[name] = init_matrix(rng, cfg.width, size)
113
114 def forward(self, x):
115 p = self.params
116 q, k, v = x @ p["q"], x @ p["k"], x @ p["v"]
117 T = x.shape[1]
118 scale = 1.0 / math.sqrt(q.shape[-1])
119 s = (q @ k.transpose(0, 2, 1)) * scale
120 s = np.where(np.triu(np.ones((T, T), dtype=bool), k=1), MASKED, s)
121 w = softmax(s)
122 self.cache = (x, q, k, v, w, scale)
123 return w @ v
124
125 def backward(self, dy):
126 x, q, k, v, w, scale = self.cache
127 p = self.params
128 dv = w.transpose(0, 2, 1) @ dy
129 dw = dy @ v.transpose(0, 2, 1)
130 ds = w * (dw - (dw * w).sum(axis=-1, keepdims=True)) * scale
131 dq, dk = ds @ k, ds.transpose(0, 2, 1) @ q
132 self.grads = {"q": flat(x).T @ flat(dq), "k": flat(x).T @ flat(dk), "v": flat(x).T @ flat(dv)}
133 return dq @ p["q"].T + dk @ p["k"].T + dv @ p["v"].T
134
135
136class Attention(Module):
137 def __init__(self, cfg, rng):
138 super().__init__()
139 self.heads = [Head(cfg, rng) for _ in range(cfg.heads)]
140 self.children = {f"heads.{i}": h for i, h in enumerate(self.heads)}
141 self.params["out"] = init_matrix(rng, cfg.width, cfg.width)
142
143 def forward(self, x):
144 self.z = np.concatenate([h.forward(x) for h in self.heads], axis=-1)
145 return self.z @ self.params["out"]
146
147 def backward(self, dy):
148 self.grads = {"out": flat(self.z).T @ flat(dy)}
149 dz = np.split(dy @ self.params["out"].T, len(self.heads), axis=-1)
150 return sum(h.backward(d) for h, d in zip(self.heads, dz))
90MASKED = -1.0e30
91
92def softmax_rows(X: Arr) {
93 e = exp(X - amax(X, axis=-1, keepdims=true))
94 return e / sum(e, axis=-1, keepdims=true)
95}
96
97def attention(X: Arr, Wq: Arr, Wk: Arr, Wv: Arr) {
98 T = shape(X)[-2]
99 step project: Q, K, V = X @ Wq, X @ Wk, X @ Wv
100 step scores: S = (Q @ swapaxes(K, -1, -2)) * (1.0 / sqrt(1.0 * shape(Wq)[1]))
101 step mask {
102 future = zeros(T, T)
103 for i to T {
104 future[i, i + 1:T] = 1.0
105 }
106 S = where(future, MASKED, S)
107 }
108 step weights: A = softmax_rows(S)
109 step mix: Y = A @ V
110 return Y
111}
112
113def self_attention(a: Attention, N: Arr) {
114 hd = shape(N)[-1] / len(a.heads)
115 Z = zeros_like(N)
116 for i, h in enumerate(a.heads) {
117 step head: Z[:, :, i * hd:(i + 1) * hd] = attention(N, h.Wq, h.Wk, h.Wv)
118 }
119 step project: Y = Z @ a.Wo
120 return Y
121}
(refused)
90MASKED = -1.0e30
91
92def softmax_rows(X: Arr) {
93 e = exp(X - amax(X, axis=-1, keepdims=true))
94 return e / sum(e, axis=-1, keepdims=true)
95}
96
97phrase "the attention scores of Q against K" -> (Q @ swapaxes(K, -1, -2)) * (1.0 / sqrt(1.0 * shape(Q)[-1]))
98phrase "S masked where future" -> where(future, MASKED, S)
99
100def attention(X: Arr, Wq: Arr, Wk: Arr, Wv: Arr) {
101 T = shape(X)[-2]
102 step project: Q, K, V = X @ Wq, X @ Wk, X @ Wv
103 step scores: S = take the attention scores of Q against K
104 step mask {
105 future = take the causal mask of size T
106 S = take S masked where future
107 }
108 step weights: A = take the softmax of each row of S
109 step mix: Y = A @ V
110 return Y
111}
112
113def self_attention(a: Attention, N: Arr) {
114 hd = shape(N)[-1] / len(a.heads)
115 Z = zeros_like(N)
116 for i, h in enumerate(a.heads) {
117 step head: Z[:, :, i * hd:(i + 1) * hd] = attention(N, h.Wq, h.Wk, h.Wv)
118 }
119 step project: Y = Z @ a.Wo
120 return Y
121}
7 block
153def rmsnorm(x):
154 r = np.sqrt((x * x).mean(axis=-1, keepdims=True) + EPS)
155 return x / r, r
156
157
158def rmsnorm_backward(x, r, dy):
159 dot = (dy * x).sum(axis=-1, keepdims=True)
160 return dy / r - x * dot / (x.shape[-1] * r * r * r)
161
162
163class MLP(Module):
164 def __init__(self, cfg, rng):
165 super().__init__()
166 self.params["up"] = init_matrix(rng, cfg.width, cfg.hidden)
167 self.params["down"] = init_matrix(rng, cfg.hidden, cfg.width)
168
169 def forward(self, x):
170 self.x = x
171 self.u = x @ self.params["up"]
172 self.h = np.maximum(self.u, 0.0) ** 2
173 return self.h @ self.params["down"]
174
175 def backward(self, dy):
176 dh = dy @ self.params["down"].T
177 du = dh * 2.0 * np.maximum(self.u, 0.0)
178 self.grads = {"down": flat(self.h).T @ flat(dy), "up": flat(self.x).T @ flat(du)}
179 return du @ self.params["up"].T
180
181
182class Block(Module):
183 def __init__(self, cfg, rng):
184 super().__init__()
185 self.attn, self.mlp = Attention(cfg, rng), MLP(cfg, rng)
186 self.children = {"attn": self.attn, "mlp": self.mlp}
187
188 def forward(self, x):
189 n1, self.r1 = rmsnorm(x)
190 x1 = x + self.attn.forward(n1)
191 n2, self.r2 = rmsnorm(x1)
192 self.x, self.x1 = x, x1
193 return x1 + self.mlp.forward(n2)
194
195 def backward(self, dy):
196 dx1 = rmsnorm_backward(self.x1, self.r2, self.mlp.backward(dy)) + dy
197 return rmsnorm_backward(self.x, self.r1, self.attn.backward(dx1)) + dx1
124NN_EPS = 1.0e-8
125
126def row_rms(X: Arr) {
127 return sqrt(mean(X * X, axis=-1, keepdims=true) + NN_EPS)
128}
129
130def rmsnorm(X: Arr) {
131 return X / row_rms(X)
132}
133
134def relu2(X: Arr) {
135 R = max(X, 0.0)
136 return R * R
137}
138
139def block(b: Block, X: Arr) {
140 step attention: X = X + self_attention(b.attn, rmsnorm(X))
141 step mlp: X = X + relu2(rmsnorm(X) @ b.mlp.up) @ b.mlp.down
142 return X
143}
(refused)
124NN_EPS = 1.0e-8
125
126def row_rms(X: Arr) {
127 return sqrt(mean(X * X, axis=-1, keepdims=true) + NN_EPS)
128}
129
130def rmsnorm(X: Arr) {
131 return X / row_rms(X)
132}
133
134def relu2(X: Arr) {
135 R = max(X, 0.0)
136 return R * R
137}
138
139def block(b: Block, X: Arr) {
140 step attention: X = X + self_attention(b.attn, rmsnorm(X))
141 step mlp: X = X + relu2(rmsnorm(X) @ b.mlp.up) @ b.mlp.down
142 return X
143}
8 forward
200class GPT(Module):
201 def __init__(self, cfg, rng):
202 super().__init__()
203 self.cfg = cfg
204 self.embed = Embedding(cfg, rng)
205 self.blocks = [Block(cfg, rng) for _ in range(cfg.layers)]
206 self.children = {"embed": self.embed, **{f"blocks.{i}": b for i, b in enumerate(self.blocks)}}
207 # the output map starts at zero, so the first prediction is uniform
208 self.params["unembed"] = np.zeros((cfg.width, cfg.vocab))
209
210 def forward(self, idx):
211 x = self.embed.forward(idx)
212 for block in self.blocks:
213 x = block.forward(x)
214 self.x = x
215 self.nf, self.rf = rmsnorm(x)
216 return self.nf @ self.params["unembed"]
146def gpt_forward(m: Gpt, tokens: Arr) {
147 step embed: X = embed(m, tokens)
148 for b in m.blocks {
149 step block: X = block(b, X)
150 }
151 step logits: logits = rmsnorm(X) @ m.unembed
152 return logits
153}
(refused)
146def gpt_forward(m: Gpt, tokens: Arr) {
147 step embed: X = embed(m, tokens)
148 for b in m.blocks {
149 step block: X = block(b, X)
150 }
151 step logits: logits = rmsnorm(X) @ m.unembed
152 return logits
153}
9 loss
219 def loss(self, logits, targets):
220 """The mean cross-entropy of the next token, and its gradient with respect to the logits."""
221 shape = logits.shape
222 logits, targets = flat(logits), targets.reshape(-1)
223 rows = np.arange(len(targets))
224 top = logits.max(axis=1, keepdims=True)
225 e = np.exp(logits - top)
226 total = e.sum(axis=1)
227 loss = np.mean(np.log(total) + top[:, 0] - logits[rows, targets])
228 dlogits = e / total[:, None]
229 dlogits[rows, targets] -= 1.0
230 return loss, (dlogits / len(targets)).reshape(shape)
156def cross_entropy(logits: Arr, target: Arr) {
157 classes = numel(logits) / numel(target)
158 top = amax(logits, axis=-1, keepdims=true)
159 logp = logits - top - log(sum(exp(logits - top), axis=-1, keepdims=true))
160 return -mean(take(logp, arange(numel(target)) * classes + ravel(target)))
161}
162
163def sequence_loss(m: Gpt, x: Arr, y: Arr) {
164 return cross_entropy(gpt_forward(m, x), y)
165}
(refused)
156def cross_entropy(logits: Arr, target: Arr) {
157 classes = numel(logits) / numel(target)
158 top = amax(logits, axis=-1, keepdims=true)
159 logp = logits - top - log(sum(exp(logits - top), axis=-1, keepdims=true))
160 return -mean(take(logp, arange(numel(target)) * classes + ravel(target)))
161}
162
163def sequence_loss(m: Gpt, x: Arr, y: Arr) {
164 return take the cross-entropy of gpt_forward(m, x) against y
165}
10 backward
233 def backward(self, dlogits):
234 self.grads = {"unembed": flat(self.nf).T @ flat(dlogits)}
235 dx = rmsnorm_backward(self.x, self.rf, dlogits @ self.params["unembed"].T)
236 for block in reversed(self.blocks):
237 dx = block.backward(dx)
238 self.embed.backward(dx)
168# nothing: `train` calls `value_and_grad(sequence_loss, model, ...)`, and the
169# compiler writes the backward pass from the forward pass above
(refused)
168# nothing: `train` calls `value_and_grad(sequence_loss, model, ...)`, and the
169# compiler writes the backward pass from the forward pass above
11 optimizer
241def newton_schulz(g, steps=5, a=3.4445, b=-4.775, c=2.0315):
242 """A matrix with g's singular vectors and every singular value near 1 (Muon's direction)."""
243 if g.shape[0] > g.shape[1]:
244 return newton_schulz(g.T, steps, a, b, c).T
245 norm = np.linalg.norm(g)
246 x = g / norm if norm else g
247 for _ in range(steps):
248 s = x @ x.T
249 x = a * x + b * (s @ x) + c * ((s @ s) @ x)
250 return x
251
252
253class MuonAdam:
254 """Muon on the matrices inside the blocks, Adam on the embeddings and the output map."""
255
256 def __init__(self):
257 self.state = {}
258 self.taken = 0
259
260 def step(self, model, lr_scale):
261 self.taken += 1
262 for name, p, g in model.named_parameters():
263 s = self.state.setdefault(name, {"m": 0.0, "v": 0.0})
264 if name.startswith("blocks."):
265 s["m"] = MUON_MOMENTUM * s["m"] + g
266 rows, cols = p.shape
267 p -= MUON_LR * lr_scale * math.sqrt(max(1.0, rows / cols)) * newton_schulz(s["m"])
268 else:
269 s["m"] = ADAM_BETA1 * s["m"] + (1.0 - ADAM_BETA1) * g
270 s["v"] = ADAM_BETA2 * s["v"] + (1.0 - ADAM_BETA2) * g * g
271 mean = s["m"] / (1.0 - ADAM_BETA1**self.taken)
272 spread = np.sqrt(s["v"] / (1.0 - ADAM_BETA2**self.taken))
273 p -= ADAM_LR * lr_scale * mean / (spread + ADAM_EPS)
172#| Muon's, from the nanogpt speedrun. Faster off the mark, and the fixed point
173#| is not 1.
174MUON_A, MUON_B, MUON_C = 3.4445, -4.775, 2.0315
175#| Rounds of it Muon runs, as the speedrun does.
176MUON_STEPS = 5
177
178# the matrix's polar factor, approximately: its singular values pushed toward 1
179def newton_schulz(A: Arr) {
180 f = sqrt(sum(A * A))
181 X = A if f == 0.0 else (1.0 / f) * A
182 for k to MUON_STEPS {
183 S = X @ transpose(X)
184 X = MUON_A * X + MUON_B * (S @ X) + MUON_C * ((S @ S) @ X)
185 }
186 return X
187}
188
189def muon_matrix_direction(M: Arr) {
190 if len(M) <= shape(M)[1] {
191 return newton_schulz(M)
192 }
193 return transpose(newton_schulz(transpose(M)))
194}
195
196def muon_update(w: Arr, g: Arr, mom: Arr, lr_scale) {
197 mom[:] = MUON_MOMENTUM * mom + g
198 rows, cols = shape(w)
199 size = MUON_LR * lr_scale * sqrt(max(1.0, rows / cols))
200 w[:] = w - size * muon_matrix_direction(mom)
201}
202
203def adam_update(w: Arr, g: Arr, m: Arr, v: Arr, taken, lr_scale) {
204 m[:] = ADAM_BETA1 * m + (1.0 - ADAM_BETA1) * g
205 v[:] = ADAM_BETA2 * v + (1.0 - ADAM_BETA2) * g * g
206 mean = m / (1.0 - pow(ADAM_BETA1, taken))
207 spread = sqrt(v / (1.0 - pow(ADAM_BETA2, taken)))
208 w[:] = w - ADAM_LR * lr_scale * mean / (spread + ADAM_EPS)
209}
210
211def optimizer_step(model: Gpt, g: Gpt, m: Gpt, v: Gpt, taken, lr_scale) {
212 step muon: for_leaves(muon_update, model.blocks, g.blocks, m.blocks, lr_scale)
213 step adam {
214 adam_update(model.tok, g.tok, m.tok, v.tok, taken, lr_scale)
215 adam_update(model.pos, g.pos, m.pos, v.pos, taken, lr_scale)
216 adam_update(model.unembed, g.unembed, m.unembed, v.unembed, taken, lr_scale)
217 }
218}
(refused)
172#| Muon's, from the nanogpt speedrun. Faster off the mark, and the fixed point
173#| is not 1.
174MUON_A, MUON_B, MUON_C = 3.4445, -4.775, 2.0315
175#| Rounds of it Muon runs, as the speedrun does.
176MUON_STEPS = 5
177
178# the matrix's polar factor, approximately: its singular values pushed toward 1
179def newton_schulz(A: Arr) {
180 f = sqrt(sum(A * A))
181 X = A if f == 0.0 else (1.0 / f) * A
182 for k to MUON_STEPS {
183 S = X @ transpose(X)
184 X = MUON_A * X + MUON_B * (S @ X) + MUON_C * ((S @ S) @ X)
185 }
186 return X
187}
188
189def muon_matrix_direction(M: Arr) {
190 if len(M) <= shape(M)[1] {
191 return newton_schulz(M)
192 }
193 return transpose(newton_schulz(transpose(M)))
194}
195
196phrase "the step-size scale at step t of steps" -> 1.0 - (1.0 - LR_FLOOR) * t / steps
197
198def muon_update(w: Arr, g: Arr, mom: Arr, lr_scale) {
199 mom[:] = MUON_MOMENTUM * mom + g
200 rows, cols = shape(w)
201 size = MUON_LR * lr_scale * sqrt(max(1.0, rows / cols))
202 w[:] = w - size * muon_matrix_direction(mom)
203}
204
205def adam_update(w: Arr, g: Arr, m: Arr, v: Arr, taken, lr_scale) {
206 m[:] = ADAM_BETA1 * m + (1.0 - ADAM_BETA1) * g
207 v[:] = ADAM_BETA2 * v + (1.0 - ADAM_BETA2) * g * g
208 mean = m / (1.0 - pow(ADAM_BETA1, taken))
209 spread = sqrt(v / (1.0 - pow(ADAM_BETA2, taken)))
210 w[:] = w - ADAM_LR * lr_scale * mean / (spread + ADAM_EPS)
211}
212
213def optimizer_step(model: Gpt, g: Gpt, m: Gpt, v: Gpt, taken, lr_scale) {
214 step muon: for_leaves(muon_update, model.blocks, g.blocks, m.blocks, lr_scale)
215 step adam {
216 adam_update(model.tok, g.tok, m.tok, v.tok, taken, lr_scale)
217 adam_update(model.pos, g.pos, m.pos, v.pos, taken, lr_scale)
218 adam_update(model.unembed, g.unembed, m.unembed, v.unembed, taken, lr_scale)
219 }
220}
12 training loop
276def get_batch(tokens, cfg, step, shift):
277 windows = len(tokens) - cfg.context
278 starts = [(step * cfg.batch + b) % windows for b in range(cfg.batch)]
279 return np.array([tokens[s + shift : s + shift + cfg.context] for s in starts])
280
281
282def train(model, tokens, steps):
283 opt = MuonAdam()
284 losses = []
285 for step in range(steps):
286 x, y = get_batch(tokens, model.cfg, step, 0), get_batch(tokens, model.cfg, step, 1)
287 loss, dlogits = model.loss(model.forward(x), y)
288 model.backward(dlogits)
289 opt.step(model, lr_scale=1.0 - (1.0 - LR_FLOOR) * step / steps)
290 losses.append(loss)
291 print(f"step {step + 1}: loss {loss:.6f}")
292 return losses
221def batch(tokens: Arr, c: GptConfig, t, shift) {
222 starts = [mod(t * c.batch + b, len(tokens) - c.context) + shift for b to c.batch]
223 return [tokens[s:s + c.context] for s in starts]
224}
225
226def train(model: Gpt, c: GptConfig, tokens: Arr, steps) {
227 losses = zeros(steps)
228 m, v = zeros_like(model), zeros_like(model)
229 for t to steps {
230 step batch: xs, ys = batch(tokens, c, t, 0), batch(tokens, c, t, 1)
231 step forward_backward {
232 loss, grad = value_and_grad(sequence_loss, model, xs, ys)
233 losses[t] = loss
234 }
235 step measure: gnorm = sqrt(sum(grad * grad))
236 step update: optimizer_step(model, grad, m, v, t + 1, 1.0 - (1.0 - LR_FLOOR) * t / steps)
237 }
238 return losses
239}
(refused)
223def batch(tokens: Arr, c: GptConfig, t, shift) {
224 starts = [mod(t * c.batch + b, len(tokens) - c.context) + shift for b to c.batch]
225 return [tokens[s:s + c.context] for s in starts]
226}
227
228def train(model: Gpt, c: GptConfig, tokens: Arr, steps) {
229 losses = zeros(steps)
230 m, v = zeros_like(model), zeros_like(model)
231 for t to steps {
232 step batch: xs, ys = batch(tokens, c, t, 0), batch(tokens, c, t, 1)
233 step forward_backward {
234 loss, grad = value_and_grad(sequence_loss, model, xs, ys)
235 losses[t] = loss
236 }
237 step measure: gnorm = sqrt(sum(grad * grad))
238 step update {
239 lr_scale = take the step-size scale at step t of steps
240 optimizer_step(model, grad, m, v, t + 1, lr_scale)
241 }
242 }
243 return losses
244}
13 sampling
295def generate(model, prompt, count, rng):
296 out = list(prompt)
297 for _ in range(count):
298 logits = model.forward(np.array([out[-model.cfg.context :]]))
299 out.append(rng.choice(softmax(logits[0, -1])))
300 return out
242def sample(m: Gpt, prompt: Arr, count, seed) {
243 srand(seed)
244 out = prompt
245 for k to count {
246 step context: ctx = out[max(0, len(out) - len(m.pos)):len(out)]
247 step probs: p = softmax_rows(gpt_forward(m, reshape(ctx, 1, len(ctx)))[0, len(ctx) - 1, :])
248 step draw: out = append(out, choice(p))
249 }
250 return out
251}
(refused)
247def sample(m: Gpt, prompt: Arr, count, seed) {
248 srand(seed)
249 out = prompt
250 for k to count {
251 step context: ctx = take the last len(m.pos) elements of out
252 step probs: p = softmax_rows(gpt_forward(m, reshape(ctx, 1, len(ctx)))[0, len(ctx) - 1, :])
253 step draw: out = append(out, choice(p))
254 }
255 return out
256}
14 text as tokens
303vocab = sorted(set(TEXT))
304stoi = {c: i for i, c in enumerate(vocab)}
305
306
307def encode(text):
308 return [stoi[c] for c in text]
309
310
311def decode(ids):
312 return "".join(vocab[i] for i in ids)
254vocab = unique(TEXT)
255stoi = {c: i for i, c in enumerate(vocab)}
256encode = fn(text) -> [stoi[c] for c in text]
257decode = fn(ids) -> join("", [vocab[i] for i in ids])
(refused)
259def encode(text: Str, vocab: [Str]) {
260 stoi = {c: i for i, c in enumerate(vocab)}
261 return [stoi[c] for c in text]
262}
263
264def decode(tokens: Arr, vocab: [Str]) {
265 return join("", [vocab[i] for i in tokens])
266}
15 run
315cfg = Config(vocab=len(vocab))
316model = GPT(cfg, Minstd(INIT_SEED))
317count = sum(p.size for _, p, _ in model.named_parameters())
318print(f"{count} parameters, {len(vocab)} characters")
319losses = train(model, np.array(encode(TEXT)), STEPS)
320written = decode(generate(model, encode(PROMPT), SAMPLES, Minstd(SAMPLE_SEED)))
321print(written)
260cfg = GptConfig { vocab: len(vocab) }
261srand(INIT_SEED)
262model = Gpt(cfg)
263count = [0.0]
264for_leaves(count_into, model, count)
265# story_gpt prints the four numbers bare; here every loss is printed, to six
266# decimals, so the Python run can be compared with it line by line
267print(f"{count[0]:d} parameters, {len(vocab):d} characters")
268step train: losses = train(model, cfg, encode(TEXT), STEPS)
269for t to STEPS {
270 print(f"step {t + 1:d}: loss {losses[t]:.6f}")
271}
272step generate: written = decode(sample(model, encode(PROMPT), SAMPLES, SAMPLE_SEED))
273print(written)
(refused)
269vocab = unique(TEXT)
270cfg = GptConfig {
271 vocab: len(vocab),
272 context: 8,
273 width: 16,
274 heads: 2,
275 layers: 2,
276 hidden: 64,
277 batch: 4,
278}
279srand(INIT_SEED)
280model = init_gpt(cfg)
281count = [0.0]
282for_leaves(count_into, model, count)
283# story_gpt prints the four numbers bare; here every loss is printed, to six
284# decimals, so the Python run can be compared with it line by line
285print(f"{count[0]:d} parameters, {len(vocab):d} characters")
286step train: losses = train(model, cfg, encode(TEXT, vocab), STEPS)
287for t to STEPS {
288 print(f"step {t + 1:d}: loss {losses[t]:.6f}")
289}
290step generate: written = decode(sample(model, encode(PROMPT, vocab), SAMPLES, SAMPLE_SEED), vocab)
291print(written)
16 checks
324assert count == 6560 and len(vocab) == 9
325assert abs(losses[0] - math.log(9.0)) <= 1e-15
326assert abs(losses[-1] - 0.199118) <= 1e-6
327assert written == "hello world! hello world! hel"
276test "the model has 6,560 parameters, and 9 characters make its vocabulary" {
277 assert_true(count[0] == 6560 and len(vocab) == 9)
278}
279
280test "the first loss is log 9 exactly, because the output map starts at zero" {
281 assert_near(losses[0], log(9.0), 1.0e-15)
282}
283
284test "thirty steps take the loss from 2.197 to 0.1991" {
285 assert_near(losses[STEPS - 1], 0.199118, 1.0e-6)
286}
287
288test "from the prompt hello, the model writes the line on" {
289 assert_true(written == "hello world! hello world! hel")
290}
(refused)
294test "the model has 6,560 parameters, and 9 characters make its vocabulary" {
295 assert_true(count[0] == 6560 and len(vocab) == 9)
296}
297
298test "the first loss is log 9 exactly, because the output map starts at zero" {
299 assert_near(losses[0], log(9.0), 1.0e-15)
300}
301
302test "thirty steps take the loss from 2.197 to 0.1991" {
303 assert_near(losses[STEPS - 1], 0.199118, 1.0e-6)
304}
305
306test "from the prompt hello, the model writes the line on" {
307 assert_true(written == "hello world! hello world! hel")
308}

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.py

6560 parameters, 9 characters
step 1: loss 2.197225
step 2: loss 2.102097
step 3: loss 1.879860
step 4: loss 1.575851
step 5: loss 1.150941
step 6: loss 1.198318
step 7: loss 0.840376
step 8: loss 0.673002
step 9: loss 0.686226
step 10: loss 0.541481
step 11: loss 0.403199
step 12: loss 0.305271
step 13: loss 0.563067
step 14: loss 0.306272
step 15: loss 0.314534
step 16: loss 0.195234
step 17: loss 0.539175
step 18: loss 0.290064
step 19: loss 0.191980
step 20: loss 0.334400
step 21: loss 0.255869
step 22: loss 0.229073
step 23: loss 0.151354
step 24: loss 0.341850
step 25: loss 0.196461
step 26: loss 0.141672
step 27: loss 0.266627
step 28: loss 0.179300
step 29: loss 0.143552
step 30: loss 0.199118
hello world! hello world! hel

$ viv run tiny_gpt.viv

6560 parameters, 9 characters
step 1: loss 2.197225
step 2: loss 2.102097
step 3: loss 1.879860
step 4: loss 1.575851
step 5: loss 1.150941
step 6: loss 1.198318
step 7: loss 0.840376
step 8: loss 0.673002
step 9: loss 0.686226
step 10: loss 0.541481
step 11: loss 0.403199
step 12: loss 0.305271
step 13: loss 0.563067
step 14: loss 0.306272
step 15: loss 0.314534
step 16: loss 0.195234
step 17: loss 0.539175
step 18: loss 0.290064
step 19: loss 0.191980
step 20: loss 0.334400
step 21: loss 0.255869
step 22: loss 0.229073
step 23: loss 0.151354
step 24: loss 0.341850
step 25: loss 0.196461
step 26: loss 0.141672
step 27: loss 0.266627
step 28: loss 0.179300
step 29: loss 0.143552
step 30: loss 0.199118
hello world! hello world! hel
PASS the model has 6,560 parameters, and 9 characters make its vocabulary
PASS the first loss is log 9 exactly, because the output map starts at zero
PASS thirty steps take the loss from 2.197 to 0.1991
PASS from the prompt hello, the model writes the line on