Pre-training: weight optimisation

For each iteration in our pre-training loop, we have a batch of inputs and targets. For each input/target pair, our model simultaneously calculates the probability of predicting the target token from the input tokens:

We then use a loss function called cross-entropy loss to calculate how well our model did at predicting the correct tokens. This takes all of our predictions across all inputs in a given batch, and produces a single score called the loss.

logits = model(inputs)
loss = F.cross_entropy(
    logits.flatten(0, 1),
    targets.flatten(),
)

The reason for using flatten is that PyTorch cross-entropy expects a two-dimensional collection of predictions and one target ID for each prediction. Flattening does the following:

Each flattened row of logits is still paired with exactly the same target token as before. Flattening changes only the layout expected by the loss function; it does not mix predictions and targets.

Cross-entropy gives a larger loss when the correct token receives a low score and a smaller loss when it receives a high score. By default, PyTorch averages the loss across all batch × tokens predictions.

Backpropagation

Once the cross-entropy loss has been calculated, backpropagation calculates how sensitive that loss is to each weight in the model, i.e., how much changing that weight would likely change the loss.

More precisely, it computes the partial derivative of the loss with respect to every weight. These derivatives collectively form the gradient.

To do this, the model’s calculations are treated as a chain - or computational graph - of simple mathematical operations. Backpropagation works backward through these operations, repeatedly applying the differentiation chain rule to calculate how each weight may have contributed to the loss.

Don’t worry if partial derivatives and the chain rule are unfamiliar; the key point is that backpropagation provides a systematic way to obtain the gradients needed to improve the model.

Optimisation

An optimiser then uses the resulting gradients to update the weights in a direction expected to reduce the loss. The optimiser often used in LLMs is one called AdamW.

Putting this all together, the code looks like the following:

model.to(device)
model.train()
optimiser = torch.optim.AdamW(model.parameters(), lr=learning_rate)

for epoch in range(epochs):
    total_loss = 0.0

    for inputs, targets in batches:
        inputs = inputs.to(device)
        targets = targets.to(device)

        optimiser.zero_grad()
        logits = model(inputs)
        loss = F.cross_entropy(
            logits.flatten(0, 1),
            targets.flatten(),
        )
        loss.backward()
        optimiser.step()

        total_loss += loss.item()

• optimiser.zero_grad() clears old gradients, because PyTorch accumulates gradients unless asked to reset them.
• model(inputs) performs the forward pass and records the operations needed for differentiation.
• loss.backward() runs backpropagation and stores a gradient on every parameter that contributed to the loss.
• optimiser.step() uses those gradients, AdamW’s running statistics, the learning rate, and weight decay to update the parameters.

Monitor the training loss

After each pass through the dataset, our script prints the average loss across its batches:

average_loss = total_loss / len(batches)
print(f"Epoch {epoch + 1}: loss = {average_loss:.4f}")

This number is useful to show that optimisation is working, but for a real LLM a slightly more sophisticated training system is used, where some of our text corpus (called the validation dataset) is held back from training, and used to validate the model.