Writing the train_model_simple Function

A neural network starts out with weights that produce poor predictions. Training is the process of nudging those weights again and again until the predictions improve. No single nudge is enough, and the network never sees the whole dataset as one lump. Instead, the training code is a set of nested repetitions: an outer loop that goes over the entire training set, and an inner loop that works through that set in smaller chunks.

This article follows the approach of starting from the smallest unit of learning, one weight update, and building outward. Once you can explain what a single update needs, the outer repetitions and the monitoring steps become simple additions rather than items in a list you have to memorize. The function that ties all of this together is train_model_simple. Its listing is the one piece of code this section focuses on, and your goal is to notice its pattern: which work happens once, which work repeats for every batch, and which work happens only some of the time.

Before reading the explanation, commit to a prediction. Consider only the work that happens for one single batch of training data. Three things must occur: the model's gradients need to be cleared, a loss must be computed from the model's output, and the weights must be changed using that loss. Decide which of these comes first and which comes last.

Every weight update in a PyTorch training run follows the same core cycle. Each step has a specific job, and the steps depend on one another in a fixed order.

In train_model_simple these four operations appear as gradient reset, a call to calc_loss_batch, a backward call on the loss, and optimizer.step(). The optimizer is the object that actually changes the weights, so optimizer.step() is the step that makes the learning happen.

Rather than computing one update from the entire training set at once, the training loop processes the data in batches. Each batch is a group of training examples, and the loop iterates over the batches inside every epoch. An epoch is one complete pass over the training set, and it contains one batch iteration for each batch. Each pass through a batch produces one round of the core update cycle, so the number of batches determines how many weight updates happen during one complete pass through the data.

Here $N$ is the number of training examples and $B$ is the batch size, which is how many examples each batch holds. A larger batch size means fewer batches per pass, and a smaller batch size means more batches. Because the batch count comes from this division, you can work it out before training begins.

Training usually runs for more than one epoch. The epoch loop sets the number of full passes, and within each pass the batch loop performs the core update cycle once per batch. This is why iterating over epochs is the outermost repetition: everything else lives inside it.

The signature of train_model_simple takes ten parameters: model, train_loader, val_loader, optimizer, device, num_epochs, eval_freq, eval_iter, start_context, and tokenizer. Grouping them by job makes the function easier to read.

Each pass through the inner loop handles one pair of input and target batches. The function first clears the gradients left from the previous batch. It then calls calc_loss_batch, which takes the input batch, the target batch, the model, and the device, and returns a loss. Calling backward on that loss computes the loss gradients for the model's weights. Next, optimizer.step() uses those gradients to update the weights. After the update, tokens_seen and global_step are advanced, and the evaluation branch may run.

The evaluation branch comes after the update and the counter advance, so evaluation always sees the weights produced by the batch that just finished. The diagram below traces this path for a single batch, including the decision that either runs evaluation or skips straight back to the next batch.

Before the loop starts, the function creates three empty lists for losses and tokens, sets tokens_seen to zero, and sets global_step to negative one. The starting value of negative one is deliberate: the counter is increased once per batch, so the first batch leaves it at zero.

Each batch adds the number of elements in its input tensor to tokens_seen, so tokens_seen measures how much input the model has consumed so far. global_step is different. It counts batches across all epochs rather than restarting at each epoch, which is why it can be used to decide when evaluation runs.

Notice what does not happen at an epoch boundary. global_step keeps counting from 9 to 10 when epoch 2 starts, rather than resetting to zero. tokens_seen likewise keeps growing. Both counters describe the whole run, so the epoch boundary changes which pass over the data is underway but not the running totals.

Evaluation runs only when global_step is divisible by eval_freq. Because global_step is incremented after each batch and starts at negative one, the first batch brings it to zero. Zero is divisible by any eval_freq, so evaluation also runs after the first batch. After that, evaluation occurs after every eval_freq-th batch counted from that point.

The workflow ends with two steps that are not needed for learning itself. One prints the training and validation set losses so you can watch progress, and the other generates sample text so you can inspect what the model currently produces. The source describes these as optional steps used for tracking training progress. They change nothing about the weights; they only let the person running training see how things are going.

train_model_simple calls two functions that this section does not define. You will see their calls when tracing the listing. Treat them as black boxes for now: their internals are outside the scope of this function, and the function itself is designed so that it can be read without them.