Loss Calculation and Training Functions

By the point this article begins, the chapter has already completed two large pieces of work. The first is a model that has been pre-trained and loaded into memory. The second is a dataset of instructions that has been cleaned, formatted, and turned into batches. Instruction fine-tuning is the step that connects these two pieces. The model is trained further, this time on instruction examples, so that it learns to respond to the kind of input the dataset describes.

This article looks at the moment just before that training step. The goal is to understand what has been prepared, what the fine-tuning step is allowed to rely on, and how to tell which part of the overall process you are looking at when something goes wrong.

Before any training happens, two things are already in place. The model is the one loaded in the previous section, and fine-tuning does not begin from random weights. It continues training the model that already exists. The pre-trained LLM therefore contributes its starting state: the weights it has already learned during pre-training.

The data side is the instruction dataset prepared earlier in the chapter. That dataset was processed at the start of the chapter, so the fine-tuning step does not repeat that preparation work. The instruction dataset contributes the examples that drive the training: the kinds of inputs the model should learn to respond to.

Instruction fine-tuning is organized as a three-stage process. Stage 1 prepares the dataset. Stage 2 fine-tunes the LLM. Stage 3 evaluates the LLM. The overall process is laid out as nine numbered steps in order, with each step grouped under the stage it belongs to. The diagram below shows the three stages in sequence and the steps that sit inside each one.

Notice where the fine-tuning work sits. Stage 1 and Stage 3 surround it. Stage 1 produces the batched data loaders that the training step will consume. Stage 3 looks at the responses the tuned model produces afterward. Stage 2 is the only stage in which the model's weights are trained, and it contains three steps: loading the pretrained LLM, running instruction fine-tuning, and inspecting the modeling loss.

Stage 2 is the center of the process, and its three steps follow a clear order. First the pretrained LLM is loaded. Then instruction fine-tuning runs, which continues training that loaded model on the instruction examples. Finally the modeling loss is inspected, so you can see how the training is progressing.

Fine-tuning needs a way to compute loss and run training. The chapter already contains code for both, so the fine-tuning code does not write those routines again. Instead it imports them. The two functions are calc_loss_loader and train_model_simple, and both come from chapter 5.

Four inputs meet at the fine-tuning step. The pre-trained LLM is loaded and then trained further. The instruction dataset supplies the examples for that training. The two chapter 5 functions supply the loss calculation and the training procedure. The diagram below shows how these inputs feed the single fine-tuning step.

Suppose a run of the process produces a result you did not expect. The three-stage structure gives you a way to narrow the search. Ask first which stage the problem most likely belongs to, and only then look at the steps inside that stage. The diagram below maps common symptoms to the stage where you should begin looking.