Run the Forward, Backward, and Update Cycle

Order training mode, gradient clearing, forward logits, shifted-target cross-entropy, backward, gradient inspection or clipping, and the optimizer update by their dependencies.

One training step has a dependency order. Changing the order can reuse stale gradients, clip too late, or update parameters before the current loss has been differentiated.

Read it in this order:

select training behavior;
clear gradients left by the preceding step;
calculate $(B,T,8)$ logits;
pair the flattened logits with the already-shifted targets;
backpropagate through the recorded forward computation;
inspect and, if needed, clip the current gradients;
let the optimizer update parameters and its own state.

Gradients Accumulate unless Cleared

PyTorch adds a new gradient contribution to an existing .grad tensor. This is useful when accumulation is deliberate. In this experiment, each optimizer step represents one batch, so stale gradients would change the effective update.

Using set_to_none=True leaves gradients absent until backward creates them. This makes a missing gradient easier to distinguish from a real zero gradient.

The Loss Uses Raw Logits

Cross-entropy performs a numerically stable log-softmax internally. Do not apply softmax first. The target tensor contains class IDs with shape $(BT)$ , not one-hot rows and not probabilities.

Inspect before Updating

After backward() and before step():

For deliberately frozen parameters, a missing gradient can be expected. This chapter has none. The tied embedding appears under one registered parameter name and receives the sum of both computational paths.

Q1. Repair the update order

Put these operations in order for one ordinary step: optimizer.step(), loss.backward(), optimizer.zero_grad(), forward and loss, gradient clipping.

Choose the correct order

zero_grad; forward and loss; backward; clip gradients; optimizer step

optimizer step; zero_grad; forward and loss; backward; clip gradients

zero_grad; clip gradients; forward and loss; optimizer step; backward

Select one choice, then check.

Hint

The optimizer cannot use the current gradient before backward creates it.

Solution

zero_grad; forward and loss; backward; inspect or clip gradients; then optimizer.step().

Record Step-Level Evidence

At selected steps, store batch loss, full training loss, validation loss, global gradient norm before and after clipping, and learning rate. Logging every tensor at every step would hide the trend and consume unnecessary storage.

Run the Complete Experiment Locally

Download the complete tiny-Transformer training script. It contains the frozen model, corpus sampling, one-sequence and corpus runs, metric logging, checkpoint serialization, reload, and exact next-step resume test.

Use the official PyTorch installation selector for your operating system and Python environment, then run:

python train-tiny-transformer.py

The recorded regression values were produced with Python 3.14, PyTorch 2.13.0, and CPU execution. A compatible newer environment may produce slightly different floating-point results. The structural assertions and qualitative outcomes should still hold; investigate rather than deleting a failed check.