13 Read and control the training loop

Inspect one update before tuning a hundred settings

After your first neural run, open its training loop. Find where the batch is created, where the model runs, where the loss is calculated, where backward is called, and where the optimizer updates weights. Print the shapes and dtypes once. Print how many parameters require gradients. Those few checks answer more useful questions than changing a dozen unfamiliar flags.

The byte-transformer loop uses fixed-length windows, so each microbatch contains the same number of target positions. The loss is averaged over those positions. Dividing by the accumulation count makes the accumulated gradient an average across equal-size microbatches. This simple accounting becomes more subtle when sequence lengths and supervised-token counts differ.

Suppose one microbatch contains 100 supervised tokens and another contains 1,000. Averaging their mean losses equally gives each microbatch equal influence, so each token in the shorter batch counts ten times as much. If your objective is an average over supervised tokens, accumulate summed losses with a correct total-token denominator or use a trainer whose normalization matches that objective. State what is being averaged.

Batch size is a statistical and systems choice

A batch is a group of examples used together. A microbatch is the group processed in one forward/backward pass. The effective batch is the group contributing to one parameter update. On one GPU with equal microbatches, effective examples = microbatch size times accumulation steps.

For variable-length text, examples per update can be misleading. Record input tokens and supervised tokens as well. Two batches of eight examples can differ greatly if one contains short labels and the other contains long answers. Packing multiple short sequences into a longer block can improve utilization, but only if the boundary and attention behavior match the intended objective.

A larger effective batch can make gradients less noisy, but it can also change optimization behavior and the number of updates at a fixed token budget. Do not increase batch size, learning rate, and training duration together and then attribute a change to one of them. Start with a feasible microbatch, choose an effective token budget per update, and compare controlled changes.

Learning rate and schedule

The learning rate scales parameter updates. In a pretrained model, too large a rate can destroy useful behavior quickly. In a randomly initialized model, too small a rate can make learning inefficient. The appropriate range depends on architecture, optimizer, batch, numeric precision, initialization, data, and which parameters are trainable.

Warmup gradually increases the learning rate during early updates. It can help avoid abrupt early changes while optimizer statistics develop. A schedule then changes the rate over time, such as linear decay or cosine decay. A schedule is usually defined over optimizer steps, not the number of microbatches. If you change accumulation, check whether the scheduler still covers the intended number of updates.

A resume must preserve scheduler position and its original total-step plan if exact continuation is intended. Extending a run from 1,000 to 2,000 steps while recomputing a cosine schedule can change the rate at the same nominal step. That may be a valid new experiment, but it is not an invisible continuation of the old one.

Do a short learning-rate comparison using a fixed validation procedure. If loss becomes nonfinite or jumps sharply, stop and inspect rather than allowing the run to continue indefinitely. A smaller rate may help, but first rule out malformed inputs, empty supervision, bad masks, or numeric problems.

Optimizer parameters

AdamW's beta1 controls the smoothing of the first-moment estimate, a running estimate related to gradient direction. beta2 controls smoothing of the second-moment estimate, related to squared gradient magnitude. Values close to one retain longer history. epsilon is a small stabilizing term in the denominator. These values affect update behavior; they are not accuracy thresholds.

Weight decay applies a shrinking tendency to parameters. It can act as regularization, but its effect depends on learning rate and which parameters receive it. Many recipes exclude bias and normalization parameters from decay. Copying only the headline decay value while changing parameter groups is not the same experiment.

Gradient clipping limits the norm or magnitude of a gradient before the update. Global-norm clipping preserves direction while reducing the size of an unusually large gradient vector. Log the norm before clipping; constant heavy clipping can signal an unstable setup. Clipping is a guardrail, not proof that an otherwise bad learning rate is safe.

In FP16 mixed precision with a gradient scaler, unscale gradients before clipping. Keep scaling consistent across accumulated microbatches and update the scaler at effective-batch boundaries. PyTorch's AMP examples spell out these ordering requirements. BF16 often uses a different policy, so do not copy an FP16 scaler loop mechanically. F12

Precision policy and autocast

Autocast chooses lower-precision computation for supported operations while preserving higher precision where the framework's policy requires it. It does not necessarily change the dtype of stored model parameters or optimizer state. The byte-transformer and pointer examples explicitly keep FP32 parameters and AdamW state while optionally using BF16 autocast on compatible CUDA hardware.

FP16 and BF16 both use two bytes, but divide their bits differently between range and precision. BF16 has a much wider exponent range than FP16 and lower significand precision. A model that works with one can fail or behave differently with the other. FP16 training often uses gradient scaling to keep small gradients representable; the scaler and clipping order must be correct. Do not assume a GPU supports a format merely because its model name contains RTX. Check the actual device and backend.

