One training update: from tokens to an AdamW step

When people talk about training a large language model, the exciting parts seem to be elsewhere: the trillions of tokens, the architecture, the thousands of GPUs. But all of that machinery serves one small operation, repeated a few hundred thousand times: take a batch of tokens, measure how wrong the model’s predictions are, and nudge the weights. Everything else, from data cleaning to distributed training to fine-tuning, is about feeding that operation or running it faster without changing what it means.

So in this post we follow one update from start to finish. On the way we will meet several quiet conventions that decide what the model actually learns: which positions count as targets, what each position may look at, how the losses are averaged, and how AdamW turns a gradient into a step. Get any of them wrong and training still runs. It just optimizes something you did not intend.

If backpropagation, cross-entropy or Adam are new to you, the Foundations posts on backpropagation, softmax and cross-entropy, optimisers and language models cover them.

What one update has to do

A language model reads text left to right and, at every position, guesses the next token. One training update does four things:

1. Decide which positions are targets, and what the correct next token is at each.
2. Run the model, letting each position look only at what it is allowed to see.
3. Turn the guesses into one number, the loss, by averaging how surprised the model was at each target.
4. Compute how each weight should change to lower that number, then let the optimizer take a step.

The first three steps define what the model is trained to do. The fourth decides how fast and in what shape it moves there. We take them in order.

Targets: each position predicts the next token

Take the sentence <bos> the cat sat <eos>, where <bos> and <eos> mark the beginning and end of a document. At every position the model predicts the next token:

Five tokens give four targets: the targets are the inputs shifted left by one, and the last position has nothing to predict.

In code, the targets usually start as a copy of the inputs, called labels, and the shift happens inside the loss. In Hugging Face Transformers 4.51.3, ForCausalLMLoss appends one extra ignored label and drops the first one, so position $t$ is scored against the token at $t + 1$. Any label set to -100 is ignored. Some codebases, such as nanoGPT, shift the targets in the data loader instead and use -1 as the ignore value. Either is fine; mixing the two conventions is not.

What does “ignored” mean for learning? A label of -100 removes the loss term at that output position, so that position’s prediction sends no direct gradient back. It does not freeze the token. Its input embedding and hidden states can still receive gradients through later positions that attend to it. A tiny example: if position 1 averages the inputs of positions 0 and 1, $h_1 = (x_0 + x_1)/2$, and position 1 has a loss $\ell_1 = h_1^2$, then with $x_0 = 1$ and $x_1 = 3$ we get $h_1 = 2$ and

Position 0 has no target of its own, yet its input moves, because position 1’s prediction depends on it. This is exactly what happens to prompt tokens when fine-tuning masks their labels.

A batch with packing and padding

Real batches hold several rows of the same length. Suppose three short documents must fill two rows of six positions:

• document A: <bos> the cat <eos>
• document B: <bos> cat, cut off at the end of the row
• document C: <bos> cat sat <eos>

Row 0 packs A and B together so no slot is wasted. Row 1 holds C and is padded with two <pad> tokens.

Everyday picture: packing is like two small groups sharing a minibus so no seat goes empty. The attention mask, coming next, is the partition that stops one group listening to the other’s conversation.

Three kinds of position get no target:

• the last position of each row, which has no next token;
• positions whose next token is padding;
• A’s last token, whose next token is B’s <bos>. Training it to predict the start of an unrelated document teaches the packing format, not language, so many pipelines ignore this boundary target. Others keep it; the important thing is to know which you do.

That leaves 4 targets in row 0 and 3 in row 1: seven target positions out of twelve.

Attention masks: what each position may see

Labels say what each position must predict. The attention mask says what it may look at while predicting. They are separate decisions. In a decoder, position $t$ may attend to position $s$ only if:

1. $s \le t$, so it never sees the future;
2. $s$ is not padding;
3. in a packed row, $s$ belongs to the same document.

The third rule is the one people skip, so look at a position where it matters. In row 0, position 4 is B’s <bos>, and it must predict cat. Give the five tokens it could see the illustrative values 1, 2, 3, 4 and 1, and let attention simply average whatever is allowed:

• causal mask only: position 4 sees positions 0 to 4, so it averages $(1 + 2 + 3 + 4 + 1)/5 = 2.2$;
• with the document mask: it sees only position 4, so it averages $1/1 = 1$.

