Module 7 · Autograd = automatic chain rule · train2.py · ~30 min

7.5 backward(): let the graph do calculus

microgpt has 4,192 knobs. To learn, it must know, for every single knob, “if I turn you a tiny bit, does the loss go up or down, and by how much?” Nudging each knob one at a time would mean 4,192 extra runs per step. Today you write the 14 lines that get all 4,192 answers in one backwards sweep.

1Watch

During the forward pass every Value remembers two things: its children (the values it was made from) and its local gradients, the “exchange rates” from each child to itself. For e = a × b, nudging a by 1 moves e by b, so the local gradient for a is b.data.

backward() does three things:

Sort the graph so every node comes after its children. That’s the topological sort you wrote in lesson 3.3 (socks before shoes).
Set the loss’s own gradient to 1.
Walk the sorted list backwards. Each node passes its gradient to each child: child.grad += local_grad × node.grad. That’s the chain rule: multiply the exchange rates along the path.

Two details carry the whole idea. Backwards order guarantees a node has received gradient from everyone who uses it before it passes anything on. And +=, not =: when a value is used in two places, both paths push on the loss, so their effects add up. In the widget, press “try = instead of +=” and step through: e ends at −6 instead of −2, and a and b inherit the mistake.

2Explore

topo order (walked right → left):abecdL

Forward pass done: every node knows its value. Every grad starts at 0.

step 1 / 6

3Build

The Value class below is copied from microgpt, with two lines of backward() removed:

In build_topo: add v to topo after all of its children have been visited. That means after the for loop over children, at the same indentation as that for: not before it, and not inside it.
In the backwards loop: grow the child’s gradient by local_grad * v.grad. Think hard about = vs +=.

The last check compares your gradients against plain nudging (the numerical derivative) on a messy formula with exp, log and relu. If they agree, your backward() is correct, full stop.

Stuck? In the code-along video, backward() starts at 2:30 and the step-by-step trace at 3:43. The earlier part rebuilds the Value class from lesson 7.4.

Stuck? Watch the code being written, line by line (6 min)

import math

class Value:
    __slots__ = ('data', 'grad', '_children', '_local_grads') # Python optimization for memory usage

    def __init__(self, data, children=(), local_grads=()):
        self.data = data                # scalar value of this node calculated during forward pass
        self.grad = 0                   # derivative of the loss w.r.t. this node, calculated in backward pass
        self._children = children       # children of this node in the computation graph
        self._local_grads = local_grads # local derivative of this node w.r.t. its children

    def __add__(self, other):
        other = other if isinstance(other, Value) else Value(other)
        return Value(self.data + other.data, (self, other), (1, 1))

    def __mul__(self, other):
        other = other if isinstance(other, Value) else Value(other)
        return Value(self.data * other.data, (self, other), (other.data, self.data))

    def __pow__(self, other): return Value(self.data**other, (self,), (other * self.data**(other-1),))
    def log(self): return Value(math.log(self.data), (self,), (1/self.data,))
    def exp(self): return Value(math.exp(self.data), (self,), (math.exp(self.data),))
    def relu(self): return Value(max(0, self.data), (self,), (float(self.data > 0),))
    def __neg__(self): return self * -1
    def __radd__(self, other): return self + other
    def __sub__(self, other): return self + (-other)
    def __rsub__(self, other): return other + (-self)
    def __rmul__(self, other): return self * other
    def __truediv__(self, other): return self * other**-1
    def __rtruediv__(self, other): return other * self**-1


    def backward(self):
        # Step 1: put every node in an order where a node comes AFTER
        # all of its children (topological order - socks before shoes).
        topo = []
        visited = set()
        def build_topo(v):
            if v not in visited:
                visited.add(v)
                for child in v._children:
                    build_topo(child)
                # TODO: v is finished only after all its children - add it to topo
        build_topo(self)

        # Step 2: the loss moves 1:1 with itself
        self.grad = 1

        # Step 3: walk the order BACKWARDS, so each node is complete before
        # it hands its gradient on to its children.
        for v in reversed(topo):
            for child, local_grad in zip(v._children, v._local_grads):
                pass  # TODO: chain rule - child's grad grows by local_grad * v.grad


