Fine-Tuning a Pre-Trained LLM for Spam Classification

This article works from train_classifier_simple, a function used to fine-tune a model to classify spam. Its job is to train a model on a training dataset and to evaluate it on a validation dataset as training proceeds. Instead of memorizing the function's description, treat it as a program you can trace. Ask two questions as you read: which inputs drive the loops, and which outputs does the function hand back to its caller?

The function takes the model, the two data loaders (one for training and one for validation), an optimizer, the device to run on, and four control values: num_epochs, eval_freq, and eval_iter. The first group supplies the things to be trained and the data to train on. The control values set how long training runs and how often it pauses to measure performance.

The function returns four lists and a counter. The lists record training loss, validation loss, training accuracy, and validation accuracy at each evaluation point. The counter, examples_seen, records how many training examples the model has processed in total. Together these outputs let you follow training progress over time, because each list entry corresponds to one evaluation point during the run.

The function contains two loops, one inside the other. The outer loop, for epoch in range(num_epochs), runs once per epoch. At the start of each epoch the function calls model.train(), which puts the model in training mode. The inner loop, for input_batch, target_batch in train_loader, runs once per batch of training data. Each batch passes through the same sequence of steps, and only after the inner loop finishes does the outer loop move on to the next epoch.

The loops nest in a fixed way: the outer loop controls how many full passes over the data occur, and the inner loop controls how many batches make up each pass. The batch body is the only place where weights change during training. Counting the work is therefore simple multiplication: the total number of batch passes equals num_epochs times the number of batches in the training loader.

Each batch runs four training steps in a fixed order. First, optimizer.zero_grad() resets the loss gradients left over from the previous batch iteration. Second, calc_loss_batch computes the loss for the current input and target batch. Third, loss.backward() calculates the loss gradients. Fourth, optimizer.step() updates the model weights using those gradients. The reset sits at the top of the batch body because it clears the gradients from the previous batch before new ones are calculated.

Read the diagram as the path a single batch follows. Every batch goes through the four weight-update steps. Only some batches reach the evaluation step, and the diagram shows that choice as a decision rather than a fixed step. The evaluation branch returns to the batch check, so training continues after an evaluation rather than restarting.

After the weight update, the function performs two bookkeeping actions. It adds the number of rows in the current input batch to examples_seen, using input_batch.shape[0], and it increments global_step by one. Because examples_seen counts examples rather than tokens, the source notes this as a change from the earlier pretraining code. The counter global_step is what later decides whether an evaluation happens.

The check if global_step % eval_freq == 0 controls how often the model is evaluated. The % operator gives the remainder after division, so the condition is true whenever global_step is an exact multiple of eval_freq. The value eval_freq is therefore a count of training steps between evaluations, not a count of epochs. Since global_step starts at -1 and increases with every batch across all epochs, the counter keeps running through the whole training run and does not reset at each epoch.

When the check passes, the function calls evaluate_model to get a training loss and a validation loss, appends both to their lists, and prints them with the current epoch and step. It then calculates accuracy on both loaders and records those values as well. The accuracy calculation uses num_batches=eval_iter, so only a limited number of batches from each loader are used for this measurement.

The evaluate_model function is the other half of the picture. The source states it is identical to the one used for pretraining. It switches the model into evaluation mode with model.eval(), and it runs its calculations inside a torch.no_grad() block so that no gradients are computed. It estimates loss on both datasets and does not update any weights. This is the key difference from a training step: a training step computes gradients and changes weights, while evaluate_model only measures.