PyTorch Training Loops

A training run reports steadily dropping loss, yet its saved incident router behaves erratically the moment you deploy it to staging. Nothing threw an exception. No crash log appeared in your terminal. Behind that quiet execution, one of several silent state bugs occurred: gradients carried over between mini-batches, Dropout stayed active during validation, or saving a live reference meant subsequent epochs quietly mutated your supposedly frozen checkpoint. We'll build an automated ticket router from scratch and trace every one of these state transitions directly.

Suppose an upstream NLP feature pipeline converts each incoming customer ticket into two normalized numerical features: an urgency score and a severity-language score. A binary target indicates whether the ticket requires specialist escalation (1) or standard routing (0). Untrained linear weights guess no better than a coin flip. A training loop passes examples forward through the model, measures prediction error with a loss function, calculates how each individual weight contributed to that error, and updates the weights in the direction of steepest descent.

The earlier causal reasoning lesson explains how to test whether a deployed model update actually improves user outcomes. An A/B test provides that evaluation. Neither technique builds the model itself. First, the loop has to learn it.

PyTorch doesn't hide this learning cycle behind a monolithic .fit() call: you write the loop yourself. That imperative design makes every tensor, gradient, and memory buffer inspectable at runtime. Later, when you scale up to fine-tune a large language model or run distributed training across multiple nodes, the fundamental handoffs remain identical: batches, logits, loss, gradients, optimizer updates, validation passes, and durable checkpoints.^{[1]Reference 1PyTorch: An Imperative Style, High-Performance Deep Learning Library.https://arxiv.org/abs/1912.01703[2]Reference 2Optimizing Model Parameters.https://docs.pytorch.org/tutorials/beginner/basics/optimization_tutorial.html}

The gradient descent and cross-entropy concepts build on backpropagation and softmax classification. Here we implement them in executable PyTorch code. Each code example runs self-contained on CPU unless explicitly noted otherwise. The small numbers are designed to expose tensor mechanics clearly rather than optimize classification accuracy.