# Try it on the example from the video:
a, b, c = Value(2.0), Value(-3.0), Value(10.0)
e = a * b      # -6
d = e + c      # 4
L = d * e      # -24  (e is used twice: by d and by L)
L.backward()
print('L =', L.data)
print('a.grad =', a.grad, ' b.grad =', b.grad, ' c.grad =', c.grad, ' e.grad =', e.grad)

4Check yourself

1. Why does backward() walk the topological order in REVERSE?

2. Predict: in the film graph (a = 2, b = −3, c = 10, e = a × b, d = e + c, L = d × e), what is e.grad after backward()?

3. What breaks if you write child.grad = … instead of +=?

4. Predict: e = a × b with a = 2, b = −3, and e.grad = −2. What is a.grad?

5Unlocked in microgpt

backward() (lines 59 to 72) is now complete: lines 60 to 68 are your topological sort from lesson 3.3, and lines 69 to 72 are the chain rule you just wrote. Line 172, loss.backward(), runs it once per training step, across about 60,000 nodes for a typical name. PyTorch and JAX (the tools real AI labs use to train big models) do exactly this, on big grids of numbers instead of single numbers.

Progress is saved in this browser. Sign in to keep it across devices.

microgpt.py94 / 175 lines learned

1"""

2The most atomic way to train and run inference for a GPT in pure, dependency-free Python.

3This file is the complete algorithm.

4Everything else is just efficiency.

6@karpathy

7"""

9import os # os.path.exists

10import math # math.log, math.exp

11import random # random.seed, random.choices, random.gauss, random.shuffle

12random.seed(42) # Let there be order among chaos

14# Let there be a Dataset `docs`: list[str] of documents (e.g. a list of names)

15if not os.path.exists('input.txt'):

16 import urllib.request

17 names_url = 'https://raw.githubusercontent.com/karpathy/makemore/988aa59/names.txt'

18 urllib.request.urlretrieve(names_url, 'input.txt')

19docs = [line.strip() for line in open('input.txt') if line.strip()]

20random.shuffle(docs)

21print(f"num docs: {len(docs)}")

23# Let there be a Tokenizer to translate strings to sequences of integers ("tokens") and back

24uchars = sorted(set(''.join(docs))) # unique characters in the dataset become token ids 0..n-1

25BOS = len(uchars) # token id for a special Beginning of Sequence (BOS) token

26vocab_size = len(uchars) + 1 # total number of unique tokens, +1 is for BOS

27print(f"vocab size: {vocab_size}")

29# Let there be Autograd to recursively apply the chain rule through a computation graph

30class Value:

31 __slots__ = ('data', 'grad', '_children', '_local_grads') # Python optimization for memory usage

33 def __init__(self, data, children=(), local_grads=()):

34 self.data = data # scalar value of this node calculated during forward pass

35 self.grad = 0 # derivative of the loss w.r.t. this node, calculated in backward pass

36 self._children = children # children of this node in the computation graph

37 self._local_grads = local_grads # local derivative of this node w.r.t. its children

39 def __add__(self, other):

40 other = other if isinstance(other, Value) else Value(other)

41 return Value(self.data + other.data, (self, other), (1, 1))

43 def __mul__(self, other):

44 other = other if isinstance(other, Value) else Value(other)

45 return Value(self.data * other.data, (self, other), (other.data, self.data))

47 def __pow__(self, other): return Value(self.data**other, (self,), (other * self.data**(other-1),))

48 def log(self): return Value(math.log(self.data), (self,), (1/self.data,))

49 def exp(self): return Value(math.exp(self.data), (self,), (math.exp(self.data),))

50 def relu(self): return Value(max(0, self.data), (self,), (float(self.data > 0),))

51 def __neg__(self): return self * -1

