LEARN · NEURAL NETWORKS FROM THE GROUND UP
A training run can look perfectly healthy for thousands of steps and then suddenly produce a loss that is 100 times larger than expected. Sometimes the model recovers. Sometimes gradients become infinite, optimizer state becomes corrupted, and a multi-day training job becomes unusable.
This is one of the most expensive failure modes in modern deep learning.
The challenge is that a loss spike is a symptom, not a diagnosis. It might originate from an unusually difficult batch, exploding gradients, an aggressive learning rate, optimizer-state dynamics, or a floating-point overflow. Several of these failures can produce nearly identical loss curves while requiring very different fixes.
Three tools are particularly important:
-
Gradient clipping limits the magnitude of parameter gradients before an optimizer update.
-
Mixed-precision training accelerates computation but introduces numerical-range and rounding considerations.
-
Stability monitoring makes it possible to distinguish harmless fluctuations from failures that require intervention.
The difficult part is understanding how they interact. Gradient clipping cannot repair an infinity created during the forward pass. Loss scaling cannot extend the representable range of FP16 activations. And switching to BF16 does not prevent a large but finite gradient from destabilizing an optimizer update.
This tutorial develops a systematic approach to these problems, including a runnable PyTorch experiment that deliberately triggers FP16 overflow, demonstrates why changing precision alone is insufficient, and combines BF16 with gradient clipping to keep subsequent training steps stable.
1. Understand what a loss spike actually means
Consider a model trained with mini-batch gradient descent. At each optimizer step, it computes a loss, differentiates that loss with respect to its parameters, and updates the parameters.
A simplified update is:
θ_next = θ_current − learning_rate × gradient
Real optimizers such as AdamW use momentum and adaptive second-moment estimates, but the basic stability problem remains: if the effective update is much larger than the local region in which the model behaves predictably, training can deteriorate.
A sudden increase in loss can represent several different events.
A difficult but legitimate batch
Imagine training a retail demand forecasting model. Most batches contain ordinary products with stable sales. One batch contains promotional events with unusually large sales volumes.
The loss may jump because the examples are genuinely more difficult.
If gradients remain finite and subsequent batches return to their previous loss range, this may not require intervention.
A dangerous optimizer update
A batch produces a gradient norm far larger than usual. The optimizer makes a large parameter update, and the next batch produces an even larger loss.
This creates a feedback loop:
-
A parameter update moves the model into a poorly conditioned region.
-
Activations or prediction errors increase.
-
Gradients grow.
-
The next update moves the parameters even farther.
-
The model eventually diverges or generates non-finite values.
Clipping can help by limiting the gradient magnitude used for the update.
A numerical overflow
A tensor operation produces a value beyond the maximum representable number of its floating-point format.
For example, an FP16 matrix multiplication might produce an infinity even though the mathematically correct result is finite.
The loss calculation then encounters that infinity and may produce an infinite or NaN loss.
This is a fundamentally different failure: the computation is already invalid before gradient clipping runs.
A persistent training divergence
Sometimes the loss stays elevated even after the initial spike disappears.
Possible causes include:
-
Optimizer momentum carrying the effects of an abnormal gradient into later updates.
-
An unsuitable learning-rate schedule.
-
Changes in data distribution or batch composition.
-
Accumulated numerical errors.
-
Excessive updates to particularly sensitive model components.
The first diagnostic rule is therefore simple:
Never choose a stabilization technique based on the loss curve alone.
You need to inspect gradient norms, activation ranges, numerical precision, optimizer behavior, and the batch that preceded the incident.
2. Why floating-point precision changes training stability
Mixed-precision training uses lower-precision arithmetic for selected operations while retaining higher precision where it is numerically important.
Modern PyTorch provides automatic mixed precision through torch.autocast, which chooses supported lower-precision operations while allowing numerically sensitive operations to execute in FP32.
Two common formats are IEEE FP16 and BF16.
FP16 versus BF16
|
Property |
FP16 |
BF16 |
FP32 |
|---|---|---|---|
|
Total width |
16 bits |
16 bits |
32 bits |
|
Exponent bits |
5 |
8 |
8 |
|
Fraction bits |
10 |
7 |
23 |
|
Largest finite value |
65,504 |
~3.39 × 10³⁸ |
~3.40 × 10³⁸ |
|
Precision near 1 |
Higher than BF16 |
Lower than FP16 |
Highest |
|
Main advantage |
Fine relative precision for a 16-bit format |
Very large dynamic range |
Precision and numerical headroom |
The practical distinction is that BF16 provides approximately FP32’s exponent range, but with much less significand precision.
FP16 has more fraction bits than BF16, but its narrower exponent range makes it vulnerable to overflow when activations, intermediate products, or gradients become too large.
For example, consider the finite values:
-
30,000: representable in FP16.
-
60,000: representable in FP16.
-
100,000: beyond FP16’s finite range.
-
1,000,000: within BF16’s finite range.
BF16 values will be rounded more coarsely than FP16 values at comparable magnitudes, but large finite values do not immediately become infinity merely because they exceed 65,504.
This is why BF16 is frequently useful for deep networks with widely varying activation scales.
The trap: accumulating in FP32 does not always save FP16
Many accelerated matrix multiplications use higher-precision accumulation internally.
That helps reduce summation error, but it does not guarantee that the final FP16 output is representable.
A matrix multiplication can accumulate a finite result and still overflow when that result is converted to FP16 for storage.
Similarly, casting an already overflowing input into FP16 can produce an infinity before multiplication even begins.
The important distinction is between:
-
Input representation.
-
Multiplication precision.
-
Accumulation precision.
-
Output representation.
-
Precision used by the loss and backward pass.
Each stage can have different numerical behavior.
PyTorch’s current AMP documentation describes the operator-level casting rules and mixed-precision behavior.
3. Gradient clipping: controlling update magnitude
Gradient clipping modifies gradients before the optimizer uses them.
The most common strategy is global norm clipping, which rescales the gradients when their combined norm exceeds a specified threshold.
Suppose the gradients of all trainable parameters are concatenated into one large vector.
Its L2 norm is:
gradient_norm = √(g₁² + g₂² + … + gₙ²)
For a maximum allowed norm of 1.0, the clipping rule is:
scale_factor = min(1, 1.0 / gradient_norm)
Every gradient is then multiplied by the same scale factor.
For example:
|
Original norm |
Maximum norm |
Scale factor |
Resulting norm |
|---|---|---|---|
|
0.25 |
1.0 |
1.0 |
0.25 |
|
0.80 |
1.0 |
1.0 |
0.80 |
|
5.00 |
1.0 |
0.20 |
1.00 |
|
100.00 |
1.0 |
0.01 |
1.00 |
Global norm clipping preserves the direction of the combined gradient vector while reducing its magnitude.
That distinction matters.
Clipping by individual tensor values changes the relative proportions of gradient components and may substantially change the direction of the update.
Norm clipping is generally a more natural starting point for protecting against unusually large global updates.
The current PyTorch API
PyTorch provides torch.nn.utils.clip_grad_norm_ for global norm clipping.
The function:
-
Computes a combined norm over the selected parameter gradients.
-
Modifies gradients in place when clipping is necessary.
-
Returns the original, pre-clipping norm.
-
Supports an
error_if_nonfiniteargument to detect invalid norms.
The last two features are particularly useful for production monitoring.
A norm of 50 followed by clipping to 1.0 should not be logged simply as 1.0. You want to know that the model generated a gradient 50 times larger than the clipping threshold.
Similarly, clipping a gradient containing infinity is not a legitimate recovery technique. It cannot reconstruct a meaningful finite gradient.
PyTorch documents this behavior in its current gradient clipping API reference.
Why the clipping threshold is not universal
A maximum norm of 1.0 is a useful experimental starting point, but not a magic number.
The appropriate threshold depends on:
-
Model architecture and parameterization.
-
Loss normalization.
-
Effective batch size.
-
Optimizer and learning-rate schedule.
-
Gradient accumulation.
-
How gradient norms evolve throughout training.
If typical gradients have a norm of 0.02, clipping everything above 0.1 might be reasonable.
If typical gradients have a norm of 40, setting the threshold to 0.1 could interfere with almost every optimizer step.
A better approach is to record the distribution of pre-clipping gradient norms during a stable baseline run and choose the threshold based on those observations.
Also remember that AdamW does not transform gradients into parameter updates in the same way as SGD. A bound on the raw gradient norm is not an exact bound on an AdamW parameter update.
4. The mixed-precision trap: gradient scaling is not gradient clipping
FP16 has a second major numerical problem besides overflow: underflow.
Extremely small gradients may become zero when represented in FP16.
Mixed-precision loss scaling addresses this by multiplying the loss by a scale factor before backpropagation, making small gradients easier to represent.
The gradients are then divided by the same factor before the optimizer update.
PyTorch’s torch.amp.GradScaler automates this process for FP16 training.
It also detects non-finite gradients and can skip an optimizer update while reducing its scaling factor.
However, loss scaling does not guarantee that the forward pass will remain finite.
If a matrix multiplication has already generated infinity, multiplying or dividing the loss cannot restore the original finite result.
The correct order of operations
For FP16 training, the safe sequence is:
-
Clear previous gradients.
-
Execute the forward pass with autocasting.
-
Calculate the loss.
-
Scale the loss and run backpropagation.
-
Unscale the gradients.
-
Inspect or clip the unscaled gradients.
-
Let the scaler execute or skip the optimizer update.
-
Update the scaler.
Never clip scaled gradients before unscaling them.
Suppose the real gradient norm is 2.0 and the loss scale is 65,536.
The scaled gradient norm becomes approximately 131,072.
If you clip that scaled gradient to 1.0 and then unscale, the resulting norm is approximately 0.0000153, not 1.0.
The update is therefore much smaller than intended.
This is an easy bug to miss because the training loop runs without necessarily raising an exception.
The current PyTorch automatic mixed precision examples explicitly demonstrate unscaling before clipping.
5. Hands-on experiment: manufacture an FP16 loss spike
We will now build a small experiment that reproduces three distinct outcomes:
-
FP16 overflow produces an infinite loss, and GradScaler skips the invalid optimizer update.
-
BF16 prevents that overflow, but an enormous finite gradient still destabilizes training.
-
BF16 with gradient clipping keeps the update controlled, allowing normal training to resume.
The important feature is that all three experiments use the same model, data, initialization, and learning rate.
Only the precision and clipping strategy change.
Step 1: Prepare the environment
Use Python 3.10 or newer and an up-to-date PyTorch installation.
Create a virtual environment:
python -m venv .venv
Activate it on Linux or macOS:
source .venv/bin/activate
Install PyTorch:
python -m pip install --upgrade pip python -m pip install torch
For an accelerator-specific installation, select the appropriate build from the official PyTorch installation guide.
The example also supports CPU execution, so a GPU is not required.
On machines with a BF16-capable CUDA GPU, it can use that GPU. Otherwise, it falls back to CPU.
Step 2: Understand the synthetic problem
We will train a one-parameter linear regression model.
Its prediction is:
prediction = weight × input
The target is always zero.
The model begins with a weight of 0.25 and receives normal input values of 1.0.
Under ordinary conditions, the model steadily reduces its weight, minimizing mean squared error.
At step 6, we replace the normal input with 100,000.
This is an intentionally extreme outlier, similar in spirit to a corrupted or unexpectedly rescaled numerical feature.
Why use 100,000?
Because it exceeds the maximum finite FP16 value of 65,504.
When the linear operation receives that input under FP16 autocasting, the input conversion can overflow.
BF16, however, can represent the value within its exponent range.
We use plain SGD to make the relationship between gradient magnitude and parameter updates transparent. AdamW’s adaptive normalization would obscure this particular failure mechanism.
Step 3: Create the complete training script
Save the following as stability_demo.py.
import math import torch import torch.nn.functional as F STEPS = 12 SPIKE_STEP = 6 OUTLIER_X = 100_000.0 MAX_NORM = 1.0 LR = 0.02 def run_case(name, dtype, clipping): device = ( "cuda" if torch.cuda.is_available() and torch.cuda.is_bf16_supported() else "cpu" ) model = torch.nn.Linear( 1, 1, bias=False ).to(device) with torch.no_grad(): model.weight.fill_(0.25) optimizer = torch.optim.SGD( model.parameters(), lr=LR, ) use_scaler = dtype == torch.float16 scaler = torch.amp.GradScaler( device, enabled=use_scaler, ) print(f"\n{name} ({device})") print( "step | loss | grad_norm " "| weight | status" ) for step in range(STEPS): optimizer.zero_grad(set_to_none=True) x_value = ( OUTLIER_X if step == SPIKE_STEP else 1.0 ) x = torch.tensor( [[x_value]], device=device, dtype=torch.float32, ) target = torch.zeros_like(x) with torch.autocast( device_type=device, dtype=dtype, ): predictions = model(x) loss = F.mse_loss( predictions.float(), target, ) if use_scaler: scaler.scale(loss).backward() scaler.unscale_(optimizer) else: loss.backward() gradients = [ p.grad.detach().float() for p in model.parameters() if p.grad is not None ] raw_norm = torch.linalg.vector_norm( torch.stack([ torch.linalg.vector_norm(g) for g in gradients ]) ).item() if clipping and math.isfinite(raw_norm): torch.nn.utils.clip_grad_norm_( model.parameters(), max_norm=MAX_NORM, error_if_nonfinite=True, ) if use_scaler: scale_before = scaler.get_scale() scaler.step(optimizer) scaler.update() scale_after = scaler.get_scale() status = ( "AMP overflow: skipped" if scale_after < scale_before else "updated" ) elif not math.isfinite(raw_norm): status = "non-finite: skipped" else: optimizer.step() status = ( "clipped" if clipping and raw_norm > MAX_NORM else "updated" ) weight = model.weight.detach().float().item() print( f"{step:>4} | " f"{loss.item():>12.3e} | " f"{raw_norm:>12.3e} | " f"{weight:>12.3e} | " f"{status}" ) if __name__ == "__main__": torch.manual_seed(7) print("PyTorch:", torch.__version__) run_case( "FP16 + GradScaler, no clipping", torch.float16, False, ) run_case( "BF16, no clipping", torch.bfloat16, False, ) run_case( "BF16 + norm clipping", torch.bfloat16, True, )
Step 4: Execute the experiment
python stability_demo.py
The experiment completes quickly and requires no downloaded dataset.
The model parameters stay in FP32. Autocasting controls the precision of the linear operation, and the prediction is converted to FP32 for loss calculation.
That last detail is deliberate: we want to demonstrate overflow originating in the linear operation rather than creating a second, avoidable overflow in the squared-error calculation.
The script also avoids applying a normal optimizer update when its explicitly checked BF16 gradient norm is non-finite. In FP16 mode, GradScaler handles invalid-gradient detection and optimizer-step skipping.
Step 5: Examine the results
The following are representative results from an actual CPU execution with PyTorch 2.10.0.
Selected steps, rounded
|
Mode |
Step |
Loss |
Pre-clip gradient norm |
Outcome |
|---|---|---|---|---|
|
FP16 + GradScaler |
5 |
0.0416 |
0.408 |
Updated |
|
FP16 + GradScaler |
6 |
inf |
inf |
Skipped |
|
FP16 + GradScaler |
7 |
0.0383 |
0.391 |
Updated |
|
BF16, no clipping |
5 |
0.0417 |
0.408 |
Updated |
|
BF16, no clipping |
6 |
3.79 × 10⁸ |
3.89 × 10⁹ |
Updated |
|
BF16, no clipping |
7 |
6.02 × 10¹⁵ |
1.55 × 10⁸ |
Updated |
|
BF16 + clipping |
5 |
0.0417 |
0.408 |
Updated |
|
BF16 + clipping |
6 |
3.79 × 10⁸ |
3.89 × 10⁹ |
Clipped |
|
BF16 + clipping |
7 |
0.0309 |
0.352 |
Updated |
GPU values may differ slightly because of kernel implementations and rounding.
Case A: FP16 produces infinity
At step 6, the injected input is outside FP16’s finite range.
The forward pass overflows, producing an infinite loss and non-finite gradients.
GradScaler detects the invalid gradients and skips the optimizer update.
The model parameters remain intact, and step 7 proceeds normally.
This demonstrates an important distinction:
The scaler prevented an invalid parameter update, but it did not prevent the forward overflow.
The problematic batch did not contribute a useful learning signal.
Repeated events of this kind can waste training steps or make the job effectively unable to learn from certain examples.
Case B: BF16 prevents overflow but not instability
With BF16, the outlier produces a large but finite prediction.
The loss is approximately 379 million.
The gradient norm approaches 3.9 billion.
Without clipping, SGD applies an enormous update. The model’s weight jumps from approximately 0.196 to around −77.85 million.
On the following ordinary batch, the loss is approximately 6.02 × 10¹⁵.
The original outlier is gone, but its effect on the parameters persists.
BF16 fixed the representational problem without fixing the optimization problem.
Case C: BF16 plus clipping contains the damage
In the final run, BF16 again produces a finite loss and gradient.
This time, we clip the gradient norm to 1.0 before calling the optimizer.
The gradient is rescaled by roughly:
1 / 3.89 billion ≈ 2.57 × 10⁻¹⁰
With a learning rate of 0.02, the scalar parameter update caused by the outlier is limited to approximately 0.02.
The model’s weight changes from approximately 0.196 to 0.176 rather than jumping to a magnitude of tens of millions.
At the next ordinary step, loss returns to approximately 0.031.
Crucially, the loss at step 6 remains huge. Gradient clipping does not retroactively reduce the loss that was already computed.
It prevents that loss from causing an uncontrolled parameter update.
This is the central result of the experiment:
BF16 keeps the computation representable; gradient clipping keeps the resulting optimizer update manageable.
6. Build a reliable training loop with gradient accumulation
The single-parameter experiment isolates numerical effects, but production training often uses gradient accumulation.
Gradient accumulation combines contributions from several micro-batches before making one optimizer update.
This introduces another common stability bug: clipping separately after each micro-batch.
Doing so changes the optimization objective because each contribution is rescaled independently.
Instead, accumulate the gradients for the complete effective batch, then clip once.
A runnable BF16 accumulation example
The following script generates neutral synthetic regression data and trains a small model with BF16 autocasting, four-way accumulation, global norm clipping, and finite-gradient checks.
import torch import torch.nn.functional as F torch.manual_seed(17) device = ( "cuda" if torch.cuda.is_available() and torch.cuda.is_bf16_supported() else "cpu" ) generator = torch.Generator( device="cpu" ).manual_seed(17) features = torch.randn( 64, 8, generator=generator ).to(device) targets = ( features.sum( dim=1, keepdim=True, ) * 0.2 ) model = torch.nn.Linear( 8, 1 ).to(device) optimizer = torch.optim.AdamW( model.parameters(), lr=0.01, ) accumulation_steps = 4 micro_batch_size = 4 num_micro_batches = 16 for start in range( 0, num_micro_batches, accumulation_steps, ): optimizer.zero_grad( set_to_none=True ) for offset in range(accumulation_steps): batch_index = start + offset left = batch_index * micro_batch_size right = left + micro_batch_size x = features[left:right] y = targets[left:right] with torch.autocast( device_type=device, dtype=torch.bfloat16, ): predictions = model(x) loss = F.mse_loss( predictions.float(), y, ) normalized_loss = ( loss / accumulation_steps ) normalized_loss.backward() pre_clip_norm = ( torch.nn.utils.clip_grad_norm_( model.parameters(), max_norm=1.0, error_if_nonfinite=True, ) ) optimizer.step() print( f"Update {start // accumulation_steps}: " f"gradient_norm={pre_clip_norm.item():.4f}" )
Notice four design choices.
First, the loss is divided by the accumulation count.
Because each micro-batch has the same size and uses mean reduction, dividing by four makes the accumulated gradient equivalent, up to numerical differences, to the average gradient over the combined batch.
With unequal micro-batch sizes, use sample-weighted normalization instead.
Second, clipping happens once.
The four backward passes contribute to the same parameter gradients before the clipping operation.
Third, non-finite gradient norms raise an exception.
The example uses error_if_nonfinite=True to fail fast. This is appropriate for the BF16 configuration shown here, where we want invalid gradients to trigger investigation rather than silently flowing into the optimizer.
Fourth, BF16 autocasting does not require FP16-style loss scaling.
BF16’s large exponent range reduces the underflow concerns that motivate ordinary FP16 gradient scaling.
This does not mean BF16 is mathematically exact or immune to non-finite computations. Large reductions, badly conditioned operations, and invalid mathematical inputs can still fail.
How the loop changes with FP16
When using FP16 and GradScaler, the effective batch must be accumulated at a consistent loss scale.
The correct sequence is:
-
Scale each normalized micro-batch loss before backpropagation.
-
Accumulate all micro-batch gradients.
-
Unscale once, after accumulation finishes.
-
Clip the unscaled gradients.
-
Call the scaler’s optimizer-step operation.
-
Update the scale once for the completed effective batch.
Do not call unscale_ after every micro-batch, and do not update the scale partway through accumulation.
If gradients are non-finite in the FP16 path, GradScaler’s skip-and-backoff behavior should be preserved. An unconditional fail-fast clipping check can prevent that automatic recovery policy, so choose explicitly between diagnostic failure and managed skipping.
7. How to diagnose an actual instability
A useful monitoring system records more than training loss.
The most important principle is to gather enough information to identify the first abnormal event, not just the last visible symptom.
Minimum telemetry for stable training
|
Metric |
What it reveals |
Warning sign |
|---|---|---|
|
Raw training loss |
Immediate optimization behavior |
Abrupt unexplained jump |
|
Validation loss |
Generalization and recovery |
Persistent degradation |
|
Pre-clipping gradient norm |
Gradient magnitude |
Large deviation from baseline |
|
Clipping frequency |
How often updates are constrained |
Sustained increase |
|
Learning rate |
Update aggressiveness |
Spikes around schedule changes |
|
AMP loss scale |
FP16 numerical behavior |
Repeated downward adjustments |
|
Non-finite gradient count |
Numerical failure |
Any unexpected occurrence |
|
Activation maxima |
Intermediate value growth |
Rapidly expanding range |
|
Parameter norm |
Weight magnitude |
Unexplained drift |
|
Batch metadata |
Data-related anomalies |
Outliers or distribution changes |
The word baseline matters. A gradient norm of 10 is not automatically dangerous, and a norm of 0.1 is not automatically safe.
Compare measurements against the model’s recent stable behavior.
Step 1: Locate the first abnormal signal
Suppose the loss spikes at step 20,001.
Look at the gradient norm from step 20,000.
If the gradient norm rose dramatically before the loss spike, the previous update may be responsible.
If the loss was already infinite during the forward pass at step 20,001, investigate activations and numerical operations first.
Be careful with step indexing: many logging systems record metrics before the optimizer update, while others report them after it.
A one-step mismatch can lead to an incorrect diagnosis.
Step 2: Inspect the problematic batch
Check whether the batch contains:
-
Unexpected input magnitudes.
-
Missing normalization.
-
Non-finite input values.
-
Incorrect targets.
-
A distribution shift.
-
Extreme sequence lengths or other unusual structural properties.
For tabular regression, a common example is inconsistent feature units.
If a model normally receives an energy-usage feature in kilowatt-hours but one pipeline starts supplying watt-hours, numerical magnitudes increase by a factor of 1,000.
The model may see a valid number that is nevertheless completely outside its training distribution.
Clipping may protect parameters, but correcting the data pipeline is the real solution.
Step 3: Compare precisions
Run the same batch through a controlled FP32 baseline.
Then compare it with BF16 and FP16 autocasting.
Look for:
-
Outputs that are finite in FP32 but infinite in FP16.
-
Operations producing NaNs in multiple precisions.
-
Exceptionally large activation values.
-
Unexpected differences caused by rounding or reduction behavior.
If FP16 fails while FP32 and BF16 produce finite outputs, numerical range becomes a strong suspect.
If every precision fails, investigate the mathematical operation, data, model state, and optimization history.
Step 4: Inspect gradient distribution by layer
A single global norm is useful but incomplete.
A transformer might have reasonable average gradient magnitudes while one projection matrix produces an exceptional spike.
Collect per-layer information such as:
-
Gradient L2 norm.
-
Maximum absolute gradient.
-
Fraction of non-finite values.
-
Parameter norm.
-
Update-to-parameter norm ratio.
These measurements can reveal whether instability is global or concentrated in a specific component.
That distinction becomes especially important for attention projections, embeddings, and specialized layers.
8. Cherry on the cake: a real 70-billion-parameter loss-spike incident
The most revealing large-model training stories are not always about a better optimizer. Sometimes they are about building systems that recognize when training is becoming unsafe.
A concrete example appears in the December 2025 K2-V2 technical report, published by the Institute of Foundation Models at Mohamed bin Zayed University of Artificial Intelligence.
The report documents the development of the 70-billion-parameter K2 model, including practical training stability problems and an automated recovery system.
In Figure 9, the team shows a severe training-loss spike around step 464,000.
The accompanying training alert identifies consecutive problematic iterations, beginning at 464,821, with reported loss values climbing to approximately 8.71.
The monitoring system automatically rolled the job back to checkpoint step 464,500 and restarted training.
This was not merely a retrospective plot showing that something went wrong. The team had operational machinery that detected the event and initiated recovery.
What made the recovery system interesting?
Rather than restarting whenever a single loss value became unusually high, the researchers differentiated transient fluctuations from sustained divergences.
Their detector examined a sliding window of recent losses and used two conditions:
-
Persistence: the minimum loss within the window had to exceed a sustained-loss threshold.
-
Severity: the maximum loss within the window had to exceed a separate critical threshold.
Both conditions had to hold before the event qualified for a particular level of intervention.
The system also used different sensitivities for early warnings and automatic restarts.
That design captures a production lesson: treating every isolated loss spike as catastrophic can waste compute, but ignoring sustained instability risks wasting considerably more.
The team also analyzed spikes using robust statistics derived from rolling medians and median absolute deviations. Their analysis found that most spikes were concentrated in the earlier portions of pretraining.
The surprising detail
The report describes BF16 precision being enabled in the training configuration.
In other words, BF16 did not eliminate the need for stability monitoring and checkpoint recovery.
This is particularly relevant to our small experiment.
We demonstrated that BF16 fixes a representational overflow, while clipping prevents an excessive finite update.
But real large-model training has additional failure mechanisms involving optimizer dynamics, model architecture, schedules, and data interactions.
The K2-V2 incident illustrates why numerical precision must be part of a broader resilience strategy, not treated as a complete solution.
The original report is available as K2-V2: A 360-Open, Reasoning-Enhanced LLM, particularly sections 3.5.2 and 3.5.3.
9. Why clipping and BF16 are still not enough
Even a correctly implemented BF16 training loop with gradient clipping can experience severe instability.
There are several reasons.
Clipping does not correct the gradient direction
Global norm clipping reduces magnitude while preserving direction.
A large gradient pointing in a harmful direction remains harmful, although its effect is reduced.
If the gradient is dominated by a mislabeled example or a systematically incorrect data transformation, repeatedly clipping it does not solve the underlying problem.
Clipping does not undo earlier optimizer damage
Suppose a large update was applied at step 10,000 and shifted the optimizer’s internal state.
Adding clipping at step 10,001 does not restore the previous parameter values or momentum estimates.
When recovery is required, loading a known-good checkpoint can be safer than trying to continue from a corrupted state.
For AdamW, checkpoint recovery should generally restore both model parameters and optimizer state.
For FP16 training, preserving GradScaler state is also important for consistent continuation.
Clipping does not guarantee finite activations
You can constrain parameter gradients perfectly and still calculate invalid forward values.
For example, an unstable exponential operation can overflow before backpropagation.
The correct solution might be:
-
Rewriting the operation in a numerically stable form.
-
Computing a sensitive sub-operation in FP32.
-
Normalizing intermediate values.
-
Constraining logits or activation magnitudes.
-
Fixing an invalid input distribution.
One particularly common example is manually computing softmax using exponentials without appropriate numerical stabilization.
Use tested framework implementations of softmax and log-softmax rather than reconstructing unstable variants.
BF16 does not eliminate precision error
BF16’s larger dynamic range comes at the cost of fewer fraction bits.
That can matter for:
-
Very small parameter updates.
-
Subtraction of nearby quantities.
-
Accumulation of rounding errors.
-
Sensitive normalization calculations.
-
Algorithms with poor numerical conditioning.
This is one reason autocasting is preferable to indiscriminately converting every model parameter and operation to BF16.
With standard mixed-precision training, parameters can remain FP32 while suitable computation kernels use lower precision.
10. Advanced stability strategies
Once basic clipping, precision configuration, and logging are correct, several additional techniques become useful.
Learning-rate warmup
Aggressive learning rates early in training can produce unstable updates before optimizer estimates and model representations settle.
Warmup gradually increases the learning rate from a small initial value to the intended peak.
It can be particularly useful for large transformers and models with sensitive initialization.
However, warmup is not a substitute for clipping.
A training run can still encounter an outlier late in training after the learning rate has stabilized.
Adaptive gradient clipping
A fixed global clipping threshold assumes that a single norm limit is appropriate throughout training.
That assumption can fail when gradient magnitudes differ substantially across layers or evolve over time.
Adaptive clipping methods estimate thresholds using local parameter or historical gradient statistics.
For instance, recent research on AdaGC investigates per-parameter adaptive clipping thresholds based on exponential moving averages of gradient norms.
The motivation is that one global threshold may over-constrain some components while providing inadequate protection to others.
Adaptive methods are useful research directions, but they introduce extra state, hyperparameters, and operational complexity.
They should be evaluated against a correctly tuned fixed-clipping baseline rather than adopted automatically.
Stable optimizer configurations
Optimizer selection and hyperparameters influence the effect of gradient spikes.
Relevant considerations include:
-
Momentum coefficients.
-
Second-moment decay rates.
-
Numerical epsilon values.
-
Weight decay.
-
Learning-rate magnitude.
-
Parameter groups with different sensitivities.
In Adam-style optimizers, estimates of gradient history may lag sudden changes in gradient statistics.
This can create situations where the effective update does not respond well to an abrupt change in the optimization landscape.
A stability investigation should therefore examine optimizer state and update magnitudes, not only raw gradients.
Selective FP32 operations
Some operations may require more precision or range than the surrounding computation.
In those cases, selectively disabling autocasting for a sensitive operation can be more efficient than abandoning mixed precision throughout the entire model.
This is especially relevant to custom losses, unusual reductions, and custom kernels.
However, the operands must also be converted appropriately; disabling autocasting does not magically restore information already lost through an earlier overflow.
Reproducible checkpoint recovery
For long-running training jobs, the recovery system is part of the training algorithm in practice.
A checkpoint should contain the information required to resume consistently, including:
-
Model parameters and relevant buffers.
-
Optimizer state.
-
Learning-rate scheduler state.
-
Current global step and consumed-data position.
-
Random-number-generator states where required.
-
Gradient-scaler state for FP16 runs.
-
Relevant distributed training and sampler state.
Checkpoint integrity and recovery procedures should be tested before an expensive run begins.
A checkpoint that cannot be restored under realistic failure conditions provides little protection.
11. Common mistakes and their consequences
|
Mistake |
Consequence |
Better approach |
|---|---|---|
|
Clipping scaled FP16 gradients |
Excessively small effective updates |
Unscale before clipping |
|
Using only GradScaler to handle forward overflow |
Bad forward computations remain invalid |
Inspect operations and try BF16 or selective FP32 |
|
Clipping after the optimizer step |
Cannot protect the update just applied |
Clip before the step |
|
Clipping each accumulation micro-batch independently |
Changes the effective accumulated gradient |
Clip once per effective batch |
|
Logging only post-clipping norms |
Conceals large raw gradients |
Record pre-clipping norms |
|
Treating every spike as divergence |
Unnecessary restarts |
Check persistence and recovery |
|
Ignoring isolated non-finite events |
Can hide a recurring numerical issue |
Count and investigate occurrences |
|
Converting all parameters directly to FP16 |
Risks loss of precision and range |
Prefer autocasting with FP32 parameters |
|
Assuming BF16 guarantees stability |
Misses finite-gradient and optimizer problems |
Combine precision choice with monitoring |
|
Choosing a clipping threshold without measurement |
Can suppress learning or fail to contain spikes |
Calibrate from stable-run telemetry |
12. A practical incident-response playbook
When a training run first exhibits a severe anomaly, work through the following sequence.
Phase A: Confirm the symptom
Determine whether the event is:
-
A finite, isolated loss increase.
-
A sustained increase in loss.
-
An infinite or NaN loss.
-
A large gradient norm preceding a parameter update.
Check whether validation behavior has also changed.
Avoid immediately lowering the learning rate without first identifying the failure mode.
Phase B: Identify the numerical failure boundary
Inspect the earliest operation or tensor that became abnormal.
If inputs are finite but a particular FP16 operation produces infinity, test BF16 and FP32.
If the forward pass is finite but backpropagation produces non-finite gradients, inspect backward operations, loss scaling, and gradient magnitudes.
If gradients are finite but the next step becomes unstable, inspect the optimizer update and clipping behavior.
Phase C: Choose the intervention
Match the intervention to the evidence.
-
FP16 overflow: consider BF16 or selective FP32.
-
Finite gradient explosion: introduce or tune gradient clipping.
-
Repeated anomalous batches: investigate input preprocessing and data composition.
-
Aggressive parameter updates: review the learning rate, optimizer configuration, and schedule.
-
Persistent divergence: consider restoring a known-good checkpoint.
-
Intermittent instability: add monitoring before making unnecessarily disruptive changes.
Phase D: Validate the fix
A successful fix should do more than eliminate NaNs.
Compare the modified run against a stable baseline using:
-
Training and validation trajectories.
-
Raw gradient norms.
-
Clipping frequency.
-
Effective optimizer updates.
-
Throughput and memory usage.
-
Recovery behavior after outlier batches.
A configuration that never diverges because almost every update is suppressed is not necessarily a successful training configuration.
The goal is stable learning, not merely finite numbers.
13. Final checklist for production training
Before launching a substantial training job, verify the following.
- The precision format is supported by the chosen hardware.
- Mixed precision uses autocasting rather than unnecessary full-model FP16 conversion.
- FP16 GradScaler is configured correctly when needed.
- Gradients are unscaled before clipping in FP16.
- Clipping happens after gradient accumulation and before the optimizer step.
- Global clipping is appropriate for the distributed or sharded training setup.
- Raw gradient norms are logged before clipping.
- Non-finite losses and gradients are monitored.
- The model has been tested with extreme but valid input magnitudes.
- Known-good checkpoints include the state necessary for recovery.
- The restart procedure has been tested.
- Validation metrics are monitored alongside training metrics.
- Stability changes are evaluated against convergence quality, not just the absence of spikes.
14. What to remember
There are three fundamental distinctions that make training instability much easier to reason about.
Numerical range is not optimization stability. BF16 can represent values that FP16 cannot, but finite computations can still create disastrously large parameter updates.
Gradient clipping is not numerical error correction. It can prevent large gradients from producing excessive updates, but it cannot restore an invalid forward computation or repair corrupted optimizer state.
A loss spike is not automatically a training failure. Short-lived outliers, persistent divergence, and numerical overflows require different diagnoses and operational responses.
The runnable experiment demonstrates all three points. FP16 produced an infinite loss, BF16 allowed the computation to remain finite but exposed an enormous update, and BF16 combined with clipping kept subsequent training stable.
The larger engineering lesson is that robust training requires several layers of protection: numerically suitable operations, controlled updates, observability, and reliable recovery.
Your next step
Run stability_demo.py and experiment with three parameters: the outlier magnitude, learning rate, and maximum gradient norm.
Start by lowering the outlier below FP16’s representable limit. Then increase it until overflow appears. Compare the resulting gradients under FP16 and BF16, and experiment with clipping thresholds of 0.1, 1.0, and 10.0.
Finally, add pre-clipping gradient-norm logging and a non-finite detector to one of your own training pipelines.
The objective is not to hide every spike. It is to understand why each spike happens, prevent avoidable damage, and make recovery predictable when something genuinely goes wrong.