Train It in Python

Your TinyGPT is assembled, but its 26,771 weights are random. Before training it spreads its guess for the next character almost evenly over all 19 characters, and the loss sits near 3. Training changes that with the loop you know from the PyTorch chapter: take a batch, score it, measure the loss, and move every weight a little. 400 passes take under 2 seconds.

The same loop on a tiny model: here it teaches nn.Linear(1, 1) that y is twice x.

xs = torch.tensor([[1.0], [2.0], [3.0]])
ys = torch.tensor([[2.0], [4.0], [6.0]])
line = nn.Linear(1, 1)
mse = nn.MSELoss()
opt = torch.optim.SGD(line.parameters(), lr=0.05)
for step in range(200):
    err = mse(line(xs), ys)
    opt.zero_grad()
    err.backward()
    opt.step()

Three calls do the learning. optimizer.zero_grad() empties the slopes left from the last pass, because PyTorch adds new slopes to old ones. loss.backward() works out the slope (the gradient) of every weight, back through the blocks, the heads and the embeddings: the calculus you wrote by hand in numpy. optimizer.step() moves each weight a small step against its slope; the starting code made optimizer an Adam with lr=0.01.

Remember the shapes from the bigram lesson: the model gives 19 scores for each of the 8 positions of the 32 windows, shape [32, 8, 19], and cross-entropy wants one row of scores per question. logits.reshape(-1, vocab_size) gives 256 rows of 19 scores, and y.reshape(-1) the 256 right answers. The loop prints the loss every 100 steps; then a check measures the loss on the whole text, compares it with the counting model, and asks how likely a name is to end after three starts. The run prints:

step 0 loss: 2.986
step 100 loss: 1.446
step 200 loss: 1.431
step 300 loss: 1.392
loss on the whole text: 1.38
the counting model's loss: 1.496
after 'ka', the name ends with probability 0.002 | counting model: 0.393
after 'kami', the name ends with probability 0.894 | counting model: 0.455
after 'kamiro', the name ends with probability 0.893 | counting model: 0.429

Your task: inside the for step in range(400) loop, above the if that prints, write one training step: take a batch into x, y, store the scores in logits and the cross-entropy in loss, then clear the old slopes, call backward() and step the optimizer.

Assignment
Follow the Coding Tutorial to write the six lines of one training step inside the loop.

Hint
Order matters: zero_grad() before backward(), and step() last. If the loss stays near 3, the weights are not moving: check that optimizer.step() is there.