Training Stability › Debugging Training Instability
Language AI HandbookAll chapters Training Stability › Debugging Training InstabilityOne section · Sections in this chapter7 / 15 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
When instability occurs, the debugging process requires methodical isolation. A spike in training loss could stem from many causes, and the symptoms often look similar regardless of root cause.
A Systematic Debugging Workflow
Effective debugging follows a structured approach that eliminates hypotheses one by one.
Step 1: Isolate the timeline. Identify exactly when the instability began by examining gradient norm logs alongside loss. If the gradient norm spikes before the loss spike, the cause is an optimizer issue. If loss spikes without a gradient norm precursor, the cause may be a data issue or numerical overflow.
Step 2: Examine the triggering batch. If your training loop logs data indices, retrieve the batch that preceded the spike and inspect it. In language modeling, look for sequences with unusual token distributions, very long lengths, or repeated patterns. These outlier batches are common triggers.
Step 3: Check for numerical issues. NaN or Inf values in gradients or activations always indicate a numerical problem. Adding NaN checks after loss computation and after backpropagation immediately narrows the search:
def check_for_nan(model, step, loss):
"""Check for NaN or Inf values in loss, gradients, and parameters."""
issues = []
if torch.isnan(loss) or torch.isinf(loss):
issues.append(f"Step {step}: NaN/Inf in loss (value={loss.item()})")
for name, param in model.named_parameters():
if param.grad is not None:
if torch.any(torch.isnan(param.grad)):
issues.append(f"Step {step}: NaN gradient in {name}")
if torch.any(torch.isinf(param.grad)):
issues.append(f"Step {step}: Inf gradient in {name}")
if torch.any(torch.isnan(param.data)):
issues.append(f"Step {step}: NaN weight in {name}")
return issuesStep 4: Bisect the architecture. If numerical issues are found in specific layers, reduce the model to just those layers and reproduce the problem with a minimal example. This makes the root cause much easier to identify and fix.
Step 5: Check initialization. Poorly initialized models are fragile in early training. Re-running the first 1000 steps with a different random seed can determine whether instability is systematic (same architecture, same data, same behavior) or stochastic (different seed resolves it). Systematic instability requires architectural or hyperparameter fixes; stochastic instability may simply require better luck or a more reliable initialization scheme.
Layer-Wise Gradient Analysis
When the global gradient norm shows a spike, it does not tell you which layers are responsible. Layer-wise gradient analysis decomposes the gradient norm by component, revealing which parts of the network are generating the instability. This is valuable for architectural debugging.
In a transformer, the most common sources of gradient spikes are:
- The embedding layer, which can receive very large gradients when the vocabulary is large and some tokens are rare
- The final projection (language model head), which connects the full model width to the vocabulary size and amplifies gradients proportionally
- Attention output projections, especially in early layers where gradients from all subsequent layers accumulate
def layer_gradient_analysis(model):
"""Compute per-layer gradient norms for architectural debugging."""
layer_norms = {}
for name, param in model.named_parameters():
if param.grad is not None:
layer_norms[name] = param.grad.norm(2).item()
return layer_norms
# Run a single forward-backward pass and analyze layer gradients
torch.manual_seed(42)
debug_model = SimpleTransformerBlock(d_model=128, n_heads=4)
debug_opt = optim.AdamW(debug_model.parameters(), lr=1e-3)
x_debug = torch.randn(16, 32, 128)
target_debug = torch.randn(16, 32, 128)
debug_opt.zero_grad()
out_debug = debug_model(x_debug)
loss_debug = criterion(out_debug, target_debug)
loss_debug.backward()
layer_norms = layer_gradient_analysis(debug_model)Loss: 2.0107 Per-layer gradient norms: Layer Grad Norm -------------------------------------------------------------- norm2.weight 0.177542 ffn.2.weight 0.058973 ffn.0.weight 0.030297 norm2.bias 0.016122 attention.in_proj_weight 0.012976 attention.out_proj.weight 0.012630 attention.out_proj.bias 0.007798 ffn.2.bias 0.007755 norm1.bias 0.007732 norm1.weight 0.007320 attention.in_proj_bias 0.004454 ffn.0.bias 0.002660
This output reveals the relative contribution of each layer to the total gradient signal. If one layer consistently shows gradient norms 10-100x larger than others, that is the component deserving architectural attention.
Monitoring Infrastructure
Effective debugging requires proper monitoring infrastructure set up before problems occur, not after. The minimum viable monitoring for a large training run includes:
- Loss curve: Training loss (and validation loss if available) at every step or every N steps
- Gradient norm: Computed and logged before each optimizer step
- Parameter norms: The Frobenius norm of each weight matrix layer, logged every 100-200 steps. Rapidly growing parameter norms indicate the weight decay coefficient is too small.
- Learning rate: The current learning rate at each step (especially during warmup and decay phases)
- Gradient-to-parameter ratio: The ratio characterizes how large updates are relative to parameter magnitudes. Values above 0.1 are often associated with instability.
def compute_monitoring_metrics(model, loss, step):
"""Compute comprehensive monitoring metrics for one training step."""
metrics = {"step": step, "loss": loss.item()}
# Gradient norm
total_grad_norm_sq = 0.0
for param in model.parameters():
if param.grad is not None:
total_grad_norm_sq += param.grad.data.norm(2).item() ** 2
metrics["grad_norm"] = total_grad_norm_sq**0.5
# Parameter norm (Frobenius across all params)
total_param_norm_sq = 0.0
for param in model.parameters():
total_param_norm_sq += param.data.norm(2).item() ** 2
metrics["param_norm"] = total_param_norm_sq**0.5
# Gradient-to-parameter ratio
if metrics["param_norm"] > 0:
metrics["grad_param_ratio"] = (
metrics["grad_norm"] / metrics["param_norm"]
)
else:
metrics["grad_param_ratio"] = float("nan")
return metrics
# Run a monitored training loop
torch.manual_seed(42)
monitored_model = SimpleTransformerBlock(d_model=128, n_heads=4)
monitored_opt = optim.AdamW(
monitored_model.parameters(), lr=1e-3, weight_decay=0.01
)
monitoring_log = []
for step in range(60):
x = torch.randn(16, 32, 128)
target = torch.randn(16, 32, 128)
monitored_opt.zero_grad()
out = monitored_model(x)
loss_val = criterion(out, target)
loss_val.backward()
metrics = compute_monitoring_metrics(monitored_model, loss_val, step)
torch.nn.utils.clip_grad_norm_(monitored_model.parameters(), max_norm=1.0)
monitored_opt.step()
monitoring_log.append(metrics) Step Loss Grad Norm Param Norm g/p Ratio
----------------------------------------------------
0 2.0107 0.1917 26.5670 0.00722
10 1.9674 0.1889 26.5752 0.00711
20 1.9495 0.1889 26.5815 0.00711
30 1.9310 0.1885 26.6000 0.00709
40 1.8933 0.1896 26.6370 0.00712
50 1.8609 0.1924 26.7158 0.00720
