Databricks Says Async Save Cut 20B-Parameter PyTorch Checkpoints to 9 Seconds on 32 H100s

Large PyTorch training jobs do not just lose time when a GPU fails. They also lose it while checkpoints block or while accelerators wait for the next batch. Databricks describes AI Runtime components for both problems: distributed asynchronous checkpoint I/O for Unity Catalog volumes and a cached data path intended to keep training moving through failures and remote-storage delays.

When failures arrive, checkpoint design sets the recovery bill

Databricks says automatic recovery becomes necessary as jobs use more GPUs and run longer. Its approach starts with PyTorch distributed checkpointing: each training rank writes a state shard in parallel, and the checkpoint can later be reloaded on a different number of GPUs. That avoids a single rank gathering and synchronously writing the entire model state.

Asynchronous saves then separate a quick staging step from the slower upload. PyTorch stages checkpoint data and uploads it in the background while training continues; on AI Runtime, UCVolumeWriter and UCVolumeReader implement that distributed checkpoint path for Unity Catalog volumes, using local NVMe for staging. A checkpoint is marked complete only after its data has landed, giving a restarted job a way to avoid a partially written save.

The input path can become the next bottleneck

Recovery is only half the utilization problem. Training data stored in Unity Catalog volumes is accessed through network mounts; reading files from that remote path on every access can tie each step to network latency and download the same files again in later epochs. Databricks’ UCVolumeDataset copies files to local NVMe on first access, partitions files across ranks and workers, and fetches upcoming files while the GPU computes.

The company’s image-classification example reported 6,590 images per second in the second epoch with UCVolumeDataset and its DataLoader, compared with 371.6 for a stock PyTorch DataLoader reading directly from Unity Catalog. That is a company-supplied workload result, not a general performance guarantee, but it illustrates why the cache matters more after its first pass is populated.

A restart can still change what the model sees

The final safeguard is less visible than a checkpoint file. Databricks warns that restoring model and optimizer state without restoring the data pipeline’s position can silently corrupt the training-data sequence or degrade model quality after a restart. Its DataLoader logs pipeline metrics to MLflow, including fetch_seconds, the time required to produce a batch when a GPU may be idle.

That makes the package a broader reliability pattern, not simply a faster save call: parallel checkpoint shards, background persistence, selection of a completed checkpoint, and a resumable input pipeline must work together. Databricks has supplied benchmark examples and the runtime components; teams adopting them still need to preserve the full training state.