Implementing the Modified Training Function

A language model does not become useful on its own. It has to be trained, and the work of training is organized around a function. A training function is a function you write to train a model: it takes a model and training data and returns the model in a trained state.

The phrase implementing a training function points to the fact that training is not automatic. Someone has to decide how the training process is organized and write it as a function. Once it exists as a function, the same procedure can be called again whenever the model needs to be trained, which keeps the process consistent and easy to repeat.

This article follows one idea from the chapter on training and evaluating large language models: how a training function is structured, and in particular how its learning rate is computed. The central question to keep in mind is this: what goes into a training function, and what state does the model hold when the function finishes? The source for this topic is brief, so the focus here is on the core mechanism and what it means for structuring your code.

The easiest way to build intuition for a training function is to follow the model's state from start to finish. Before training, the model is in its initial condition. During training, the function repeatedly processes data and updates the model. After training, the model is in a new condition that reflects the data it was trained on. Each stage is a snapshot of the same model at a different point in time.

The diagram shows the shape of the idea. The training data and the model both go into the training function, and what comes back is the model after training. The model's internal state is the part that changes, which is why the trained model is the changed element. Practically, this means you should check after the call that the model's state actually changed, because a function that returns without updating the model has not trained anything.

The source names two linked ideas: implementing a training function, and pretraining the LLM. Pretraining is the stage where a large language model is trained on its base data, and the training function is the code that carries out that stage. In other words, the function is the mechanism, and pretraining is the job it is used for.

A learning rate controls how large a change the optimizer makes to a model's weights on each update. In many simple training loops it stays at one value for the entire run. The training function in this article does something different: it recalculates the learning rate on every step of the loop, and it uses one rule for the early steps and another rule for the rest of training. The two rules are called the warmup phase and the cosine annealing phase.

The word linear describes the shape of the warmup change. The learning rate rises by the same amount on every step, so plotted against steps it forms a straight line. The increment is a fixed quantity computed once, before the training loop begins. The learning rate itself is recalculated inside the innermost loop, once per batch, because the step counter changes once per batch.

The increment added on each warmup step is the total rise divided by the number of warmup steps. The total rise is the distance from initial_lr up to peak_lr, and the number of warmup steps spreads that rise evenly across the warmup period.

Once the increment is known, the learning rate at any warmup step follows directly. The step counter global_step counts the steps taken so far, so the learning rate is the starting value plus global_step times the increment.

The diagram shows the warmup rising in equal steps of 0.1. Each batch adds the same amount, so the learning rate climbs in a straight line. When global_step reaches 4, which equals warmup_steps, the warmup branch no longer applies and the learning rate is taken from the cosine schedule, which starts from the peak.

The branch that handles warmup begins with the condition global_step < warmup_steps. When that condition is true, the code uses the linear formula. When it is false, the code uses the cosine formula. Once global_step reaches warmup_steps, the training function stops using the linear rule and starts using cosine annealing, a technique used to adjust the learning rate after the warmup phase.

The two phases are connected through the same step counter and the same peak_lr. Warmup climbs toward the peak, and cosine annealing starts from that peak and decreases toward the minimum. The variable progress measures how far training has moved through the post-warmup period. The source defines it as the number of steps completed after warmup divided by the number of steps remaining after warmup, which is the total number of training steps minus warmup_steps. Because progress is measured from the end of warmup, the cosine schedule begins from the peak learning rate and moves toward min_lr as training continues.

When a learning rate looks wrong during training, the fastest path to the cause is to ask three questions in order. First, which branch applied at that step: was global_step below warmup_steps? Second, what were the values of initial_lr, peak_lr, and warmup_steps that fed into the increment? Third, for a cosine step, what were total_training_steps and warmup_steps that defined progress? Each question points to one line of the training function, so the divergence can be located rather than guessed.

The snippet is an illustration of the structure described in the source, not a complete training function. The lines that matter for debugging are the increment line, the branch condition, and the warmup assignment. If a warmup value is wrong, check whether the increment was computed from the values you expected. If the value is wrong only after warmup ends, check the cosine branch and the progress inputs.