- What: New performance and resiliency features for PyTorch training on Databricks AI Runtime.
- Core Metric: Focus on "goodput," the ratio of productive GPU computation to total training time.
- Key Technology: Implementation of PyTorch Distributed Checkpoint (DCP) and asynchronous saving.
- Infrastructure: Integration with Unity Catalog (UC) volumes for efficient state management.
Databricks has announced a suite of upgrades to its AI Runtime designed to enable fast, fault-tolerant PyTorch training at scale. By introducing more efficient checkpointing mechanisms and optimized data pipelines, the company aims to maximize "goodput"—the proportion of time GPUs spend on productive computation rather than waiting for I/O or recovering from inevitable hardware failures. These enhancements target the high-stakes environment of large-scale model training, where frequent interruptions can lead to significant resource waste and increased costs.
The Challenge of GPU Failures at Scale
As AI models grow in complexity, the infrastructure required to train them becomes increasingly prone to interruptions. Databricks identifies "goodput" as the critical metric for training efficiency, noting that at scale, GPU failures should be treated as an expected occurrence rather than an anomaly.
According to a model cited by Databricks, even with a modest 1% annualized failure rate per GPU, the probability of a job surviving its full duration drops rapidly as the cluster size increases. A 256-GPU job running for 30 days has approximately a 19% chance of experiencing a failure. When scaled to 1,024 GPUs, that probability climbs to 57%. To ground these statistics in real-world data, Databricks referenced the Delta supercomputer, where a cluster of 608 H100 GPUs experienced failures every 1.9 hours on average.
Without robust recovery mechanisms, these failures force developers to rerun hours of work, leading to massive amounts of idle GPU time and ballooning compute budgets.
Moving Beyond Monolithic Checkpointing
The primary lever for maintaining high goodput is the checkpointing mechanism—the process of snapshotting the model's state so a job can resume after a crash. Traditionally, many teams rely on a monolithic torch.save approach on "rank 0" (the primary worker). However, Databricks identifies this as a major bottleneck for two reasons:
- Rank 0 Bottleneck: In distributed training, all states must be gathered to a single process, which then writes a single file.
- Synchronous Blocking: The entire training process must pause while the checkpoint is written to storage, such as a remote object store. This leaves expensive accelerators idle for the duration of the transfer.
To solve this, Databricks AI Runtime now leverages PyTorch’s Distributed Checkpoint (DCP) API. Unlike the monolithic approach, DCP allows every rank to write its own distinct shard in parallel. This reduces the time required to save a checkpoint by a factor of roughly 1/N, where N is the number of ranks.
Asynchronous Saving and Unity Catalog Integration
Even with parallelized writing, synchronous saves can still block training for tens of seconds while data is moved to durable storage. Databricks has implemented an async_save feature to mitigate this "checkpoint tax."
The async_save mechanism splits the checkpointing operation into two phases:
- Phase 1: A fast copy of the model state to a staging buffer.
- Phase 2: A background upload to persistent storage that overlaps with the resumption of training.
By moving the heavy I/O to the background, the training loop only pays for the time required for the staging copy, making frequent checkpointing nearly "free" in terms of GPU idle time. Frequent checkpoints are essential for resilience; if a team checkpoints every hour instead of once a day, a failure only results in a maximum of 60 minutes of lost progress rather than 12 hours.
On the AI Runtime, these features are supported by UCVolumeWriter and UCVolumeReader, which implement DCP specifically for Databricks Unity Catalog (UC) volumes. This allows for seamless integration with Databricks' existing data governance and storage infrastructure.
Impact on Developers and Industry
The transition to DCP and asynchronous saving provides several direct benefits for AI developers:
- Resource Efficiency: By reducing the time GPUs spend idle during both failures and routine checkpointing, developers can significantly lower the total cost of ownership (TCO) for training large models.
- Flexible Recovery: Because DCP records the global layout of tensors in a metadata file, checkpoints can be reloaded onto a different number of GPUs. If a cluster loses nodes and must resume with reduced capacity, the system can automatically re-plan the data loading.
- Future-Proofing: While some may assume DCP is only for complex sharded models (like FSDP or tensor parallelism), Databricks notes it is also beneficial for standard Distributed Data Parallel (DDP) jobs. Adopting DCP early allows teams to build resilience code that will scale as they move to more advanced parallelization techniques.
- Data Pipeline Reliability: The company emphasizes that a data pipeline that cannot keep pace with accelerators will "silently starve" GPUs, emphasizing that both checkpointing and data ingestion are equally vital to maintaining goodput.
What's Next
Databricks has released a companion "Training performance and resiliency guide" containing code pointers and examples for developers looking to implement these features. Additionally, the company has published technical documentation on how its infrastructure detects and isolates unhealthy GPUs to prevent them from impacting active training jobs. No specific timeline for further feature releases was announced.
Sources
- Databricks Blog
- Databricks Engineering: How we keep GPUs reliable (referenced in source)