52 def __radd__(self, other): return self + other

53 def __sub__(self, other): return self + (-other)

54 def __rsub__(self, other): return other + (-self)

55 def __rmul__(self, other): return self * other

56 def __truediv__(self, other): return self * other**-1

57 def __rtruediv__(self, other): return other * self**-1

59 def backward(self):

60 topo = []

61 visited = set()

62 def build_topo(v):

63 if v not in visited:

64 visited.add(v)

65 for child in v._children:

66 build_topo(child)

67 topo.append(v)

68 build_topo(self)

69 self.grad = 1

70 for v in reversed(topo):

71 for child, local_grad in zip(v._children, v._local_grads):

72 child.grad += local_grad * v.grad

74# Initialize the parameters, to store the knowledge of the model

75n_layer = 1 # depth of the transformer neural network (number of layers)

76n_embd = 16 # width of the network (embedding dimension)

77block_size = 16 # maximum context length of the attention window (note: the longest name is 15 characters)

78n_head = 4 # number of attention heads

79head_dim = n_embd // n_head # derived dimension of each head

80matrix = lambda nout, nin, std=0.08: [[Value(random.gauss(0, std)) for _ in range(nin)] for _ in range(nout)]

81state_dict = {'wte': matrix(vocab_size, n_embd), 'wpe': matrix(block_size, n_embd), 'lm_head': matrix(vocab_size, n_embd)}

82for i in range(n_layer):

83 state_dict[f'layer{i}.attn_wq'] = matrix(n_embd, n_embd)

84 state_dict[f'layer{i}.attn_wk'] = matrix(n_embd, n_embd)

85 state_dict[f'layer{i}.attn_wv'] = matrix(n_embd, n_embd)

86 state_dict[f'layer{i}.attn_wo'] = matrix(n_embd, n_embd)

87 state_dict[f'layer{i}.mlp_fc1'] = matrix(4 * n_embd, n_embd)

88 state_dict[f'layer{i}.mlp_fc2'] = matrix(n_embd, 4 * n_embd)

89params = [p for mat in state_dict.values() for row in mat for p in row] # flatten params into a single list[Value]

90print(f"num params: {len(params)}")

92# Define the model architecture: a function mapping tokens and parameters to logits over what comes next

93# Follow GPT-2, blessed among the GPTs, with minor differences: layernorm -> rmsnorm, no biases, GeLU -> ReLU

94def linear(x, w):

95 return [sum(wi * xi for wi, xi in zip(wo, x)) for wo in w]

97def softmax(logits):

98 max_val = max(val.data for val in logits)

99 exps = [(val - max_val).exp() for val in logits]

100 total = sum(exps)

101 return [e / total for e in exps]

102

103def rmsnorm(x):

104 ms = sum(xi * xi for xi in x) / len(x)

105 scale = (ms + 1e-5) ** -0.5

106 return [xi * scale for xi in x]

107

108def gpt(token_id, pos_id, keys, values):

109 tok_emb = state_dict['wte'][token_id] # token embedding

110 pos_emb = state_dict['wpe'][pos_id] # position embedding

111 x = [t + p for t, p in zip(tok_emb, pos_emb)] # joint token and position embedding

112 x = rmsnorm(x) # note: not redundant due to backward pass via the residual connection

113

114 for li in range(n_layer):

115 # 1) Multi-head Attention block

116 x_residual = x

117 x = rmsnorm(x)

118 q = linear(x, state_dict[f'layer{li}.attn_wq'])

119 k = linear(x, state_dict[f'layer{li}.attn_wk'])

120 v = linear(x, state_dict[f'layer{li}.attn_wv'])

121 keys[li].append(k)

122 values[li].append(v)

123 x_attn = []

124 for h in range(n_head):

125 hs = h * head_dim

126 q_h = q[hs:hs+head_dim]

127 k_h = [ki[hs:hs+head_dim] for ki in keys[li]]

128 v_h = [vi[hs:hs+head_dim] for vi in values[li]]

