翻訳待ち：Fast, fault-tolerant PyTorch training on AI Runtime

AI サービスが一時的に利用できないため、復旧後に翻訳を補完します。

Fast, fault-tolerant PyTorch training on AI Runtime | Databricks Blog Skip to main content 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 JobSavings of async_save over torch.save DDP LLM with 2.8B parameters on 32xH1001.8x (36s vs 66s) FSPD LLM with 20B parameters on 32xH10058x (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 minutes, you only spend 2.15 hours, a goodput of 91%. The recovery must also be automatic. On restart, the job should find the most recent checkpoint that finished writing, skipping any left half-written by the crash, and resume from it with no human in the loop. DCP makes this reliable: the .metadata file is written only after all shards land, so its presence is a trustworthy "this save is complete" marker to select on. Impact 3: Dataloading decides whether your GPUs are ever idle A training job proceeds at the speed of its slowest input. When accelerators wait on the next batch, your goodput is reduced as your GPUs are simply idle. The only way to fix this issue is to ensure that your input pipeline overlaps data preparation for the next step with computation on the current one as seen in the figure below: We often see customers that shift to overlapping dataloading with compute see a 20–50% decrease in wall-clock time. The cost of reading straight from remote storage On a governed platform, training data lives in remote object storage. On AI Runtime, Unity Catalog (UC) volumes are surfaced as network mounts. Reading files directly from that mount on every access binds your step time to network latency and re-downloads the same files every epoch. The fix is a dataloader that copies each file to fast local storage on first access, serves subsequent reads from that local cache, and fetches upcoming files in parallel while the GPU computes. With AI Runtime, UCVolumeDataset and DataLoader do exactly this (see the guide for code examples) . UCVolumeDataset streams files from a UC volume, caching each one to local NVMe on first access, and partitions files across ranks and workers so every accelerator gets a disjoint, non-overlapping slice. Our DataLoader is a drop-in subclass of the PyTorch DataLoader whose defaults are tuned for this path, so files are fetched and cached concurrently while the GPU computes instead of one at a time on the training thread. Example: training an image model off UC files Consider a straightforward image-classification workload: decode JPEGs from a UC volume, augment, and train a vision model. Let’s look at two ways to do this on the same GPU, model, and batch size: the stock PyTorch Dataset reading from a UC volume versus UCVolumeDataset plus the Databricks DataLoader defaults. Metric (per GPU, steady state)Stock PyTorch DataLoader, reading directly from UCUCVolumeDataset + databricks DataLoader Epoch 1 Throughput (images/sec)57.2417 Epoch 2 Throughput (images/sec)371.66590 GPU utilization (%)12.6%53.3% You don't have to guess where the time goes As part of engineering DataLoader, we’ve ensured that it logs its metrics to MLFlow, making it easy to tell at a glance if your data pipeline is blocking training. The metric fetch_seconds measures explicitly how long it takes the dataloader to produce a batch and during this time your GPU is sitting idle. Impact 4: Forgetting the data pipeline silently corrupts your model There is one last resilience bug that produces no error message, no crash, and no failed job, just a model that is subtly worse than it should be. It happens when you checkpoint the model, optimizer, and step, but not the position of your data pipeline within the dataset. Consider a job interrupted partway through an epoch. It restores the model correctly and resumes the training loop but the dataloader starts over from the beginning of the dataset. The resumed job re-trains on examples it already saw this epoch and potentially skips the ones it hadn't reached yet. Across the many restarts that scale makes routine, this silently biases your data distribution. The model still trains; it just trains on the wrong sampling of your data, precisely the kind of silent failure that’s the most costly, because the job completes and nobody sees a problem until the metrics are disappointing. The fix is to treat data position as part of the checkpoint. Depending on your pipeline, that means tracking a sample or shard offset and skipping ahead on resume, having a custom dataset serialize its own position, or checkpointing at epoch boundaries. All of these rest on one prerequisite: determinism. Shuffling and augmentation draw from random number generators, so those seeds and RNG states must be part of the checkpoint too, otherwise the data order after a restart won't match the order before it, and a saved position points at the wrong samples. Seed, reproducible order, and resumable data pipeline are three expressions of a single idea. The guide covers each strategy with code. Summary Fast, fault-tolerant training comes from a handful of decisions that compound: Use Distributed Checkpoint instead of torch.save, even for DDP, so saves are parallel and cheap rather than a serial bottleneck. Save asynchronously so checkpoints are nearly free, which lets you save often. Recover automatically to the most recent valid checkpoint, so a failure costs minutes of recomputation, not hours. Overlap data loading with compute by caching and prefetching from remote storage so accelerators never idle waiting for input. This is recurring GPU-hours saved on every step. Checkpoint the data pipeline and RNG state, so a resumed job continues on the correct data instead of silently corrupting your model. The unifying principle: frequent, inexpensive, complete checkpoints turn a hardware failure from a job-ending event into a rounding error, [truncated for AI cost control]