Fast, fault-tolerant PyTorch training on AI Runtime

Summary

At scale, training jobs face frequent GPU interruptions, while remote-storage input pipelines can leave accelerators idle. Distributed Checkpoint (DCP) has each rank write a shard plus a .metadata file, allowing parallel saves and reloads onto a different number of GPUs; DCP is presented as useful for DDP as well as sharded training. Asynchronous saves stage data locally and upload in the background; cited measurements report 1.8x savings for a 2.8B-parameter DDP model and 58x for a 20B-parameter FSPD model on 32 H100s, though the comparison excludes torch.save network storage time. UCVolumeDataset and the accompanying DataLoader cache files on local NVMe and prefetch disjoint rank/worker slices while computation runs. Correct recovery also requires checkpointing data position plus seeds and RNG state so resumed sampling remains consistent.

Context

Large training jobs are likely to experience hardware failures, making recovery cost important to goodput and GPU spend. Separately, training data stored on remote object storage can make GPUs wait for input. Resuming without the data pipeline's position and random-state information can also retrain or skip examples and silently change the data distribution.

Approach / What changed

Use PyTorch Distributed Checkpoint so ranks write shards in parallel, combine it with asynchronous saves and automatic selection of the most recent completed checkpoint, and use UCVolumeWriter and UCVolumeReader with local NVMe staging. Use UCVolumeDataset and the tuned DataLoader to cache remote files, partition them across ranks and workers, and prefetch during computation. Include data position, seeds, and RNG state in checkpointed state.

Takeaways

• DCP writes checkpoint shards in parallel across ranks and uses the .metadata file as a completion marker; its recorded global layout allows recovery on a different number of GPUs.
• The stated example estimates 64% goodput when checkpointing every two hours with 8.6 interruptions per day, versus 91% when checkpointing every 30 minutes.
• A resumable data pipeline must preserve its position and the relevant seeds and RNG states; otherwise a restart can repeat or skip samples and silently alter training data distribution.

Related reading