Superpower Daily
Databricks Says Async Save Cut 20B-Parameter PyTorch Checkpoint...Databricks Files $5B Equity Offering With Nearly All of It Alre...OpenAI’s 1,200 Test Agents Built a Covert Network and Breached...Snowflake Open-Sources Semi-Persistence With Sub-Second Single-...Midjourney Opens V8.2 Image-Editing Tests With Up to Four Refer...Harvard Uses 7 AI Faculty Avatars in a $699 Founder Bootcamp, W...Databricks Says Async Save Cut 20B-Parameter PyTorch Checkpoint...Databricks Files $5B Equity Offering With Nearly All of It Alre...OpenAI’s 1,200 Test Agents Built a Covert Network and Breached...Snowflake Open-Sources Semi-Persistence With Sub-Second Single-...Midjourney Opens V8.2 Image-Editing Tests With Up to Four Refer...Harvard Uses 7 AI Faculty Avatars in a $699 Founder Bootcamp, W...
Toolspublished

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

The AI Runtime components pair distributed recovery with local-NVMe data caching. The remaining operational test is whether training jobs save their input position as well as their model state.

By 3 min read
Databricks Says Async Save Cut 20B-Parameter PyTorch Checkpoints to 9 Seconds on 32 H100s
Sign in to save and react

Listen to this story

The audio brief

About 0:13
0:000:13
Read transcript
Databricks says it cut checkpoint saves for a 20-billion-parameter PyTorch model from 522 seconds to 9 seconds, using asynchronous distributed saves across 32 H100 GPUs. The result comes from the company’s AI Runtime, and it changes where the waiting happens: each training process stages its checkpoint shard on local NVMe, then uploads it in the background while training continues. A save is considered recoverable only after all of its data has landed, which helps prevent a restart from selecting a partially written checkpoint. The comparison is not perfectly apples to apples: Databricks says the 522-second torch.save figure excludes network-storage time. Still, the design addresses a real scaling problem. With more GPUs and longer jobs, failures make recovery expensive, and synchronous persistence can leave expensive accelerators idle. Databricks is also targeting the other side of that pause: input data. Its UCVolumeDataset caches files from Unity Catalog volumes on local NVMe, distributes work across ranks and workers, and prefetches upcoming files while the GPU computes. In one company-supplied image-classification benchmark, second-epoch throughput rose from 371.6 to 6,590 images per second. That is a specific workload result, not a general guarantee. The key constraint is restart fidelity. Restoring weights and optimizer state is not enough; the data pipeline’s position must also be restored, or training quality can suffer. MLflow metrics such as fetch_seconds expose whether data access is leaving GPUs idle. The open question is whether teams can preserve that full training state in production.

Story brief

3 key points

Databricks reports that its AI Runtime reduced checkpoint-save time for a 20-billion-parameter FSDP model from 522 seconds with `torch.save` to 9 seconds using asynchronous distributed saves on 32 H100 GPUs. The system stages shards on local NVMe, uploads them in the background, and marks only fully persisted checkpoints as recoverable. A companion cached input pipeline reportedly improved second-epoch image...

  1. 01

    Async checkpointing separates local staging from background upload, allowing training to continue during persistence.

  2. 02

    Databricks’ UCVolumeDataset caches remote files on local NVMe and prefetches upcoming data across ranks and workers.

  3. 03

    The reported 6,590-versus-371.6 images-per-second result applies to a specific image-classification workload, not all training jobs.

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.

Diagram showing data preparation overlapping with GPU computation during PyTorch training.
Databricks presents overlapping data preparation with computation as the way to prevent accelerators from waiting for the next batch. Source: databricks.com.

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.

Sources

  1. databricks.comFast, fault-tolerant PyTorch training on AI Runtime

Related coverage

Continue reading