Same token, same target, different input to the prediction. Without the document mask, B’s prediction depends on an unrelated document A. Llama 3, §3.2: “We use an attention mask that prevents self-attention between different documents within the same sequence.” The paper reports that this had limited impact in standard pre-training but was important in continued pre-training on very long sequences, so it is worth measuring on your own data rather than assuming either way.

The padding at the end of row 1 never reaches a real position, even without rule 2: padding sits after every real token, so no real position satisfies $s \le t$ for a padded $s$. Rule 2 matters for left padding, used when generating from a batch of prompts of different lengths, where the padding comes first.

Position IDs are a third, separate choice: the position number each token gets for positional encoding. Isolating documents is the attention mask’s job and does not require restarting position IDs. With rotary position embeddings (RoPE), the attention score between two positions depends only on the distance between them, so shifting a whole document’s positions by a constant leaves its internal attention unchanged. Whether a codebase restarts positions at each document matters more for learned absolute positions and for how the model sees long contexts; codebases differ, so check yours.

The loss for one position

The model outputs a score, a logit, for every token in the vocabulary, and softmax turns the scores into probabilities. The loss at one position is how surprised the model is by the correct token: minus the log of the probability it gave that token.

Two quick cases with tiny vocabularies:

• Four tokens, all logits equal: each gets probability $1/4$, so the loss is $\ln 4 \approx 1.386$, whatever the target. This is a model that knows nothing.
• Three tokens with logits $(\ln 2,\ 0,\ 0)$: the probabilities are $\frac{2}{2+1+1}, \frac{1}{4}, \frac{1}{4} = \left(\frac12, \frac14, \frac14\right)$. If the target is the first token, the loss is $\ln 2 \approx 0.693$.

In general, for logits $z$ over the vocabulary and target token $y$,

• $z_v$: the logit of token $v$; $y$: the target token.
• The first term, the log-sum-exp, measures the total mass of all scores; the loss falls as $z_y$ rises relative to it.

Implementations compute the log-sum-exp after subtracting the largest logit, which changes nothing mathematically and prevents overflow.

Averaging: every token, or every row?

The batch has seven target positions, each with a loss. The update needs one number, and there are two natural ways to average. Use made-up losses to see the difference: row A has three targets with loss 1 each, and row B has one target with loss 3.

• Token mean: average over all targets, $(1 + 1 + 1 + 3)/4 = 1.5$.
• Mean of row means: average each row, then the rows, $\left(\frac{1+1+1}{3} + \frac{3}{1}\right)/2 = (1 + 3)/2 = 2$.

They differ because they weight tokens differently. The token mean gives each of the four tokens weight $1/4$. The row mean gives each of A’s tokens weight $\frac12 \cdot \frac13 = \frac16$ and B’s single token weight $\frac12$: three times as much as one of A’s.

Everyday picture: a school averaging exam marks. The token mean averages every paper in the school; the row mean averages each class first, so one paper in a class of 3 counts ten times as much as one paper in a class of 30.

Neither is wrong in itself, but they are different objectives, not two estimates of one. Pre-training normally intends the token mean, so that every token of text matters equally:

• $\theta$: the model’s weights; $b$ indexes rows, $t$ positions.
• $\mathcal{T}$: the set of target positions, seven in our batch; $N = |\mathcal{T}|$.
• $p_\theta$: the model’s probability for the true next token, computed under the attention mask.

Hugging Face’s ForCausalLMLoss uses the token mean by default: the cross-entropy is averaged over the labels that are not -100.

Gradient accumulation must keep the same objective

When a batch does not fit in memory, the step is split into microbatches: run each one, add up the gradients, and take one optimizer step at the end. Suppose each row of the example above is one microbatch.

The natural code averages each microbatch and divides by the number of microbatches:

which is the mean of row means, not the token mean. The fix is to divide each microbatch’s summed loss by the total number of targets in the whole step, $N = 4$:

Because gradients add up linearly, the summed gradients then equal the gradient of the token mean exactly. Dividing before backward needs $N$ in advance; equally valid is to backpropagate the plain sums and divide the accumulated gradients by $N$ once at the end, before clipping. Both are mathematically the same objective. They need not be bitwise identical, because floating-point additions happen in a different order.

How large can the distortion get? If one microbatch has 1,000 target tokens and another has 10, mean-then-divide gives each token of the small microbatch $\frac12 \cdot \frac{1}{10}$ and each token of the large one $\frac12 \cdot \frac{1}{1000}$: a 100-fold difference. Packing rows to a fixed length makes counts similar and hides this; variable-length batches expose it.

