A Transformer forward pass is only part of training. How does text become targets? Which tokens count toward the loss? How are gradients aggregated, and how do we know whether an update helped? This chapter follows that sequence.

First, understand one small step on one device

If you have not written a training loop, run the two-example PyTorch update first. Separate forward, backward, and parameter updates before returning to token targets. Read initialization, optimizers, and training diagnosis when questions about gradient scale or generalization arise.

Consider two sequences. One has 2 valid targets and a summed loss of 4; the other has 6 valid targets and a summed loss of 6.

The token-weighted mean is $(4+6)/(2+6)=1.25$. Averaging the two sequence means instead gives $(2+1)/2=1.5$. The data did not change, but each token's weight in the objective did. Gradient accumulation, distributed training, and padding all bring this small issue back.

Establish the target, mask, and denominator before asking how to compute it faster.

Reading order Question to follow Check afterward
1. Language-model objectives Which position predicts which token? Why can training run in parallel while generation usually cannot?
2. The pretraining pipeline How are texts cleaned, deduplicated, mixed, and packed? Are EOS, attention masks, and loss masks the same?
3. One training step How does a batch change parameters? Does changing batch partitioning preserve the objective?
4. Multi-token prediction What supervision do extra heads receive? Can targets cross document boundaries or leak answers?

Then consider precision, memory, and multiple devices

These interventions act at different levels. BF16 changes number representation; recomputation changes stored intermediates; FSDP changes where state lives; TP partitions a matrix computation. All can affect memory, but they are not interchangeable switches.

Symptom Check first Continue with
Inf / NaN loss Numeric range, scaling, and clipping order Precision and memory
Weights fit but backward runs out of memory Activations, optimizer state, temporary tensors, sequence length Precision and memory
Results change with device count Data partitioning, valid-token counts, synchronization, randomness Distributed training
High utilization but poor throughput Useful-token throughput, recomputation, communication, waiting Distributed training · How MFU is computed

On a first pass, the overview's distinction between DDP, FSDP, TP, and PP is enough. For the mechanics inside a layer, continue with two GPUs computing one FFN: calculate [2, −1] into [9, 6], then follow the input gradient to see where communication happens. Framework configuration names can wait. Once data and model parallelism are clear, read DDP, FSDP2, and training stacks. Follow parameter gathering and gradient updates to see what sharding saves and why it can cost time. The note also separates Trainer, Accelerate, NeMo RL, and generation engines by responsibility.

Chapter contents

01

Pretraining & training mechanics

18 notes

Connect the objective, data, and one batch to precision, memory, and distributed execution.

Training foundations4

Diagnose learning failures, then examine activations, initialization, and optimizer updates.

  1. Generalization & diagnosis
  2. Activations & initialization
  3. SGD to AdamW
  4. Muon: updating a matrix
01 · Objectives, data & updates5

Establish the prediction target, then follow the data into a batch and a parameter update.

  1. Training route
  2. LM objective
  3. Pretraining pipeline
  4. One training step
  5. Multi-token prediction
02 · Precision, memory & parallelism4

After one training step, separate numeric formats, stored state, recomputation, and multi-device partitioning.

  1. Precision & memory
  2. Distributed training
  3. Tensor parallelism, worked through
  4. DDP, FSDP2 & training stacks
03 · Scaling laws & hyperparameters (draft)5
  1. Overview
  2. Compute-optimal size and data
  3. Training past Chinchilla
  4. LR, batch size & schedule
  5. Review

Put it to use

Run forward, backward, and one update on a fixed tiny batch and save the parameter changes. Change only one condition, such as accumulation steps or precision. Compare outputs, gradients, and the update within appropriate tolerances before scaling up.

Worked numbers test mechanisms; they are not hardware benchmarks. Report model version, sequence length, effective batch, dtype, hardware, and timing boundaries for actual memory or speed measurements.

Continue with SFT and post-training, or turn to generation and inference.