Training Diagnostics — infra.bea.ai

Training a large model for weeks on thousands of GPUs exposes failure modes that don’t appear in small-scale experiments. This chapter covers the symptoms, causes, and solutions for the most common training pathologies.

1. Loss Spikes

Symptoms

Loss suddenly increases 2-100× for one or more steps, then may or may not recover.

Common Causes

Cause	Pattern	Fix
Bad data batch	Single spike, recovers	Filter training data more aggressively
Learning rate too high	Repeated spikes, may diverge	Reduce LR, use warmup
Gradient explosion	Loss → NaN/Inf	Gradient clipping (`max_norm=1.0`)
Numerical instability	Spikes after many steps	Check LayerNorm, softmax in FP32
Bit flip (silent corruption)	Random spike, no data correlation	ECC monitoring, redundant computation

Monitoring

# Log gradient norm every step
grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
wandb.log({"grad_norm": grad_norm, "loss": loss.item()})

# Alert on anomalies
if loss.item() > 10 * running_avg_loss:
    alert(f"Loss spike at step {step}: {loss.item():.4f} (avg: {running_avg_loss:.4f})")

Recovery from Loss Spikes

Option 1: Skip the bad batch, continue training
Option 2: Rollback to checkpoint before spike, skip the data range
Option 3: Reduce learning rate and resume

Meta’s Llama-3 training encountered “a handful of loss spikes” during 405B pre-training and handled them by rolling back to a checkpoint and skipping the problematic data.

2. Gradient Anomalies

Vanishing Gradients

Symptom	Cause	Fix
Gradient norm → 0	Too many layers without residual connections	Pre-norm architecture, residual scaling
Some params never update	Dead neurons (ReLU at 0)	Use SiLU/GELU, check initialization
Loss plateaus	LR too low for current loss landscape	LR warmup, cosine schedule

Exploding Gradients

Symptom	Cause	Fix
Gradient norm → ∞	Unstable attention (large logits)	QK-norm, attention scaling
NaN in specific layers	Numerical overflow in FP16	Use BF16, check loss scaling
Step time increases	Inf/NaN handling is slow	Fix root cause, don’t just clip

# QK-norm: normalize Q and K before attention
# Prevents attention logits from growing unboundedly
q = F.normalize(q, dim=-1) * math.sqrt(head_dim)
k = F.normalize(k, dim=-1)
attn = q @ k.transpose(-2, -1)  # bounded by √(head_dim)

3. Silent Data Corruption

The most dangerous failure — incorrect computation without any error signal.

Causes

Cause	Detection	Frequency
GPU SRAM bit flip	Uncorrectable ECC error	Rare but real
HBM bit flip	ECC correctable (masked) or uncorrectable	Rare
Network bit flip	CRC catches most, but not all	Very rare
Software bug (CUDA, NCCL)	Regression tests	Varies

Detection Strategies

# Strategy 1: Periodic gradient checksum
# Every N steps, compute gradient on 2 different GPUs and compare
if step % 1000 == 0:
    grad_checksum = compute_gradient_checksum(model)
    all_checksums = dist.all_gather(grad_checksum)
    if not all_equal(all_checksums):
        alert("Gradient mismatch — possible silent corruption!")

# Strategy 2: Monitor per-layer gradient norms
for name, param in model.named_parameters():
    if param.grad is not None:
        grad_norm = param.grad.norm().item()
        if math.isnan(grad_norm) or grad_norm > 1000:
            alert(f"Anomalous gradient in {name}: {grad_norm}")

Meta’s Approach (from Llama-3 paper)

“We developed tools for online detection of anomalies including: NaN checks in every microbatch, comparison of per-layer gradient norms across data-parallel replicas, and automated rollback when divergence is detected.”

4. Thermal Throttling

Symptoms

Step time gradually increases over minutes/hours
Some GPUs report lower clock speeds than others
nvidia-smi shows GPU temperature >80°C

Diagnosis

# Check GPU temperatures
nvidia-smi --query-gpu=temperature.gpu --format=csv -l 5

# Check power draw
nvidia-smi --query-gpu=power.draw --format=csv -l 5

# Check clock speed throttling
nvidia-smi --query-gpu=clocks_event_reasons.hw_thermal_slowdown --format=csv

Mitigation

Action	Effect
Improve data center cooling	Prevents throttling entirely
Reduce power limit	Lower temps at cost of compute
Stagger workloads	Avoid thermal hotspots
Replace thermal paste	Fix degraded cooling on specific GPUs

5. Learning Rate Scheduling Issues

LR scheduling directly affects training stability and efficiency:

Common Schedules

Schedule	Formula	Use case
Linear warmup	$\text{lr} = \text{lr}_{max} \times t / t_{warmup}$	First 1-5% of training
Cosine decay	$\text{lr} = \text{lr}_{min} + 0.5(\text{lr}_{max} - \text{lr}_{min})(1 + \cos(\pi t / T))$	Standard for pre-training
WSD (Warmup-Stable-Decay)	Warmup → constant → linear decay	Used by MiniCPM, some prefer over cosine

Symptom	Likely cause	Fix
Loss doesn’t decrease from start	LR too low or warmup too long	Increase LR, shorten warmup
Loss diverges after warmup	LR too high	Reduce max LR
Loss plateaus mid-training	LR decayed too early	Extend constant phase (WSD)
Loss spikes at specific steps	LR schedule discontinuity	Smooth transitions
Different results at different scales	LR not adjusted for batch size	Scale LR with $\sqrt{\text{batch}}$ or linearly

Batch Size and LR Coupling

When scaling to more GPUs (larger effective batch), the learning rate must be adjusted:

\text{lr}_{\text{scaled}} = \text{lr}_{\text{base}} \times \sqrt{\frac{B_{\text{new}}}{B_{\text{base}}}}

Linear scaling ( $\text{lr} \propto B$ ) also works for smaller scale-ups but becomes unstable at very large batches. LAMB and LARS optimizers handle this automatically.

6. Network Congestion

Symptoms

AllReduce time increases without model/config changes
Some steps take 2-10× longer than others
nccl-tests show lower bandwidth than expected

Diagnosis

# Run nccl-tests during off-hours (baseline) vs during training (contention)
mpirun -np 16 ./all_reduce_perf -b 1M -e 1G -f 2 -g 1

# Check for InfiniBand errors
perfquery -x  # port counters
ibdiagnet      # comprehensive network diagnostics

Common Causes

Cause	Fix
Oversubscribed fabric	Add spine switches or reduce DP degree
Faulty NIC/cable	Replace hardware (check `ibdiagnet` errors)
PFC storms (RoCE)	Tune ECN/DCQCN parameters
Competing traffic	Isolate training network from storage/management

For communication performance analysis, see Communication Overhead Estimation.

Self-Check

Diagnose a loss spike: identify likely cause from the pattern
Implement gradient norm monitoring with anomaly detection
Explain silent data corruption and two detection strategies
Diagnose thermal throttling from nvidia-smi output
Identify network congestion from nccl-tests results

References

The Llama 3 Herd of Models (Grattafiori et al., 2024) — training diagnostics
PaLM: Scaling Language Modeling (Chowdhery et al., 2022) — loss spike handling