Two edge cases need explicit handling:

• A step with no targets. If every label in a step is ignored, $N = 0$ and the token mean is $0/0$. Skip the optimizer step and record that it happened, rather than dividing by zero. With several GPUs, $N$ must be the count over all of them, and all of them must make the same decision.
• Where $N$ comes from. In Transformers 4.51.3, the Trainer counts targets when the model accepts loss keyword arguments (as LlamaForCausalLM does) or when you pass your own loss function. It fetches all of a step’s microbatches, counts labels.ne(-100).sum() and passes that as num_items_in_batch, and the loss divides its sum by it. The count is taken on the unshifted labels, while the loss drops each row’s first label when shifting. With two rows of six tokens whose first label is <bos>, the count is 10 although only 8 positions are predicted. The ratio of true to counted targets depends on how many rows and targets each step has, so it is not a constant factor; it slightly reweights steps against each other. For pre-training rows of thousands of tokens it is negligible; for short fine-tuning sequences it is larger. Check the rule in the version you use.

The backward pass

Start with one position. Using the three-token example, the probabilities are $p = \left(\frac12, \frac14, \frac14\right)$ and the target is the first token, written as the one-hot vector $e_y = (1, 0, 0)$. The gradient of the loss with respect to the logits is

Read it as instructions: raise the target’s logit by $\frac12$, and lower each wrong token’s logit in proportion to the probability it was given. The entries sum to zero, as $p - e_y$ always does. Under the token mean each position’s loss is divided by $N$, so with $N = 4$ this position contributes $\left(-\frac18, \frac1{16}, \frac1{16}\right)$. Ignored positions contribute nothing.

Now one layer of weights. If the logits come from a hidden vector $h$ through an output matrix $W$, $z = hW$, the chain rule gives an outer product. With $h = (1, 2)$:

Each weight’s gradient is the size of the hidden feature it reads times the error of the logit it writes. Over a whole batch, the gradients from every target position are summed:

In a real Transformer the same $p - e_y$ signal enters the output layer and then flows back through every block.

Clipping the gradient

Before the optimizer step, most pre-training recipes cap the overall size of the gradient. The global norm treats every gradient entry of every parameter as one long vector. With just two entries, $(3, 4)$:

• the norm is $\sqrt{3^2 + 4^2} = 5$;
• with a cap of 1, every entry is multiplied by $1/5$, giving $(0.6,\ 0.8)$, which has norm 1 and the same direction.

In general, with gradient vector $g$ and cap $c$,

PyTorch 2.6’s clip_grad_norm_ computes exactly this, so the result is $(0.6, 0.8)$ up to the tiny $10^{-6}$ in the denominator. Because the norm is global, one parameter with a huge gradient shrinks every other parameter’s gradient too.

Three practical points:

• Clip once per step, after all microbatches are accumulated and before the optimizer step. With FP16 loss scaling, unscale the gradients first, so that the cap applies to their true size.
• Clipping a gradient is not shrinking the update. Adam divides each step by a running estimate of the gradient’s size, so a smaller gradient does not mean a proportionally smaller step. On the very first step, as the next section shows, the update is the same whether or not the gradient was clipped.
• Clipping cannot repair a broken gradient. In PyTorch 2.6, one NaN entry makes the norm NaN and turns every gradient into NaN; one Inf entry makes the norm Inf, which zeroes the finite entries and leaves NaN where the Inf was. By default no error is raised. Check that the returned norm is finite, and skip the step if it is not.

One AdamW step on one number

Adam keeps two running averages for each parameter: $m$, of the gradients, and $v$, of the squared gradients, both starting at zero. Follow one parameter $\theta = 1$ with learning rate $\gamma = 0.1$, $\beta_1 = 0.9$ and $\beta_2 = 0.99$, and ignore the tiny $\epsilon$ for now.

Step 1, gradient $g = 0.2$:

• $m = 0.9 \cdot 0 + 0.1 \cdot 0.2 = 0.02$ and $v = 0.99 \cdot 0 + 0.01 \cdot 0.04 = 0.0004$.
• Both started at zero, so early on they are too small. Bias correction divides by $1 - \beta^k$ at step $k$: $\hat m = 0.02 / 0.1 = 0.2$ and $\hat v = 0.0004 / 0.01 = 0.04$.
• The step is $\gamma \cdot \hat m / \sqrt{\hat v} = 0.1 \cdot 0.2 / 0.2 = 0.1$, so $\theta$ becomes $0.9$.

