One training step, taken apart

Four things happen, in order

Everything in this course is a variation on one step, repeated a few thousand times. The step has four parts: a forward pass, a loss, a backward pass, and an update. If you can narrate these with real shapes, every hyperparameter later stops being folklore.

Take a concrete batch: 4 sequences of 1,024 tokens each, on an 8B model with a vocabulary of 128,256.

1. Forward

The token ids go in as a tensor of shape (4, 1024). The embedding table turns each id into a vector of 4,096 numbers, giving (4, 1024, 4096). That tensor passes through every transformer block (attention, then a feed-forward network) and comes out the same shape at the end.

The last layer projects it to vocabulary size: (4, 1024, 128256). These are the logits, one score per possible next token, at every position.

Notice the size of that tensor. Four times 1,024 times 128,256 is 525 million numbers. In fp32 that is 2.1 GB for a single forward pass, and the backward pass needs room for its gradient too. On many real runs the logits are the largest single allocation in the whole job, bigger than the weights of the layers that produced them. This is why fused or chunked cross-entropy kernels (Liger’s, or cut cross-entropy) are worth switching on: they compute the loss without ever materialising the full logits tensor, and they routinely cut peak memory by 30 to 50% on models with large vocabularies.

2. Loss

At every position the model predicted a distribution over the next token, and you know what the next token actually was. Cross-entropy measures the gap: it is the negative log of the probability the model assigned to the correct token.

# labels are the inputs shifted left by one; -100 marks positions to ignore
loss = F.cross_entropy(logits[:, :-1].reshape(-1, V),
                       labels[:, 1:].reshape(-1),
                       ignore_index=-100)

That shift by one is the whole self-supervised trick: the target for position i is the token at position i+1, so a sequence of 1,024 tokens contains 1,023 training examples and you never had to label anything.

The ignore_index=-100 is how masking works. Positions labelled -100 contribute nothing to the loss and therefore nothing to the gradient. In instruction tuning that is where you put the prompt tokens, so the model is scored only on what it was supposed to write.

3. Backward

The loss is one number. loss.backward() walks the graph in reverse, and for every trainable parameter computes the partial derivative of that number with respect to that parameter. The result is a gradient tensor the same shape as the parameter.

This is why gradients cost as much memory as the weights they belong to, and why freezing 99% of the model, which is what LoRA does, removes 99% of that cost.

The backward pass also needs the activations from the forward pass, because the chain rule refers to them. Those activations were kept in memory throughout. That is the third major consumer of memory, and the reason gradient checkpointing exists: throw the activations away, recompute them during the backward pass, pay roughly 30% more compute to save a great deal of memory.

4. Update

The optimiser turns gradients into a change in the weights. In its simplest form:

w = w - learning_rate * gradient

AdamW, which you will almost always use, keeps two running averages per parameter and divides by the square root of the second, so a parameter with consistently small gradients still moves. Those two averages are why AdamW costs 8 bytes per trainable parameter in state alone.

Then the gradients are zeroed and the next batch begins.

What to take from this

Three facts do most of the work later:

• The gradient is computed from one batch. It is a noisy estimate of the direction you want. Larger batches give a less noisy estimate, which is why batch size and learning rate are linked.
• Masked positions are invisible. Not down-weighted, not de-emphasised: absent. If your mask is wrong, the model is learning a different task from the one you think, and the loss curve will look perfectly healthy while it does.
• Only trainable parameters get gradients and optimiser state. Everything about parameter-efficient fine-tuning follows from this one sentence.

The honest limitation of this picture: it describes what happens, not why it works. Nobody can tell you in advance which weights will change or what the model will have learned when 3,000 of these steps are done. That is why this course spends a whole module on measuring the result rather than predicting it.