10.1 Troubleshooting ML Training Failures
Key Takeaways
- An interactive shell on a running serverless training job lets you inspect the container, run py-spy or nvidia-smi, and check permissions with gcloud auth list.
- Interactive shell access ends as soon as the job or trial finishes, and files created in the container don't persist.
- GPU out-of-memory errors are usually fixed by smaller batches, gradient accumulation, mixed precision, gradient checkpointing, a larger-memory accelerator, or model-parallel sharding.
- A NaN or exploding loss usually points to a learning rate that is too high, bad or unnormalized input data, or numerical overflow; gradient clipping and loss scaling help.
- Training that looks too good to be true on validation data often signals label leakage or overlap between training and validation sets.
Last updated: September 2026
The exam guide lists troubleshooting ML model training failures under Section 3.2. Scenario questions describe a symptom, such as a job stuck in a pending state, a crash with exit code 137, or a loss that turns to NaN, and ask for the most likely cause or the best next step.
A Systematic Approach
- Classify the failure: did the job fail to start, crash while running, or finish with a bad model?
- Read the evidence: job state and error message, Cloud Logging output from each replica, Cloud Monitoring CPU, GPU, and memory utilization, and TensorBoard loss curves.
- Reproduce small: run the same container locally or on a tiny data sample with one replica.
- Change one thing at a time, and record each attempt as an experiment run.
Failures Before or at Startup
| Symptom | Likely cause | Fix |
|---|---|---|
| Job waits long or fails to get resources | Not enough accelerator quota in the region, or temporary capacity shortage | Request a quota increase, try another supported region, use Dynamic Workload Scheduler (FLEX_START), or use reservations |
STOCKOUT error with Spot VMs | Spot capacity reclaimed or unavailable | Checkpointing plus automatic retries, or on-demand VMs for deadline-critical jobs |
| Permission denied reading data or writing output | The job's service account lacks bucket or table roles, or the user lacks iam.serviceAccounts.actAs on a custom service account | Grant least-privilege roles to the job's service account. Grant actAs to the submitter |
| Image pull failure | Wrong image URI, image in a repository the service agent can't read, or an unsupported tag on a prebuilt image | Fix the URI, grant Artifact Registry read access, use a supported prebuilt tag |
| Invalid machine configuration | Accelerator type not available in the region, or an incompatible machine type and GPU combination | Choose a supported combination for the region |
Crashes While Running
| Symptom | Likely cause | Fix |
|---|---|---|
ModuleNotFoundError or version conflicts | Dependencies differ from the local environment | Pin versions in the package or container. Test the container locally with the same command |
GPU out-of-memory (CUDA out of memory) | Batch too large, sequence too long, or model too big for one device | Smaller batch plus gradient accumulation, mixed precision (bfloat16 or float16), gradient checkpointing, a bigger-memory GPU, or model-parallel / sharded training (Chapter 11) |
| Process killed (exit code 137) | The process got SIGKILL, commonly from host-memory exhaustion (for example, loading a whole dataset into RAM) | Stream data instead of loading it all, use a higher-memory machine, reduce loader workers or buffer sizes |
| Hangs or NCCL errors in multi-node jobs | Mismatched framework or NCCL versions across replicas, network setup, or a replica that crashed while others wait | Use the same image for every worker pool. Check each replica's logs. Enable timeouts and retries |
| Crash after hours of training | Preemption, a bad record late in the data, or a disk filling up | Checkpoint to Cloud Storage, validate data upfront, write large outputs to Cloud Storage instead of local disk |
Finishes, but the Model Is Bad
| Symptom in curves or metrics | Likely cause | Fix |
|---|---|---|
| Loss becomes NaN or explodes | Learning rate too high, unnormalized features, invalid values (log of 0, division by zero), float16 overflow | Lower the learning rate, normalize inputs, clean data, use gradient clipping, use loss scaling with mixed precision |
| Loss flat, not decreasing | Learning rate too low, a bug (labels out of order, frozen layers), a dead activation, or data that doesn't carry signal | Check the label pipeline, try to overfit a tiny batch (it should reach near-zero loss), then raise the learning rate |
| Training loss falls while validation loss rises (overfitting) | Model too complex for the data, or too many epochs | Early stopping, regularization (L1, L2, dropout), data augmentation, more data, a simpler model |
| Both training and validation loss stay high (underfitting) | Model too simple, weak features, too little training | Larger model, better features, longer training, lower regularization |
| Validation metrics look unrealistically high | Leakage: future information, duplicates across splits, or a feature derived from the label | Rebuild splits (time-based or group-based), remove leaky features |
| Great offline, poor in production | Training-serving skew or data drift | Share preprocessing, monitor skew and drift (Chapters 15 and 19) |
Tools for Diagnosis on Agent Platform
Interactive shell
Enable web access on a CustomJob, HyperparameterTuningJob, or custom TrainingPipeline to open a shell inside the running container. With it you can:
- Run
nvidia-smito check GPU utilization and memory. - Profile Python with
py-spywithout changing code. - Run
gcloud auth listand trybqorgcloud storagecommands to confirm what the job's service account can reach.
Requirements: bash in the image (prebuilt containers have it), a supported region, and permissions such as aiplatform.customJobs.get, plus iam.serviceAccounts.actAs if a custom service account is attached. When the job ends, the shell closes immediately. To debug a crash, catch the exception, log the traceback, and sleep before exiting. That time is billed.
Logs, metrics, and profiling
- Cloud Logging: per-replica stdout and stderr, with severity filters.
- Cloud Monitoring: CPU, memory, and accelerator utilization. Low GPU utilization usually means an input bottleneck.
- TensorBoard: loss and metric curves, and profiling of training steps. Cloud Profiler supports TensorFlow 2.4 and later on custom training.
Worked Scenario
A TensorFlow image model on one A100 crashes with GPU out-of-memory at batch size 256. The team must keep the effective batch at 256 for training stability.
- Use batch size 64 with 4 gradient accumulation steps, which keeps the same effective batch and cuts memory per step.
- Enable mixed precision, which roughly halves activation memory on supported GPUs.
- If that's still not enough, use data parallelism across 4 GPUs at batch 64 each (Chapter 11).
- Confirm the fix in TensorBoard: the loss curve should track the earlier runs.
Test Your Knowledge
A custom training job fails immediately with a permission error when it writes model artifacts to a Cloud Storage bucket. The job uses a custom service account. What is the most likely fix?
A
B
C
D
Test Your Knowledge
Training loss drops steadily while validation loss starts rising after epoch 6. What is happening, and what is a reasonable response?
A
B
C
D
Test Your Knowledge
An engineer wants to inspect GPU memory and profile Python in a running custom training job, but the shell disconnects as soon as the job crashes. What should they do?
A
B
C
D