Now repeat with a gradient a hundred times smaller, $g = 0.002$. Every quantity shrinks, but $\hat m / \sqrt{\hat v}$ is again $0.002 / 0.002 = 1$, and the step is again $0.1$. On the first step, every parameter with a nonzero gradient moves by the learning rate, whatever the gradient’s size, in the direction of its sign.

Everyday picture: a crowd told “everyone take one step towards your exit”. On the first step only the direction each person faces matters, not how hard they were pushed. Adam learns how much to trust each push only as its averages fill up.

That is the special case. In general, at step $k$ with gradient $g$:

• $\gamma$: learning rate; $\beta_1, \beta_2$: how slowly the two averages forget; $\epsilon$: a small constant for numerical safety; $k$: the step number, starting at 1.

Step 2 shows that history matters. Continue from step 1 with gradient 0.2:

• If the second gradient is also 0.2, the step is again 0.1.
• If both gradients were halved, $(0.1, 0.1)$, the step is still 0.1: scaling the whole gradient history by the same positive factor leaves $\hat m / \sqrt{\hat v}$ unchanged.
• If only the current gradient is halved, $(0.2, 0.1)$: $m = 0.9 \cdot 0.02 + 0.1 \cdot 0.1 = 0.028$ and $v = 0.99 \cdot 0.0004 + 0.01 \cdot 0.01 = 0.000496$. After bias correction, $\hat m = 0.028/0.19 \approx 0.147$ and $\hat v = 0.000496/0.0199 \approx 0.0249$, so the step is about $0.1 \times 0.147/0.158 \approx 0.093$.

So Adam is insensitive to a consistent rescaling of all gradients, but not to rescaling one step relative to the others. That is why the Trainer’s varying count above is not harmless just because Adam is “scale invariant”.

Weight decay, in two conventions

AdamW adds weight decay: each step, every decayed weight shrinks slightly towards zero, separately from the gradient step. PyTorch 2.6 applies it first, multiplying the weight by $1 - \gamma\lambda$, where $\lambda$ is the weight_decay argument. Continuing the scalar example with $\lambda = 0.1$:

The original paper writes it differently. Loshchilov and Hutter (2019), Algorithm 2 update $\theta \leftarrow \theta - \eta_t\left(\alpha\,\hat m/(\sqrt{\hat v}+\epsilon) + \lambda\theta\right)$, where $\alpha$ is the base learning rate and $\eta_t$ a schedule multiplier. The decay is multiplied by the schedule, but not by the learning rate.

To convert, compare the shrink per step. In PyTorch the learning rate at step $t$ is $\gamma_t = \alpha\,\eta_t$, so the shrink is $\alpha\,\eta_t\,\lambda_{\text{PyTorch}}$. In the paper it is $\eta_t\,\lambda_{\text{paper}}$. They match when

A recipe of “learning rate $3\times10^{-4}$, weight decay 0.1” in PyTorch corresponds to $\lambda_{\text{paper}} = 3\times10^{-5}$. The same name, “weight decay”, differs by a factor of the learning rate between the two conventions, which matters whenever you copy a recipe from one codebase to another.

Which parameters decay is a separate convention. nanoGPT, for example, decays every parameter with two or more dimensions (the weight matrices and embeddings) and leaves biases and normalization weights alone. Recipes differ, so record the choice.

The same step in real code

A minimal step for a Hugging Face causal language model, written so that accumulation keeps the token mean and bad steps are skipped:

num_targets = sum((mb["labels"][:, 1:] != -100).sum() for mb in microbatches)
if num_targets == 0:
    skip_step()                       # nothing to learn from; never divide by zero
else:
    for mb in microbatches:
        logits = model(input_ids=mb["input_ids"],
                       attention_mask=mb["attention_mask"]).logits
        loss_sum = F.cross_entropy(
            logits[:, :-1].reshape(-1, logits.size(-1)).float(),  # predict t+1 from t
            mb["labels"][:, 1:].reshape(-1),
            ignore_index=-100, reduction="sum")
        (loss_sum / num_targets).backward()        # gradients add up to the token mean
    norm = torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
    if torch.isfinite(norm):
        optimizer.step()
    scheduler.step()
    optimizer.zero_grad(set_to_none=True)