These examples were validated on PyTorch 2.x on CPU. Random seeds make the fixtures repeatable, but floating-point numbers can vary across different architectures or CUDA backends. Check the stated invariants, such as monotonic loss reduction or matching reloaded logits, rather than treating a last-decimal discrepancy as an error.^{[3]Reference 3Reproducibilityhttps://docs.pytorch.org/docs/stable/notes/randomness.html}

Turn a ticket into tensors

We'll bypass raw text tokenization for now. The feature extraction pipeline has converted each incident ticket into two normalized floating-point numbers:

Every training instance consists of a two-element feature vector and an integer class target. A mini-batch stacks those rows into a two-dimensional tensor. Six incident tickets produce a feature tensor of shape (6, 2) and a label tensor of shape (6,).

That output represents a strict interface contract. Floating-point features (torch.float32) pass through linear matrix multiplications, while integer targets (torch.long, equivalent to 64-bit integers) supply class indices. Mismatched shapes or floating targets will cause downstream loss functions to fail immediately.

Now that inputs and targets match the expected shapes, inspect how the model scores each ticket before introducing optimization.

Logits are scores, not probabilities

A linear classification head ending in nn.Linear(2, 2) produces two unconstrained scores for each ticket:

These unnormalized outputs are logits. Logits can be positive, negative, or zero, and their values don't sum to one. For multi-class classification, PyTorch's nn.CrossEntropyLoss expects these raw logits alongside integer target indices. Under the hood, it combines log-softmax with negative log-likelihood in a single, numerically stabilized operation.^{[4]Reference 4CrossEntropyLosshttps://docs.pytorch.org/docs/stable/generated/torch.nn.CrossEntropyLoss.html}

Consider raw ticket logits [2.0, -1.0] for an example whose ground-truth class is 0. To see how cross-entropy scores this prediction, convert the logits to probabilities:

$$p_0 = \frac{e^{2.0}}{e^{2.0} + e^{-1.0}} = \frac{7.389}{7.389 + 0.368} \approx 0.9526,\qquad -\log(p_0) \approx 0.0486.$$

The second ticket outputs logits [-0.5, 1.5] with ground-truth class 1. The correct class wins, but with less margin: $p_1 \approx 0.8808$, yielding loss $-\log(p_1) \approx 0.1269$. The average loss over both tickets is $(0.0486 + 0.1269) / 2 \approx 0.0878$. PyTorch's default reduction="mean" computes exactly this batch average.

Failure case: Don't pass torch.softmax(logits, dim=-1) into CrossEntropyLoss. The loss function expects unnormalized logits. If you pass probabilities, CrossEntropyLoss treats them as raw scores and applies a second softmax normalization, distorting the objective and corrupting gradient calculations. A naive manual calculation like torch.log(torch.softmax(logits)) can also overflow or underflow when logits are large or small, whereas CrossEntropyLoss uses the log-sum-exp trick internally: $\log \sum_j e^{z_j} = c + \log \sum_j e^{z_j - c}$, where $c = \max_j z_j$.

Why should CrossEntropyLoss receive raw logits rather than probabilities from a separate softmax?

Answer

It combines log-softmax and negative log-likelihood in a numerically stable operation using the log-sum-exp trick. Applying softmax first makes the loss treat probabilities as unnormalized scores and apply another normalization, altering the objective and its gradients.

CrossEntropyLoss summarizes batch predictions into a single scalar value. That scalar quantifies the model's current error, but it hasn't changed any weights. To update parameters, we need to trace one batch through the training loop's five execution phases.

The five execution phases in a training step

Every iteration in a PyTorch training step executes five distinct phases in a precise order. Before checking the table below, determine which specific step actually alters the model's weight tensors:

Calling loss.backward() doesn't modify weight values. It traverses the dynamic computation graph in reverse via the chain rule, computing partial derivatives of the scalar loss with respect to every trainable leaf tensor (requires_grad=True), and stores those values in each parameter's .grad field.^{[5]Reference 5Deep Learning.https://www.deeplearningbook.org/[6]Reference 6Autograd mechanicshttps://docs.pytorch.org/docs/stable/notes/autograd.html}

The parameter values themselves remain completely unchanged until optimizer.step() runs. The optimizer reads the accumulated gradients, applies its update rule (such as SGD or Adam), and updates the parameter storage in place.

Why use set_to_none=True in optimizer.zero_grad()? Setting .grad = None rather than filling tensors with zeros provides three practical advantages:

1. Memory efficiency: Setting .grad = None immediately deallocates the gradient tensor memory, lowering peak memory consumption between iterations.
2. Autograd efficiency: When a parameter's .grad is None, the autograd engine during backward() simply assigns the newly computed gradient tensor (p.grad = grad) instead of performing an in-place addition (p.grad.add_(grad)), saving memory bandwidth and kernel overhead.^{[7]Reference 7Optimizer.zero_gradhttps://docs.pytorch.org/docs/stable/generated/torch.optim.Optimizer.zero_grad.html}
3. Unused parameter behavior: In architectures with conditional routing (such as mixture of experts), parameters not used in a forward pass keep .grad = None. Optimizers like Adam skip parameters with None gradients, whereas a zero-filled tensor would cause Adam's weight decay and momentum buffers to update those parameters anyway.

Trace this step with exact arithmetic. Consider two symmetric tickets: ticket [1.0, 1.0] has label 1 (escalate), and ticket [-1.0, -1.0] has label 0 (standard). Let $W$ be a 2x2 weight matrix where row 0 scores class 0 and row 1 scores class 1. Initializing $W = 0$ with zero bias yields logits [0.0, 0.0] for both tickets. The predicted class probabilities are [0.5, 0.5], producing a cross-entropy loss of $-\log(0.5) \approx 0.6931$ for each row.

For cross-entropy loss with mean reduction, each row's gradient contribution to weight $W_{c,j}$ is $(p_c - y_c) \cdot x_j / N$, where $N=2$ is the batch size. Both tickets produce positive gradient values for class 0 weights and negative gradient values for class 1 weights, yielding an aggregate gradient matrix $dW = \begin{bmatrix} 0.5 & 0.5 \\ -0.5 & -0.5 \end{bmatrix}$. The bias gradients cancel to zero. loss.backward() leaves $W$ at zero while writing $dW$ into $W\text{.grad}$. Finally, SGD with learning rate $\eta = 0.2$ updates the weights:

$$W_1 = W_0 - \eta \cdot dW = \begin{bmatrix} 0 & 0 \\ 0 & 0 \end{bmatrix} - 0.2 \cdot \begin{bmatrix} 0.5 & 0.5 \\ -0.5 & -0.5 \end{bmatrix} = \begin{bmatrix} -0.1 & -0.1 \\ 0.1 & 0.1 \end{bmatrix}.$$

Let's execute this exact mathematical calculation in PyTorch to confirm every state transition.

The output confirms the state boundaries: loss.backward() populated .grad without touching parameter values, and optimizer.step() moved the parameters. What happens if leftover gradients from an earlier iteration remain in .grad when the next backward() executes?

Clearing gradients and micro-batch accumulation

PyTorch adds newly computed gradients into existing .grad buffers rather than overwriting them (param.grad += new_grad). That accumulation default is intentional: it enables gradient accumulation, letting multiple micro-batches build up an effective batch update when memory is constrained. In standard training where each mini-batch should trigger an independent step, failing to clear gradients causes stale derivatives to pollute future updates.

To observe this accumulation directly, differentiate $y = w^2$ twice at $w = 1.0$. The derivative is $\frac{dy}{dw} = 2w = 2.0$. Calling backward() a second time without clearing gradients adds another $2.0$, producing an accumulated gradient of $4.0$. Resetting weight.grad = None restores the clean single-step value:

Without clearing, the derivative silently doubles. Calling optimizer.zero_grad(set_to_none=True) inside the batch loop ensures each update depends solely on its own data.^{[7]Reference 7Optimizer.zero_gradhttps://docs.pytorch.org/docs/stable/generated/torch.optim.Optimizer.zero_grad.html}

Gradient accumulation across micro-batches

When a model or batch size exceeds available GPU memory, gradient accumulation lets you simulate a larger effective batch size $B_{\text{eff}} = K \times B_{\text{micro}}$ using $K$ smaller micro-batches of size $B_{\text{micro}}$.

The loss function usually averages errors over the current micro-batch: $L_k = \frac{1}{B} \sum_{i=1}^B \ell(x_{k,i}, y_{k,i})$. The true mean loss over the combined batch of $K \cdot B$ examples is:

$$L_{\text{eff}} = \frac{1}{K \cdot B} \sum_{k=1}^K \sum_{i=1}^B \ell(x_{k,i}, y_{k,i}) = \frac{1}{K} \sum_{k=1}^K \left( \frac{1}{B} \sum_{i=1}^B \ell(x_{k,i}, y_{k,i}) \right) = \frac{1}{K} \sum_{k=1}^K L_k.$$

Because autograd sums gradients across multiple backward() calls, you must divide each micro-batch loss by $K$ before backpropagating: (loss / K).backward(). If you skip this division, the accumulated gradient will be $K$ times larger than the true average, which acts like multiplying your learning rate by $K$ and frequently causes training to diverge.^{[8]Reference 8Automatic Mixed Precision Exampleshttps://docs.pytorch.org/docs/stable/notes/amp_examples.html}

In distributed training using DistributedDataParallel (DDP), every backward() call triggers an AllReduce network synchronization across all GPUs. Calling AllReduce on every micro-batch wastes network bandwidth. PyTorch provides the model.no_sync() context manager to suppress gradient communication during the first $K-1$ micro-batches, synchronizing gradients across nodes only on the final step right before optimizer.step().^{[9]Reference 9DistributedDataParallelhttps://docs.pytorch.org/docs/stable/generated/torch.nn.parallel.DistributedDataParallel.html}

What happens if optimizer.zero_grad() runs only once per epoch instead of once per mini-batch in standard training?

Answer

Each backward pass adds newly computed derivatives into existing .grad buffers. Later updates combine current gradients with stale gradients from earlier batches, destabilizing optimization. Reset inside the batch loop unless accumulation is deliberate.

One update on a synthetic fixture verifies autograd mechanics, but it doesn't constitute a complete loop. We need abstractions to manage datasets and partition rows into mini-batches.

Batching and row-weighted epoch loss

A dataset stores samples and their corresponding labels. A data loader coordinates batching, worker multiprocessing, and shuffling between epochs. Shuffling ensures the optimizer encounters diverse gradient directions across iterations, preventing cyclical oscillation.

Using our six tickets with batch_size=4, the data loader yields two batches: the first contains four rows, and the second contains the remaining two rows.

The final batch contains only two rows. When calculating average loss across an epoch, don't simply average the reported batch loss numbers. Because batch sizes differ, taking an unweighted mean of batch averages assigns excessive influence to the smaller batch.

To compute the exact epoch loss, weight each batch mean by its row count: sum the batch losses multiplied by their batch sizes, then divide by the total number of rows across the entire dataset. For instance, batches of size 4 and 2 with average losses 0.20 and 0.50 yield a true epoch average of $(4 \times 0.20 + 2 \times 0.50) / 6 = 0.30$, whereas an unweighted average of the two numbers would report $0.35$.^{[4]Reference 4CrossEntropyLosshttps://docs.pytorch.org/docs/stable/generated/torch.nn.CrossEntropyLoss.html}

With data loading established, we can construct a minimal training loop. Before evaluating generalization on held-out data, we test whether the model can memorize a small, clean batch.

Overfitting one batch as an isolation test

Before training on large datasets, run an overfit-one-batch test. Take a single batch of clean, separable examples and verify that the model drives training loss to near zero. If the model can't memorize eight clean rows, don't debug data distributions or regularization: inspect label encodings, tensor shapes, loss configurations, learning rate scale, and the order of optimizer operations.

We'll define a multi-layer classifier: nn.Linear(2, 8) expands the two inputs to eight hidden activations, nn.ReLU() introduces non-linearity by zeroing negative values, and nn.Linear(8, 2) produces class logits.

The loss drops from 0.6521 to 0.0056 and classification accuracy reaches 100%. This passes the basic wiring check. Now we can separate training from held-out validation.

Isolating validation from training updates

Memorizing training rows proves that gradients flow through the model, but it provides zero evidence that the classifier generalizes to unseen tickets. To assess generalization, split the workflow into two phases:

• Training phase: Mini-batches flow forward, autograd constructs the computation graph, loss.backward() calculates derivatives, and optimizer.step() updates parameters.
• Validation phase: Held-out examples evaluate the current weights without updating parameters or recording gradient history.

Validation relies on two independent controls:

The script below trains on eight tickets and tracks validation loss across four held-out tickets over 60 epochs.

The checkpoint snapshot aliasing trap

Notice the call to copy.deepcopy(model.state_dict()) in line 432 above. A common bug is writing best_state = model.state_dict(). Calling model.state_dict() returns a dictionary containing direct references to the live parameter tensors. As future epochs update model weights, those changes mutate the referenced tensors. Your saved best_state quietly drifts toward the final epoch's weights rather than preserving the best checkpoint.^{[10]Reference 10Saving and Loading Modelshttps://docs.pytorch.org/tutorials/beginner/saving_loading_models.html}

We can demonstrate this reference leak directly:

The validation loop called model.eval() even though the model lacked Dropout or BatchNorm layers. Adding a mode-sensitive layer makes the purpose of that call immediately visible.

Operational modes and context managers

In a network composed strictly of linear transformations and ReLU activations, model.train() and model.eval() yield identical outputs because neither layer maintains mode-dependent state. In production architectures, however, several layers change their behavior across modes:

• Dropout: During training (train()), Dropout randomly zeroes activations with probability $p$ and scales surviving activations by $1 / (1 - p)$ to preserve expected magnitude. During evaluation (eval()), Dropout acts as an identity pass-through.
• BatchNorm: During training, BatchNorm computes mean and variance across the current mini-batch and updates its exponential moving average running buffers (running_mean, running_var). During evaluation, it freezes buffer updates and normalizes activations using the accumulated running statistics.
• LayerNorm: Unlike BatchNorm, LayerNorm normalizes across the hidden channel dimensions of each sample independently. Its forward pass is identical in both train() and eval().

The snippet below introduces Dropout to our classifier, runs the same input ticket twice in training mode, and then runs it twice in evaluation mode.

The output illustrates the separation of responsibilities: model.eval() controls layer-specific operational logic, while torch.no_grad() controls autograd tape construction. Calling one without the other addresses only half of validation.

Why do validation loops require both model.eval() and torch.no_grad()?

Answer

model.eval() switches mode-sensitive layers (like Dropout and BatchNorm) into deterministic evaluation behavior. torch.no_grad() separately stops the autograd engine from constructing a computation graph. Neither replaces the other.

Saving weights to disk produces the deployable artifact. Checkpoints also protect ongoing training against sudden hardware faults and preemption.

Resilient checkpointing and atomic state snapshots

A trained model isn't production-ready until its parameters and configurations are saved in a durable, reloadable format. Never save the Python model instance directly using torch.save(model, path). That approach uses Python's pickle to serialize class definitions and file paths, binding the checkpoint to your exact codebase structure.

Save the model's state dictionary (state_dict) instead. A state_dict is an ordered Python dictionary mapping parameter and buffer names to their underlying tensor storage.

The complete resumption bundle

If training will resume later, saving only model weights is insufficient. A complete checkpoint must capture all state required for exact continuation:

1. model_state_dict: Trainable parameters and persistent buffers (such as BatchNorm running stats).
2. optimizer_state_dict: Optimization state, including Adam's first moment ($m$) and second moment ($v$) vectors alongside step counts. If you omit optimizer state upon resumption, momentum resets to zero, causing abrupt parameter oscillations that degrade convergence.^{[10]Reference 10Saving and Loading Modelshttps://docs.pytorch.org/tutorials/beginner/saving_loading_models.html}
3. lr_scheduler_state_dict: Current learning rate step counter and cycle phase.
4. scaler_state_dict: The dynamic loss scale factor and consecutive finite step count when using Automatic Mixed Precision.
5. rng_state: Random number generator states across PyTorch CPU, CUDA (torch.cuda.get_rng_state_all()), Python random, and NumPy to ensure exact data shuffling and dropout repeatability upon resumption.^{[3]Reference 3Reproducibilityhttps://docs.pytorch.org/docs/stable/notes/randomness.html}
6. Metadata: The epoch index, global step count, validation metrics, git commit hash, and dataset split identifier.

Preventing corruption with atomic file writes

On compute clusters with spot instances or preemptible cloud nodes, a training process can be terminated without warning. If preemption strikes while torch.save(checkpoint, "checkpoint.pt") is mid-write, the target file is left truncated. When the job restarts, torch.load crashes with EOFError: Ran out of input, permanently destroying days of training progress.

To make checkpointing resilient against preemption, write atomically:

1. Serialize the state bundle to a temporary staging file in the same filesystem directory (checkpoint.pt.tmp).
2. Flush application buffers and invoke os.fsync(file_descriptor) to guarantee the operating system commits write caches to physical storage.
3. Call os.replace("checkpoint.pt.tmp", "checkpoint.pt"). On POSIX systems, os.replace is an atomic filesystem operation. At any instant, the filesystem points either to the previous valid checkpoint or the new complete checkpoint, eliminating the risk of corrupted files.

In PyTorch 2.6 and later, torch.load(..., weights_only=True) is the default standard, restricting unpickling to tensors, primitive values, and basic containers to prevent arbitrary code execution vulnerabilities.^{[11]Reference 11Serialization Semanticshttps://docs.pytorch.org/docs/stable/notes/serialization.html}

Testing reloaded logits against original outputs using torch.testing.assert_close verifies that parameter tensors loaded faithfully.

Diagnosing failures and bounding gradients

When a training script runs without throwing exceptions, subtle bugs can still compromise learning. Use this triage matrix to isolate root causes:

Five-step numerical triage when loss becomes NaN

When loss becomes non-finite, don't immediately reduce the learning rate blindly. Follow this ordered diagnostic sequence to find the earliest invalid value:

1. Inputs: Check raw batch inputs (torch.isfinite(xb).all()). Unsanitized nulls or extreme feature values corrupt the forward pass immediately.
2. Logits: Check model output scores (torch.isfinite(logits).all()). If inputs are finite but logits are NaN, examine numerical stability in custom layers or activations.
3. Loss: Check loss output (torch.isfinite(loss)). If logits are finite but loss is Inf, verify class target indices and label smoothing parameters.
4. Gradients: After backward(), check parameter gradients (torch.isfinite(p.grad).all()). If loss was finite but gradients exploded, examine deep layer backpropagation paths.
5. Parameters: After optimizer.step(), check updated weights (torch.isfinite(p).all()).

Bounding gradients with global norm clipping

When dealing with deep recurrent networks or transformer architectures, gradient vectors can intermittently surge in magnitude. Gradient clipping enforces an upper bound on gradient size before the optimizer applies parameter updates.

PyTorch's nn.utils.clip_grad_norm_ calculates the total $\ell_2$ norm across all model parameters concatenated into a single vector:

$$g_{\text{norm}} = \sqrt{\sum_{p \in \Theta} \|\nabla_p L\|_2^2}.$$

If $g_{\text{norm}}$ exceeds max_norm, the function scales all gradients down proportionally:

$$\nabla_p L \leftarrow \nabla_p L \cdot \frac{\text{max\_norm}}{\max(g_{\text{norm}}, \text{max\_norm})}.$$

Global vector norm clipping preserves the exact direction of the gradient vector in parameter space, adjusting only its step magnitude. In contrast, per-parameter value clipping (clip_grad_value_) truncates individual elements independently, which distorts the search trajectory.^{[12]Reference 12clip_grad_norm_https://docs.pytorch.org/docs/stable/generated/torch.nn.utils.clip_grad_norm_.html}

Setting error_if_nonfinite=True causes clip_grad_norm_ to raise a runtime error immediately if any gradient is NaN or Inf, halting training before non-finite updates corrupt model weights.

Inference mode versus no-grad contexts

During deployment or batch inference, the model generates predictions without ever calculating gradients. Autograd tracking is entirely unnecessary.

PyTorch offers two context managers for non-gradient execution:

• torch.no_grad(): Disables dynamic graph recording, reducing compute and memory usage. However, it continues updating internal tensor version counters used for in-place mutation checks.
• torch.inference_mode(): Disables both autograd graph recording and tensor version counter tracking. This provides higher execution speed and lower memory overhead, making it the preferred context manager for production serving, validation inference, and token generation loops.^{[6]Reference 6Autograd mechanicshttps://docs.pytorch.org/docs/stable/notes/autograd.html}

Hardware accelerators alter this floating-point arithmetic to unlock higher training throughput.

Automatic mixed precision and dynamic loss scaling

Standard deep learning models perform arithmetic using 32-bit floating-point numbers (float32). On modern GPU architectures (such as NVIDIA Volta, Ampere, and Hopper), Tensor Cores execute 16-bit matrix multiplications with significantly higher throughput while cutting activation memory in half.

Automatic Mixed Precision (AMP) matches operations to their ideal precision: compute-intensive matrix multiplications run in lower precision (float16 or bfloat16), while sensitive operations (like softmax and reductions) remain in float32.^{[13]Reference 13Mixed Precision Training.https://arxiv.org/abs/1710.03740[14]Reference 14Automatic Mixed Precision package - torch.amphttps://docs.pytorch.org/docs/stable/amp.html}

The float16 underflow hazard and GradScaler

Standard float16 reserves 5 bits for its exponent and 10 bits for its mantissa. That structure limits its minimum representable positive normal value to $2^{-14} \approx 6.1 \times 10^{-5}$. During backpropagation, parameter gradients frequently fall below $10^{-5}$. In pure float16, these tiny values underflow directly to zero, starving the optimizer of updates.

torch.amp.GradScaler resolves underflow through dynamic loss scaling:

1. Forward pass: Operations execute inside torch.amp.autocast('cuda'). CrossEntropyLoss runs in float32 for numerical stability.
2. Loss scaling: The scaler multiplies the scalar loss by a large scale factor $S$ (initial default $2^{16} = 65,536$). By linearity of differentiation, $\nabla_W (S \cdot L) = S \cdot \nabla_W L$. This shifts gradient values upward by $S$, keeping them within the representable range of float16.
3. Backward pass: Calling scaler.scale(loss).backward() calculates scaled gradients ($S \cdot \nabla W$).
4. Unscaling before clipping: Before calling clip_grad_norm_, invoke scaler.unscale_(optimizer). This divides stored gradients by $S$ so they reflect true magnitude. Order is critical here: if you clip before unscaling, you evaluate $S \cdot \nabla W$ against max_norm, effectively clipping true gradients at $\text{max\_norm} / S$ and crushing updates to near zero!
5. Conditional optimizer step: scaler.step(optimizer) inspects unscaled gradients. If all gradients are finite, it invokes optimizer.step(). If any gradient contains Inf or NaN (due to floating-point overflow), the scaler skips the optimizer step entirely, shielding weights from corruption.
6. Dynamic scale update: scaler.update() adjusts the scale factor. If the step was skipped, it halves the scale ($S \leftarrow S \times 0.5$) and resets the consecutive clean counter. If training proceeds for 2,000 consecutive steps without non-finite gradients, it doubles the scale factor ($S \leftarrow S \times 2.0$) to maintain maximum numerical precision.^{[8]Reference 8Automatic Mixed Precision Exampleshttps://docs.pytorch.org/docs/stable/notes/amp_examples.html}

When training with bfloat16 on compatible hardware, the format's 8-bit exponent matches float32's dynamic range ($\sim 10^{38}$), which eliminates gradient underflow. Consequently, pure bfloat16 training typically omits GradScaler entirely.

On CPU, enabled=False acts as a transparent pass-through: autocast is a no-op, scaler.scale(loss) returns the raw loss, unscale_ does nothing, and scaler.step(optimizer) executes optimizer.step() directly.

Graph compilation after eager verification

torch.compile optimizes PyTorch execution by capturing Python operations into a computational graph via TorchDynamo and compiling them into fused kernels via TorchInductor.^{[15]Reference 15PyTorch 2: Faster Machine Learning Through Dynamic Python Bytecode Transformation and Graph Compilationhttps://doi.org/10.1145/3620665.3640366[16]Reference 16torch.compilehttps://docs.pytorch.org/docs/main/generated/torch.compile.html}

Compilation accelerates execution, but it doesn't fix logic errors. If your eager loop has misordered optimizer steps, unhandled NaN gradients, or leaky validation splits, torch.compile will simply compile those bugs. Always verify your loop's eager execution and numerical correctness before applying compilation.

We test graph capture below using backend="eager", which runs the captured graph through standard PyTorch operations to verify graph validity without requiring a GPU compiler.

If compiled outputs diverge from eager outputs, look for unsupported Python control flow or in-place tensor mutations. Graph breaks hurt runtime performance rather than producing numerical errors. Profile steady-state loop iterations separately to exclude one-time compilation warm-up overhead.^{[17]Reference 17Troubleshooting torch.compilehttps://docs.pytorch.org/docs/main/user_guide/torch_compiler/torch.compiler_troubleshooting.html}

Artifacts and audit receipts for the training run

Before deploying a model to production, assemble an audit trail verifying training integrity:

Practice: break and repair the loop

Solidify your understanding of loop mechanics by intentionally breaking individual steps in our runnable examples. Predict the failure symptom before running each test:

1. In 03-one-update.py, comment out optimizer.step(). Predict whether changed by step remains True.
2. In 04-zero-grad.py, invoke backward() three times before clearing gradients. Predict the resulting gradient magnitude.
3. In 07b-checkpoint-snapshot.py, replace copy.deepcopy(...) with model.state_dict(). Predict which assertion fails when live weights change.
4. In 08-train-eval-modes.py, leave the model in training mode during the second pair of predictions. Predict whether repeated forward passes match.
5. In 12-amp-ordering-pattern.py, explain why calling scaler.unscale_(optimizer) before clip_grad_norm_ is mandatory.
6. In 13-compile-after-eager.py, modify an eager weight after calling torch.compile and predict whether compiled output changes.

Expected observations

1. Without optimizer.step(), autograd populates .grad, but parameter values remain unchanged (changed by step: False).
2. The gradient accumulates to 6.0: each backward pass adds another derivative of 2.0 into .grad.
3. copied state stayed at zero evaluates to False. Both dictionaries reference live parameter storage, so both follow the weight update.
4. Dropout zeroes random neurons on each pass in training mode, causing repeated forward passes on identical inputs to disagree.
5. GradScaler multiplies loss by $S$ before backpropagation. Clipping before unscale_ compares $S \cdot \nabla W$ against max_norm, effectively clipping true gradients at $\text{max\_norm} / S$ and eliminating updates.
6. torch.compile wraps the original module rather than copying weights, so in-place edits to eager parameters reflect immediately in compiled outputs.