JJ

Pretraining Loop

3 min read

Pretraining a language model is a loop of five repeated steps. The data is already curated at this point, so this note is about what happens to it once training starts.

Basic training cycle

One batch follows five operations: clear old gradients, calculate predictions, calculate loss, calculate gradients, and update parameters. Repeat this cycle across the training data.

For next-token learning, input tokens and target tokens are shifted by one position. For example, input [1, 2, 3] has target [2, 3, 4]. Each output position predicts the token that follows its input position.

Five steps

  • batch tokens, so shuffle documents, pack windows and create the masks
  • forward, in mixed precision, with activation checkpointing if memory is short
  • backward, accumulating microbatches and synchronising gradients
  • update, with AdamW, gradient clipping, and a warmup followed by decay
  • monitor tokens per second, loss, gradient norm and validation perplexity

Packing

  • documents are joined, separated by an end-of-sequence token, then cut into windows of a fixed context length
  • every window is full, so no compute is wasted on padding
  • this is why one document can straddle two windows

Progress

  • the stable unit of progress is tokens processed, and not epochs or steps
  • two runs with different batch sizes are comparable once both are read in tokens

Debugging route

  • loss does not fall, so check the target shift, the learning rate, and whether the gradients are non-zero
  • model sees the future, so inspect the causal mask orientation and its boolean semantics
  • padding dominates, so use a padding mask and exclude padded targets from the loss
  • NaNs appear, so inspect the logits, lower the learning rate, clip the gradients and use stable mixed precision
  • samples repeat, so check overfitting, the decoding temperature and the data diversity

First try to fit one tiny batch. If its loss fails to fall, inspect the data, model capacity, loss, and update steps before scaling up.

The masks created in the first step are compared in Attention Masks.

Code

PyTorch: one complete small training loop

import torch
from torch import nn
 
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
X = torch.tensor([[1, 2, 3], [2, 3, 4]], dtype=torch.long).to(device)
y = torch.tensor([[2, 3, 4], [3, 4, 5]], dtype=torch.long).to(device)
model = nn.Sequential(nn.Embedding(6, 8), nn.Linear(8, 6)).to(device)
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
 
for epoch in range(10):
    model.train()
    optimizer.zero_grad()
    logits = model(X)
    loss = criterion(logits.reshape(-1, 6), y.reshape(-1))
    loss.backward()
    optimizer.step()
 
model.eval()
with torch.no_grad():
    predictions = model(X).argmax(dim=-1)
print(tuple(predictions.shape))  # (2, 3)

This small model predicts from the current token only. It demonstrates the update loop; a recurrent or attention layer adds wider context.

Read each operation:

  • .to(device) places the model, inputs, and targets on the same device
  • model.train() selects training behaviour for relevant layers
  • zero_grad() clears gradients, which otherwise accumulate
  • logits has shape (2, 3, 6): two sequences, three positions, six scores
  • reshape(-1, 6) makes six prediction rows; y.reshape(-1) gives their six target IDs
  • backward() calculates gradients and step() updates weights
  • eval() selects evaluation behaviour; no_grad() separately stops gradient recording
  • argmax(dim=-1) selects the predicted token ID at each position

TensorFlow/Keras: training managed by fit

import tensorflow as tf
 
X = tf.constant([[1, 2, 3], [2, 3, 4]])
y = tf.constant([[2, 3, 4], [3, 4, 5]])
model = tf.keras.Sequential([
    tf.keras.Input(shape=(3,), dtype="int32"),
    tf.keras.layers.Embedding(6, 8),
    tf.keras.layers.Dense(6),
])
model.compile(
    optimizer=tf.keras.optimizers.Adam(learning_rate=0.01),
    loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
)
model.fit(X, y, batch_size=2, epochs=10, verbose=0)
predictions = tf.argmax(model(X, training=False), axis=-1)
print(tuple(predictions.shape))  # (2, 3)

fit performs the forward, loss, gradient, and update steps. The sparse categorical loss accepts the integer target IDs. Final performance must be evaluated on separate data.

Graph View

Backlinks