How data and checkpointing choices decide your GPU utilization, recovery cost, and training bill at scale and the AI Runtime APIs that get them right.
At scale, your training efficiency is determined by a single metric: "goodput", the proportion of time your GPUs spend on productive computation rather than waiting or recovering from failures. Because GPU failures are the expected case at scale, the ability to rapidly and automatically recover from a failure is the only way to maintain high goodput and manage your total GPU spend.
Two subsystems make or break that recovery, yet both are routinely treated as afterthoughts: the data pipeline that feeds your accelerators, and the checkpointing mechanism that snapshots state so a job can resume. Get either one wrong and every failure costs you far more idle GPU time than it should. Even outside of failure scenarios, a data pipeline that can't keep pace with your accelerators will silently starve your GPUs and erode goodput just as surely as a crash would. We'll walk through the mechanisms and trade-offs of both, and how each one shapes your goodput and total GPU spend. See the companion Training performance and resiliency guide for code pointers and examples.
For the infrastructure side of the same problem, how a fleet detects and isolates unhealthy GPUs before they take down a job, see the companion post, How we keep GPUs reliable across Databricks AI. As the number of GPUs in a job grows, the probability that it survives its full duration without an interruption falls rapidly. A useful back-of-the-envelope model from the companion Databricks post assumes each GPU carries roughly a 1% annualized failure rate. Under that assumption, the post notes that "a 256-GPU job running for 30 days has about a 19% chance of seeing a failure. At 1,024 GPUs, that climbs to 57%." and these are just infrastructure level issues.
To ground that estimate in reality, the 608 H100 GPUs delta super computer saw failures every 1.9 hours, this means that for a 32 GPU job, the average time to failure would be 36 hours. The main take away, is that your training job will likely fail at some point and making the correct decisions can make your model resilient and reduce the total time lost when it happens.
Checkpointing is where resilience is won or lost, and the mechanism you choose has a first-order effect on how frequently you can save. This is the single biggest lever on your goodput: if you checkpoint once a day, then a failure requires rerunning on average 12 hours of duplicate work to bring your back to the state it was in when the failure occurred.
The first checkpoint most teams write is a simple torch.save on rank 0. Depending on how your model is trained, potentially two issues:
This blocking behaviour leaves your GPUs idle, reducing your goodput. But there is a way to reduce the amount of time your GPU spends checkpointing: Torch’s distributed checkpoint API.
PyTorch's distributed checkpoint inverts the design. Every rank writes its own distinct shard in parallel, alongside a small .metadata
file describing how the shards compose into the full tensors.
Saving time decreases roughly as 1/N with the number of ranks and, because the .metadata
file records the global layout, the same checkpoint can reload onto a different number of GPUs. DCP re-plans which bytes each new rank needs, so recovering onto a reduced-capacity cluster after losing nodes just works.
A common assumption is that DCP is only for sharded models, that a data-parallel (DDP) job, where every rank holds an identical replica of the weights, has nothing to gain. Not so, DCP shards the model state and writes it in parallel across each worker even for DDP training tasks.
It is also the same API you will need the day you move to FSDP or tensor parallelism, so adopting it early means you never rewrite resilience code at the worst possible time.
Even with parallel writes, a synchronous save blocks training until the bytes are durable in storage, for a large checkpoint to a remote volume, tens of seconds of idle accelerator time. async_save
splits the operation: a fast copy to a staging buffer, then a background upload that overlaps continued training.
The training loop pays only for the staging copy, not the upload. A checkpoint that used to cost tens of seconds of idle time now costs almost nothing, which is exactly what makes the frequent checkpointing in the next section affordable.
On AI Runtime, UCVolumeWriter
and UCVolumeReader
implement DCP against UC volumes, staging I/O through local NVMe and marking a checkpoint complete only once its data has fully landed. See the performance and resiliency guide for full details and code examples.
| Training Job | Savings of async_save over torch.save |
|---|---|
| DDP LLM with 2.8B parameters on 32xH100 | 1.8x (36s vs 66s) |
| FSPD LLM with 20B parameters on 32xH100 | 58x (522s vs 9s) |
The above excludes the network storage time for torch.save.
This is where the pieces compound. When a job fails, it loses everything since the last valid checkpoint and must recompute it. So the expected wasted work per failure is about half the checkpoint interval and cheap async saves let you make that interval small.
Cutting the interval by a factor of 10 cuts expected time to recover by a factor of 10. Recall the Llama 3 figure of ~8.6 interruptions per day: at that failure rate, checkpointing every 2 hours means you expect to waste 8.6 hours per day on retraining, a goodput of 64%. Checkpointing every 30 minutes, you only spend 2.15 hours, a goodput of 91%.
The recovery must also be automatic. On restart, the job should find the most recent checkpoint that finished writing, skipping any left half-written by the crash, and resume from it with no human in the loop. DCP makes this reliable: the .metadata
file is written only after all shards land, so its presence is a trustworthy "this save is complete" marker to select on.
A training job proceeds at the speed of its slowest input. When accelerators wait on the next batch, your goodput is reduced as your GPUs are simply idle. The only way to fix this issue is to ensure that your input pipeline overlaps data preparation for the next step with computation on the current one as seen in the figure below:
We often see customers that shift to overlapping data with compute see a 20–50% decrease in wall-clock time.
On a governed platform, training data lives in remote object storage. On AI Runtime, Unity Catalog (UC) volumes are surfaced as network mounts.
Reading files directly from that mount on every access binds your step time to network latency and re-downloads the same files every epoch. The fix is a data that copies each file to fast local storage on first access, serves subsequent reads from that local cache, and fetches upcoming files in parallel while the GPU computes.
With AI Runtime, UCVolumeDataset
and Data
do exactly this (see the guide for code examples) . UCVolumeDataset
streams files from a UC volume, caching each one to local NVMe on first access, and partitions files across ranks and workers so every accelerator gets a disjoint, non-overlapping slice. Our Data
is a drop-in subclass of the PyTorch Data
whose defaults are tuned for this path, so files are fetched and cached concurrently while the GPU computes instead of one at a time on the training thread.
Consider a straightforward image-classification workload: decode JPEGs from a UC volume, augment, and train a vision model. Let’s look at two ways to do this on the same GPU, model, and batch size: the stock PyTorch Dataset
reading from a UC volume versus UCVolumeDataset
plus the Databricks Data
defaults.
| Metric (per GPU, steady state) | Stock PyTorch Data, reading directly from UC | UCVolumeDataset + databricks Data |
|---|---|---|
| Epoch 1 Throughput (images/sec) | 57.2 | 417 |
| Epoch 2 Throughput (images/sec) | 371.6 | 6590 |
| GPU utilization (%) | 12.6% | 53.3% |
As part of engineering Data
, we’ve ensured that it logs its metrics to MLFlow, making it easy to tell at a glance if your data pipeline is blocking training.
The metric fetch_seconds
measures explicitly how long it takes the data to produce a batch and during this time your GPU is sitting idle.
There is one last resilience bug that produces no error message, no crash, and no failed job, just a model that is subtly worse than it should be. It happens when you checkpoint the model, optimizer, and step, but not the position of your data pipeline within the dataset.
Consider a job interrupted partway through an epoch. It restores the model correctly and resumes the training loop but the data starts over from the beginning of the dataset.
The resumed job re-trains on examples it already saw this epoch and potentially skips the ones it hadn't reached yet. Across the many restarts that scale makes routine, this silently biases your data distribution. The model still trains; it just trains on the wrong sampling of your data, precisely the kind of silent failure that’s the most costly, because the job completes and nobody sees a problem until the metrics are disappointing.
The fix is to treat data position as part of the checkpoint. Depending on your pipeline, that means tracking a sample or shard offset and skipping ahead on resume, having a custom dataset serialize its own position, or checkpointing at epoch boundaries. All of these rest on one prerequisite: determinism. Shuffling and augmentation draw from random number generators, so those seeds and RNG states must be part of the checkpoint too, otherwise the data order after a restart won't match the order before it, and a saved position points at the wrong samples.
Seed, reproducible order, and resumable data pipeline are three expressions of a single idea. The guide covers each strategy with code.
Fast, fault-tolerant training comes from a handful of decisions that compound:
torch.save
, even for DDP, so saves are parallel and cheap rather than a serial bottleneck.The unifying principle: frequent, inexpensive, complete checkpoints turn a hardware failure from a job-ending event into a rounding error, and an overlapped input pipeline keeps the accelerators busy in between. Cheap (async) saves make frequency affordable; complete saves (model, data, and RNG) make recovery correct. With both in place, and a fleet that detects and isolates failing hardware, your effective training time approaches the ceiling the hardware allows, regardless of how flaky the cluster underneath it is.
Ready to try it? See the Training performance and resiliency guide in the Databricks AI Runtime docs for the full code, and read How we keep GPUs reliable across Databricks AI for the infrastructure side of the story.
Subscribe to our blog and get the latest posts delivered to your inbox.