Fast, fault-tolerant PyTorch training on AI Runtime

Fast, fault-tolerant PyTorch training on AI Runtime | Databricks Blog Skip to main content

Summary

At scale, GPU failures are the expected case, not the exception, code must be built to survive them.

Torch’s distributed asynchronous checkpoint saves make frequent checkpointing nearly free, enabling more frequent checkpointing and cutting recovery cost.

Model checkpointing by itself is not sufficient, checkpointing the data pipeline prevents silent training-data corruption on resume.

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 . Why failures are the expected case at scale 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. Impact 1: Checkpoint format decides how often you can afford to save 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 monolithic torch.save bottleneck The first checkpoint most teams write is a simple torch.save on rank 0. Depending on how your model is trained, potentially two issues: For distributed training, it gathers all states to rank 0 and writes a single file. A single process writes the entire checkpoint to sync synchronously. This can be blocked on things like network transfers when saving to remote object stores like Unity Catalog (UC).

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. Distributed checkpoint (DCP): every rank writes its own shard 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. DCP is worth it even for plain data-parallel jobs 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. Asynchronous saves make frequency nearly free 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. Impact 2: Checkpoint frequency decides your recovery cost 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...