A quantized four-bit frozen base is a separate storage strategy. It does not mean that all activations, adapters, gradients, or optimizer states are four-bit. Write down the dtype of each state category when budgeting or comparing recipes. F12

Dropout and other regularization

Dropout randomly suppresses selected activations during training according to a probability. It should be disabled during ordinary evaluation. Its purpose is not to make outputs random at inference. A model with dropout left active can appear inconsistent even under deterministic decoding.

Data augmentation creates modified training inputs that should preserve the intended target. Flipping an image may preserve an object label but invalidate left-versus-right labels. Cropping can remove the target object. Adding noise to audio may preserve event identity up to a point but destroy intelligibility. An augmentation is valid only when its label-preservation assumption is valid.

Early stopping chooses a checkpoint when validation performance stops improving under a predefined rule. Patience is the number of evaluations tolerated without sufficient improvement. Choose the monitored metric and minimum meaningful change in advance. If your validation set is small and noisy, reacting to every tiny fluctuation can select unstable results.

Regularization cannot compensate for a fundamentally incorrect target. A model trained on mislabeled masks will not learn the intended boundary because dropout is well chosen. Inspect data before making optimization more elaborate.

Attention masks and loss masks solve different problems

An attention mask controls which input positions can influence each representation. A causal mask prevents future-token access. A padding mask excludes nonexistent padded input positions from attention. A custom branch mask may isolate independent questions sharing a context.

A loss mask controls which predictions contribute to the training objective. In assistant-only fine-tuning, prompt positions may be visible to attention while excluded from loss. Padding targets are commonly assigned an ignore index such as -100 for cross-entropy. That value is a convention of the loss implementation, not a token that should be inserted into the actual input vocabulary. F13

A training batch can have a correct attention mask and an incorrect loss mask. For example, it may teach the model to reproduce user prompts rather than only the assistant response. It can also contain no supervised tokens after truncation, producing meaningless or nonfinite loss. Count valid target positions and inspect decoded supervised spans.

Do not blindly set all occurrences of an end-of-sequence token to ignored when padding uses the same ID. Real end-of-sequence targets may then disappear from supervision, weakening stopping behavior. Construct masks from actual padding positions and the intended response spans, not from a simplistic equality test that erases legitimate tokens.

Packing and truncation are modeling decisions

Packing combines short examples to use a context window efficiently. If unrelated examples can attend to each other, the training distribution includes cross-example context that may not exist in deployment. Some objectives allow a continuous stream; others need boundaries or block-diagonal attention. Know which your trainer implements.

Truncation discards tokens beyond a length limit. It can remove the answer, the tool schema, or the evidence needed to answer. A model trained on such examples may appear to learn impossible tasks. Log how many records are truncated, from which side, and how many supervised tokens remain. Reject or deliberately transform examples when the task would no longer make sense.

Padding extends shorter sequences to a common length with placeholder positions. It adds computation unless the implementation handles variable lengths efficiently. Bucketing similar lengths can reduce waste. However, a padding optimization is not worth a silent change to target semantics.

Epochs steps and stopping

An epoch means one pass through a defined dataset, usually after shuffling. Streaming or weighted sampling may not have a simple epoch interpretation. Steps count parameter updates. Samples and tokens count presentations. Report the quantities that actually define your experiment.

If a dataset has 10,000 examples and an effective batch of 20 with no dropped remainder, one epoch is about 500 updates. With different-length examples, the token count per epoch is another useful quantity. If sampling with replacement, a given example may appear more than once before another is seen.

More epochs can improve fit and then worsen generalization. A small instruction dataset repeated many times can teach formatting quickly while causing memorization or forgetting. Preserve checkpoints and compare outputs. Do not assume the last checkpoint is best because it consumed the most computation.

A controlled tuning sequence

First verify one-batch overfitting on a tiny clean subset. The model should be able to reduce loss substantially when repeatedly shown the same examples, unless the objective or architecture makes that impossible. Failure here points to pipeline, gradient, target, or capacity issues. Success proves only that the model can fit those examples.

Next run a short representative baseline configuration. Change one major factor at a time: learning rate, data quality, number of updates, trainable modules, or context length. Record the hypothesis for each change. Prefer a small number of interpretable comparisons over a large blind sweep that you cannot afford to evaluate properly.

Finally select the candidate using validation metrics and guardrails, then run the final test once. Save the failed experiments' summaries too. Knowing which assumptions failed is part of the reusable result.

Exercises

A run has microbatch size 2 and accumulation 8. It logs 300 optimizer steps. How many example presentations occurred if every microbatch is full? The answer is 4,800. Explain why that does not tell you the number of unique examples.

Create two artificial microbatches with different numbers of supervised tokens. Calculate the difference between averaging their mean losses and averaging all token losses. Decide which objective your project intends.

Print one example's attention mask and supervised target positions. Explain why the model may attend to a user question while receiving no loss for reproducing that question.