3.8 Troubleshooting ML Model Training Failures

3.8 Troubleshooting ML Model Training Failures

Blueprint reference: Section 3.2, "Troubleshooting ML model training failures."

Exam questions in this area give you a symptom and expect a diagnosis. The skill is classifying the symptom correctly before proposing a fix, because the wrong class leads to a plausible answer that does nothing.

Classify the Failure First

Out-of-Memory on the Accelerator

The most common hard failure. Device memory holds parameters, gradients, optimizer state, and — usually dominant — activations, which scale with batch size.

Remedies in order of preference:

1. Reduce per-device batch size. Immediate and always available.
2. Gradient accumulation. Keep the effective batch size by accumulating gradients over several micro-batches before the optimizer step. Preserves convergence behaviour at a smaller memory footprint.
3. Mixed precision. Store activations in bf16/fp16 while keeping a master copy of weights in fp32. Roughly halves activation memory and usually speeds training on modern accelerators.
4. Gradient checkpointing. Recompute activations during the backward pass instead of storing them — trades compute for memory, often 30–40% slower but a large memory saving.
5. Model parallelism / sharding. When the model itself does not fit, split it across devices (see the distributed training section).
6. A larger-memory accelerator.

The classic wrong answer: "add more workers." Data-parallel workers each hold a full copy of the model, so adding workers does not reduce per-device memory at all.

NaN and Divergence

When the loss becomes NaN or Inf:

• Lower the learning rate and add linear warmup. A large initial learning rate on a fine-tuning run is the single most common cause.
• Clip gradients by global norm. Standard practice for recurrent and transformer models.
• Check for log(0) and division by zero in a custom loss. Add a small epsilon.
• Normalize inputs. Unscaled features with wildly different magnitudes destabilize the first steps.
• Check the data for NaNs. A single NaN in a feature column propagates through every gradient.
• In mixed precision, verify loss scaling is enabled; fp16 gradients underflow without it.

Underfitting Versus Overfitting

A validation loss that is lower than training loss usually means dropout or augmentation is active during training and disabled during evaluation, which is expected — not a bug.

Failures That Only Appear at Scale

• Quota. Accelerator quota is per region and per type. A job that runs with 2 GPUs and fails with 16 is usually quota, and the error will say ResourceExhausted before any training step.
• Permissions. The training job runs as a service account, not as the submitting user. A notebook that reads a bucket successfully proves nothing about whether the job's service account can.
• Distributed initialization. Every worker must agree on the cluster specification; a mismatched worker count, an unreachable chief, or a blocked port hangs the job at initialization with no useful loss output.
• Stragglers. One slow worker throttles a synchronous all-reduce step. Uneven shard sizes are a frequent cause.

Checkpointing: The Non-Negotiable

Long training runs fail. Machines preempt, quotas throttle, transient errors happen. Checkpoint to Cloud Storage at regular step or time intervals, saving model weights, optimizer state, and the step counter, then resume from the latest checkpoint on restart.

This is mandatory rather than optional when using Spot or preemptible resources, whose entire value proposition — substantially lower cost — is only realizable if reclamation costs you minutes rather than the whole run.

# Resume-safe: the job restarts into the same state it left
ckpt_dir = os.environ["AIP_CHECKPOINT_DIR"]   # provided by the training service
if latest := find_latest_checkpoint(ckpt_dir):
    model, optimizer, start_step = load_checkpoint(latest)

Profiling Before Guessing

When a job runs but is slow rather than failing, stop diagnosing and start profiling. The tooling — Cloud Logging, Cloud Monitoring, the TensorBoard profiler, and Cloud Profiler — and the read-the-trace workflow are covered in Section 6.4, Operational Troubleshooting, Performance Profiling and Cost Optimization. The rule that belongs in a failure triage is simply that slow and failed are different classes: a job that completes with poor utilization is a performance investigation, while a job that terminates is one of the failure classes tabulated above, and applying the wrong playbook wastes the cycle.

Exam Traps

• Adding workers to fix out-of-memory. Data parallelism replicates the model.
• Checking your own IAM instead of the job's service account.
• Increasing the learning rate when the loss is already diverging.
• Spot resources without checkpointing.
• Guessing at slowness instead of profiling.