The slicing logits[:, :-1] against labels[:, 1:] performs the shift; do not also pass pre-shifted labels. In Hugging Face, attention_mask marks padding, not document boundaries: isolating packed documents needs a separate mechanism, such as variable-length attention kernels driven by document lengths.

At realistic sizes the loss itself costs memory. ForCausalLMLoss upcasts the logits to 32-bit floats before the cross-entropy, a precision safeguard. With a vocabulary of 32,000 and rows of 4,096 tokens, that is $4096 \times 32000 \times 4 \approx 524$ MB per row, which is why large training codebases often compute the loss in chunks.

Failure modes

Where this update goes next

This one update is what the rest of the series keeps coming back to. Next we go upstream, to where the tokens come from: how a pre-training corpus is built, then how text becomes tokens. Later posts ask how to compute exactly this update when the batch, and then the model, no longer fit on one GPU. For a head start on that side, the paper reviews of ZeRO and Megatron-LM cover sharding the optimizer state and splitting the model across devices.

What to remember

• Targets are the inputs shifted by one. Positions labelled -100 contribute no loss and no direct gradient, but their inputs still receive gradients through later positions that attend to them.
• Label masking, attention masking and position IDs are three separate decisions. The attention mask isolates packed documents; Llama 3 does this.
• Token mean and mean of row means are different objectives: 1.5 against 2 for losses $[1, 1, 1]$ and $[3]$.
• Gradient accumulation keeps the token mean only if every microbatch’s sum is divided by the step’s total target count, before or after backward. A step with no targets must be skipped.
• The gradient with respect to the logits is $p - e_y$, divided by $N$ under the token mean; a weight matrix’s gradient is the outer product of its input and that error.
• Clip once per step, after accumulation and any unscaling. Clipping preserves direction, does not proportionally shrink Adam’s updates, and cannot repair NaN or Inf.
• AdamW’s first step moves every parameter with a nonzero gradient by the learning rate. Adam is insensitive to rescaling the whole gradient history, not one step.
• PyTorch multiplies weight decay by the learning rate; the original paper does not: $\lambda_{\text{paper}} = \alpha\,\lambda_{\text{PyTorch}}$.

Try it yourself

• Figure 1. Step to the document mask: seven of twelve positions are targets, and position 4, which predicts cat, can see only itself.
• One loss. Two tokens with equal logits: each gets probability $\frac12$, and the loss is $\ln 2 \approx$ 0.693 whichever is the target.
• Two averages. Row A has losses $[2, 2]$ and row B has $[5]$. The token mean is $9/3 =$ 3; the mean of row means is $(2 + 5)/2 =$ 3.5.
• Accumulation. Same rows as two microbatches. Dividing each sum by the step’s 3 targets gives $4/3 + 5/3 =$ 3, the token mean; mean-then-divide gives 3.5.
• Clipping. Gradients $(6, 8)$ with a cap of 2: the norm is 10, so multiply by $\frac15$ to get (1.2, 1.6).
• First Adam step. Learning rate 0.1, gradient $-5$: the parameter increases by 0.1. With gradient $-0.0005$ it also increases by 0.1.
• Weight-decay conversion. PyTorch recipe with learning rate $10^{-3}$ and weight decay 0.1: the paper’s $\lambda$ is $10^{-4}$.

Sources

• Llama Team, AI @ Meta (2024). “The Llama 3 Herd of Models.” arXiv:2407.21783, v3. §3.2 on masking attention between documents.
• I. Loshchilov and F. Hutter (2019). “Decoupled Weight Decay Regularization.” ICLR 2019. arXiv:1711.05101. Algorithm 2, AdamW.
• D. P. Kingma and J. Ba (2015). “Adam: A Method for Stochastic Optimization.” ICLR 2015. arXiv:1412.6980.
• J. Su, Y. Lu, S. Pan, A. Murtadha, B. Wen and Y. Liu (2021). “RoFormer: Enhanced Transformer with Rotary Position Embedding.” arXiv:2104.09864.
• Hugging Face Transformers 4.51.3. ForCausalLMLoss in loss_utils.py; get_batch_samples and compute_loss in trainer.py.
• PyTorch 2.6. The AdamW update in optim/adam.py; clip_grad_norm_ in clip_grad.py.
• A. Karpathy. nanoGPT, model.py (configure_optimizers).