129 attn_logits = [sum(q_h[j] * k_h[t][j] for j in range(head_dim)) / head_dim**0.5 for t in range(len(k_h))]

130 attn_weights = softmax(attn_logits)

131 head_out = [sum(attn_weights[t] * v_h[t][j] for t in range(len(v_h))) for j in range(head_dim)]

132 x_attn.extend(head_out)

133 x = linear(x_attn, state_dict[f'layer{li}.attn_wo'])

134 x = [a + b for a, b in zip(x, x_residual)]

135 # 2) MLP block

136 x_residual = x

137 x = rmsnorm(x)

138 x = linear(x, state_dict[f'layer{li}.mlp_fc1'])

139 x = [xi.relu() for xi in x]

140 x = linear(x, state_dict[f'layer{li}.mlp_fc2'])

141 x = [a + b for a, b in zip(x, x_residual)]

142

143 logits = linear(x, state_dict['lm_head'])

144 return logits

145

146# Let there be Adam, the blessed optimizer and its buffers

147learning_rate, beta1, beta2, eps_adam = 0.01, 0.85, 0.99, 1e-8

148m = [0.0] * len(params) # first moment buffer

149v = [0.0] * len(params) # second moment buffer

150

151# Repeat in sequence

152num_steps = 1000 # number of training steps

153for step in range(num_steps):

154

155 # Take single document, tokenize it, surround it with BOS special token on both sides

156 doc = docs[step % len(docs)]

157 tokens = [BOS] + [uchars.index(ch) for ch in doc] + [BOS]

158 n = min(block_size, len(tokens) - 1)

159

160 # Forward the token sequence through the model, building up the computation graph all the way to the loss

161 keys, values = [[] for _ in range(n_layer)], [[] for _ in range(n_layer)]

162 losses = []

163 for pos_id in range(n):

164 token_id, target_id = tokens[pos_id], tokens[pos_id + 1]

165 logits = gpt(token_id, pos_id, keys, values)

166 probs = softmax(logits)

167 loss_t = -probs[target_id].log()

168 losses.append(loss_t)

169 loss = (1 / n) * sum(losses) # final average loss over the document sequence. May yours be low.

170

171 # Backward the loss, calculating the gradients with respect to all model parameters

172 loss.backward()

173

174 # Adam optimizer update: update the model parameters based on the corresponding gradients

175 lr_t = learning_rate * (1 - step / num_steps) # linear learning rate decay

176 for i, p in enumerate(params):

177 m[i] = beta1 * m[i] + (1 - beta1) * p.grad

178 v[i] = beta2 * v[i] + (1 - beta2) * p.grad ** 2

179 m_hat = m[i] / (1 - beta1 ** (step + 1))

180 v_hat = v[i] / (1 - beta2 ** (step + 1))

181 p.data -= lr_t * m_hat / (v_hat ** 0.5 + eps_adam)

182 p.grad = 0

183

184 print(f"step {step+1:4d} / {num_steps:4d} | loss {loss.data:.4f}", end='\r')

185

186# Inference: may the model babble back to us

187temperature = 0.5 # in (0, 1], control the "creativity" of generated text, low to high

188print("\n--- inference (new, hallucinated names) ---")

189for sample_idx in range(20):

190 keys, values = [[] for _ in range(n_layer)], [[] for _ in range(n_layer)]

191 token_id = BOS

192 sample = []

193 for pos_id in range(block_size):

194 logits = gpt(token_id, pos_id, keys, values)

195 probs = softmax([l / temperature for l in logits])

196 token_id = random.choices(range(vocab_size), weights=[p.data for p in probs])[0]

197 if token_id == BOS:

198 break

199 sample.append(uchars[token_id])

200 print(f"sample {sample_idx+1:2d}: {''.join(sample)}")

this lessonlearnedKarpathy’s original

Was this lesson clear?

End of Module 7: check what you learnedNext available lesson → 8.1 Position embeddings