Spot GPU Training: Checkpointing, Preemption Handling and Fault-Tolerant LLM Jobs (2026)
A reserved eight-GPU node costs several times what the same node costs on the spot market, and for a month-long pretraining or fine-tuning run that gap is the difference between a funded project and a cancelled one. The catch is well known: the cloud can take the hardware back with seconds or minutes of warning. Teams that treat that as a showstopper pay full price forever; teams that treat it as an engineering constraint, bounded by arithmetic, get most of the discount. Spot GPU training is not a gamble when the job is built to lose any node at any step and lose almost no work.
This tutorial is the build guide. We start from what the clouds actually promise on interruption, then design a checkpointing path that does not stall the GPUs, choose a checkpoint interval from the Young/Daly formula with a worked example, wire up torchrun elastic restarts and Kubernetes JobSet and Kueue, and finish with the failure modes that silently ruin resumed runs: lost RNG state, a dataloader that replays data, torn checkpoints and stragglers. Numbers are labelled illustrative unless a source is cited.
What this covers: spot and preemptible semantics on AWS and Google Cloud, asynchronous sharded checkpointing with PyTorch Distributed Checkpoint (DCP), checkpoint frequency optimization, torchrun and TorchElastic behaviour, JobSet and Kueue patterns, resume correctness, a cost model, and the pitfalls that bite at scale.
Context and Background
Large-model training has three properties that make it both a perfect and a dangerous fit for interruptible capacity. It is long (days to months), it is expensive per hour (hundreds to thousands of GPUs), and it is synchronous: every data-parallel or sharded step requires every rank to participate in collectives such as all-reduce and all-gather. Lose one rank and the whole step fails. A cluster of N nodes therefore fails roughly N times as often as one node, which is the arithmetic that dominates everything below.
The clouds sell spare capacity under different names with the same bargain. AWS EC2 Spot Instances may be reclaimed when EC2 needs the capacity back. Google Compute Engine Spot VMs, the successor to preemptible VMs, can be preempted at any time, and Google advertises discounts of up to 91 percent against on-demand pricing (Compute Engine Spot VMs documentation). Azure and the GPU-specialist clouds have equivalents. Discount depth varies by region, accelerator and week, so treat any single percentage as a planning input rather than a constant.
Two other trends matter in 2026. First, GPU scarcity shifted the question from “is spot cheaper” to “can I get capacity at all”. Interruptible pools are often where H100-class and newer capacity is easiest to obtain in bulk, which is why autoscaling and bin-packing decisions matter; our comparison of Karpenter vs Cluster Autoscaler for GPU nodes covers how fleets get replaced when a node disappears. Second, frameworks matured. PyTorch Distributed Checkpoint (DCP) provides sharded, parallel, reshardable checkpoints, and torchrun provides elastic restarts, so you no longer need to build either from scratch.
What has not changed is the failure model. Hardware also fails on reserved capacity: published reports from very large training runs describe frequent GPU, network and host faults over multi-week jobs, so a fault-tolerant design is not a spot-only luxury. Spot simply raises the failure rate and makes it correlated, because a capacity crunch can reclaim many nodes in one pool at once. If you design for the spot case, the reserved case becomes easy. For the cost side of the decision beyond training, see Kubernetes cost optimization and GPU rightsizing.
A final piece of vocabulary. A checkpoint here means everything required to continue a run bit-for-bit as intended: model weights, optimizer state, learning-rate scheduler state, gradient scaler state if used, RNG states, data position and step counters. A file containing only weights is a model export, not a checkpoint. Most resume bugs are the difference between those two.
The Reference Architecture for Interruptible Training
A fault-tolerant spot training stack has four layers: a notice watcher that detects reclamation, an asynchronous sharded checkpointer that persists state without stalling the GPUs, an elastic launcher that restarts workers, and a scheduler that requeues the job. The goal is to bound lost work to one checkpoint interval plus restart time.

Figure 1: The lifecycle of an interruption in spot GPU training, from the cloud’s reclaim decision through emergency flush, requeue and resume from the last committed checkpoint.
The diagram shows two signals converging on a node-local watcher. Everything downstream is the same whether the signal was an early hint or the final notice: tell the trainer, flush what you can, let the node die, and resume elsewhere from state that is known to be complete. The rest of this section explains each layer and why it is shaped the way it is.
What the clouds actually promise
On AWS, the interruption notice is delivered roughly two minutes before the instance is stopped or terminated, through the instance metadata service and as an Amazon EventBridge event. A second, earlier signal exists: the rebalance recommendation, an indication that a Spot Instance is at elevated risk of interruption. AWS documents it as best effort, states that it can arrive together with the two-minute notice rather than before it, and says only instances launched after 5 November 2020 support it. The metadata path is events/recommendations/rebalance, which returns 404 until a signal exists, and the documentation suggests polling it about every five seconds (EC2 rebalance recommendations). The practical lesson is blunt: never design around the early warning, because you may not get it.
On Google Cloud, Spot VMs are stopped or deleted at preemption according to the termination action you configure, with stop as the default. Google documents a short shutdown window, which at the time of writing is commonly cited as 30 seconds between the preemption signal and termination; verify the figure for your machine family in the current documentation, since it has been the weakest-guaranteed part of the contract. Spot VMs have no fixed maximum runtime unless you set one, whereas legacy preemptible VMs ran for at most 24 hours. Shutdown scripts are best effort, so never put your only copy of anything in one.
Putting the two together, a conservative planning assumption is that you have between about half a minute and two minutes, and that on a bad day you have zero. A checkpoint that takes five minutes to write cannot be your interruption plan. It can only be your periodic plan.
The notice watcher and the emergency flush
The watcher is a small process (a DaemonSet on Kubernetes, a systemd unit on VMs) that polls metadata or subscribes to events and, on a signal, sends SIGTERM to the training launcher or writes a flag file the trainer polls between steps. On Kubernetes you can also rely on node drain behaviour: the kubelet sends SIGTERM to pods and waits for terminationGracePeriodSeconds, so setting that field to the notice window gives your trainer a standard hook.
What should the trainer do on the signal? Three choices, in increasing risk. The safest is nothing special: rely on the periodic checkpoint and lose up to one interval. The middle path is to finish the in-flight asynchronous checkpoint, if there is one, and skip starting a new one. The aggressive path is an emergency save of the current step, which is only worth attempting if your measured checkpoint time is comfortably below the notice window. With 30 seconds of notice and a 60-second blocking save, the emergency save will be killed mid-write, which is harmless only if your commit protocol (below) treats an unfinished checkpoint as nonexistent.
Why sharded and asynchronous
A naive checkpoint gathers the full state on rank 0 and writes one file. At 70 billion parameters with a common mixed-precision Adam setup, the state is roughly 14 bytes per parameter (2 for bf16 weights, 4 for fp32 master weights, 8 for the two Adam moments), about 1 TB. Funnelling 1 TB through one rank is slow and usually runs out of host memory. Sharded checkpointing has every rank write its own slice in parallel, so write bandwidth scales with node count. Asynchronous checkpointing then splits the work into a fast stage that copies state from GPU to pinned host memory and a slow stage that writes it out while training continues. The training loop is blocked only for the first stage.
The consequence for your interruption plan: the blocking cost C in the formulas below is the staging copy, often a few seconds, while the durable time (when the checkpoint is safe) is longer. Those are different numbers and the design must respect both.
Reliability budget in one sentence
If your cluster’s mean time between interruptions is M, your checkpoint blocks training for C seconds every T seconds, and each recovery costs R seconds, the fraction of paid GPU time that produces no progress is approximately C/T + T/(2M) + R/M. The next section derives the best T and shows what the numbers look like for a 64-node job.
Asynchronous Sharded Checkpointing with PyTorch DCP
PyTorch Distributed Checkpoint (torch.distributed.checkpoint, DCP) saves model and optimizer state as sharded files written in parallel by every rank, supports loading into a different world size, and offers async_save so that only the GPU-to-host copy blocks training. It is the standard primitive for distributed checkpointing in PyTorch.
The save path, step by step
DCP works against a Stateful abstraction. Any object that implements state_dict() and load_state_dict() can be passed in the state dictionary, and DCP calls those methods during save and load. The PyTorch recipe wraps the model and optimizer in a small class that uses get_state_dict and set_state_dict, which handle FSDP fully qualified names and the correct state-dict type (PyTorch async checkpoint recipe). The recipe lists PyTorch 2.4 or later as the prerequisite for the asynchronous API.
We extend that wrapper with everything a correct resume needs. The following is a complete, minimal trainer state object:
import random
import numpy as np
import torch
import torch.distributed as dist
import torch.distributed.checkpoint as dcp
from torch.distributed.checkpoint.stateful import Stateful
from torch.distributed.checkpoint.state_dict import get_state_dict, set_state_dict
class TrainState(Stateful):
"""Everything needed to continue the run, not just the weights."""
def __init__(self, model, optimizer, scheduler, loader_state):
self.model, self.optimizer = model, optimizer
self.scheduler = scheduler
self.loader_state = loader_state # dict with epoch, sample offset
self.step = 0
def state_dict(self):
model_sd, optim_sd = get_state_dict(self.model, self.optimizer)
return {
"model": model_sd,
"optim": optim_sd,
"sched": self.scheduler.state_dict(),
"loader": self.loader_state,
"step": self.step,
# per-rank RNG: each rank stores its own entry
"rng": {f"rank{dist.get_rank()}": {
"py": random.getstate(),
"np": np.random.get_state(),
"torch": torch.get_rng_state(),
"cuda": torch.cuda.get_rng_state(),
}},
}
def load_state_dict(self, sd):
set_state_dict(self.model, self.optimizer,
model_state_dict=sd["model"],
optim_state_dict=sd["optim"])
self.scheduler.load_state_dict(sd["sched"])
self.loader_state.update(sd["loader"])
self.step = sd["step"]
r = sd["rng"][f"rank{dist.get_rank()}"]
random.setstate(r["py"])
np.random.set_state(r["np"])
torch.set_rng_state(r["torch"])
torch.cuda.set_rng_state(r["cuda"])
A caution on the RNG block: when you resume on a different world size, rank-specific RNG entries no longer map one to one. Section “Resume correctness” covers what to do. For fixed-size resumes this layout works as written.
The save loop keeps at most one asynchronous request in flight, which is the pattern the PyTorch recipe recommends: wait for the previous future, then start the next.
CKPT_ROOT = "/mnt/shared/ckpt"
future = None
for step, batch in enumerate(loader, start=state.step):
loss = train_step(model, batch)
state.step = step + 1
if state.step % CKPT_EVERY == 0:
if future is not None:
future.result() # previous checkpoint durable
commit(prev_path) # mark it complete (see below)
prev_path = f"{CKPT_ROOT}/step_{state.step:09d}"
future = dcp.async_save({"app": state}, checkpoint_id=prev_path)
async_save returns a future; calling .result() waits for the write to finish. The recipe also documents a persistent writer with cache_staged_state_dict=True that reuses a pinned buffer between checkpoints, with the warning that the buffer persists for the life of the application and therefore raises steady host-memory pressure. In newer PyTorch (2.9 and later per the recipe), a DefaultStager makes staging itself asynchronous: the returned response exposes staging_completion and upload_completion futures, you must wait for staging before the optimizer mutates the tensors being copied, and the stager uses a background CPU thread. Check which of these your pinned PyTorch version supports before copying snippets from a newer tutorial.
The commit protocol: a checkpoint is not real until it is complete
An interrupted write leaves a directory that looks like a checkpoint and is not. On resume, the loader must never pick it up. The pattern that works on any filesystem or object store is a two-phase commit: write the data files, then write a tiny marker only after every rank reports success, then update a pointer last.
import os, json, torch.distributed as dist
def commit(path: str):
# barrier ensures every rank finished its shard before the marker exists
dist.barrier()
if dist.get_rank() == 0:
tmp = os.path.join(path, "COMMIT.tmp")
with open(tmp, "w") as f:
json.dump({"world_size": dist.get_world_size()}, f)
os.replace(tmp, os.path.join(path, "COMMIT")) # atomic on POSIX
ptr_tmp = os.path.join(CKPT_ROOT, "LATEST.tmp")
with open(ptr_tmp, "w") as f:
f.write(os.path.basename(path))
os.replace(ptr_tmp, os.path.join(CKPT_ROOT, "LATEST"))
def latest_committed():
for name in sorted(os.listdir(CKPT_ROOT), reverse=True):
p = os.path.join(CKPT_ROOT, name)
if os.path.isfile(os.path.join(p, "COMMIT")):
return p
return None
Two details deserve emphasis. os.replace is atomic on POSIX filesystems but object stores have different semantics; on S3-compatible stores, the marker object’s creation is the commit and listing consistency, not rename, is what you rely on, so verify your store’s guarantees. And scanning for the newest directory that has a COMMIT marker, rather than trusting LATEST alone, protects you from a crash between marker and pointer. Keep the last two or three committed checkpoints and delete older ones only after a newer one commits.
Storage throughput, the number that decides everything
Checkpoint time is state size divided by effective write bandwidth, and effective bandwidth is usually far below the headline figure. Take the 70-billion-parameter example: roughly 1 TB of state. If 64 nodes write in parallel to a shared parallel filesystem delivering an aggregate 40 GB/s (illustrative), the durable write takes about 25 seconds. If the same job writes to a network filesystem delivering 4 GB/s aggregate, it takes about 250 seconds, and a periodic save every 10 minutes already spends 40 percent of wall-clock in I/O unless it is asynchronous.
Three levers help. First, write locally and drain: save shards to node-local NVMe (fast, tens of GB/s per node on modern instances), then asynchronously copy to durable storage. The risk is that a reclaimed node takes its local copy with it, so a local-only checkpoint is only useful for soft failures, which is why multi-level schemes keep the durable copy at a lower cadence. Second, shrink the state: bf16 optimizer moments or 8-bit optimizers cut bytes, at a numerical cost you must validate separately. Third, avoid redundant writes in data-parallel replicas; DCP deduplicates replicated tensors so only one replica writes each, which matters for pure data-parallel and hybrid-sharded layouts.

Figure 2: The distributed checkpointing pipeline. Training resumes after the staging copy; background writes, the metadata file, an atomic commit marker and a last-updated pointer make the checkpoint durable and discoverable.
The pipeline in Figure 2 is deliberately ordered so that a crash at any point leaves the previous checkpoint intact: until the commit marker exists the new directory is ignored, and the pointer moves last.
Choosing the Checkpoint Interval: Young and Daly Applied
The checkpoint interval that minimizes expected wasted time is approximately T = sqrt(2 × C × M), where C is the time a checkpoint blocks training and M is the mean time between failures of the whole job. This is the Young/Daly first-order approximation; it is flat near the optimum, so rounding is fine.
Derivation in two lines
Checkpointing too often wastes C seconds each time; too rarely, a failure loses on average half an interval. Per unit time the overhead is C/T from saving and, since failures arrive at rate 1/M and each loses T/2 on average, T/(2M) from rework. Minimizing C/T + T/(2M) by setting the derivative to zero gives T = sqrt(2CM), and at that point the two terms are equal, so total waste is sqrt(2C/M). Daly refined the formula for higher failure rates, but the first-order version is accurate when C is much smaller than M, which holds for modern async checkpointing.
Worked example (illustrative numbers)
Suppose a 64-node job, 512 GPUs. Assume each node is interrupted on average once every 40 hours, a plausible spot figure for a contested pool but entirely workload- and region-dependent. Interruptions of different nodes are roughly independent, so the job-level mean time between interruptions is M = 40 h / 64 = 0.625 h = 2,250 s. Assume restart cost R = 300 s (replacement node, image pull, rendezvous, reload and warm-up).
| Scenario | C (blocking) | Optimal T | Save overhead C/T | Rework T/(2M) | Restart R/M | Total waste |
|---|---|---|---|---|---|---|
| Synchronous save to shared FS | 60 s | about 520 s | 11.5% | 11.6% | 13.3% | about 36% |
| Async save, staging only blocks | 8 s | about 190 s | 4.2% | 4.2% | 13.3% | about 22% |
| Async save and warm standby (R = 120 s) | 8 s | about 190 s | 4.2% | 4.2% | 5.3% | about 14% |
Three conclusions follow, and they are the practical value of the model. First, the checkpoint interval is short: three minutes, not an hour. Teams that checkpoint hourly on a 64-node spot job in this regime would waste on the order of 80 percent of the time in rework by this model (T/(2M) = 3,600/4,500), which is why people conclude spot “doesn’t work” for large jobs. Second, shrinking C is worth more than almost anything else, since waste scales with the square root of C. Cutting C from 60 s to 8 s cut the waste from the checkpoint-plus-rework terms by a factor of about 2.7. Third, once C is small, restart time dominates; the third row shows why warm standby capacity and fast image pulls matter.
Note the failure scaling: M falls linearly with node count. Doubling to 128 nodes halves M to 1,125 s and raises the optimal waste by roughly 41 percent (sqrt(2)) on the first two terms. At some size the arithmetic turns against spot for synchronous training unless you reduce per-node interruption rates by diversifying pools, using mixed on-demand and spot, or using a topology that tolerates loss of a replica.
Cost translation (illustrative)
Assume on-demand costs 4.00 dollars per GPU-hour and spot 1.60 (a 60 percent discount, illustrative). A job needing 100,000 productive GPU-hours costs 400,000 dollars on on-demand if nothing fails. On spot with 36 percent waste, the paid hours are 100,000 × 1.36 = 136,000, costing about 217,600 dollars. With async checkpointing and 22 percent waste: 122,000 hours, about 195,200 dollars. With warm standby and 14 percent waste: 114,000 hours, about 182,400 dollars. The saving versus on-demand ranges from roughly 46 to 54 percent, below the headline 60 percent discount, and the gap is the tax that checkpointing quality controls. Add storage and egress costs, which are small by comparison but not zero: keeping three 1 TB checkpoints on a fast tier for a month is a modest line item, while cross-region transfer can quietly dominate if you fail over to a different region.
Two honest caveats. The independence assumption is optimistic; correlated reclaim events (a whole zone’s pool shrinking) make interruption bursts, which hurt more than the same average rate spread evenly. And spot capacity can simply be unavailable for hours, which the formula ignores entirely. Budget for queue time, and keep an on-demand fallback path in the scheduler.
Elastic Restarts: torchrun, TorchElastic and Kubernetes JobSet
torchrun is PyTorch’s elastic launcher. When any worker fails or cluster membership changes, it stops all workers, re-forms the group through a rendezvous, and restarts every worker with fresh RANK and WORLD_SIZE values, up to a configured --max-restarts. Your script must therefore load state on every start.
What torchrun guarantees, and what it does not
The PyTorch documentation is explicit about the semantics. With --max-restarts=3, the job tolerates three failures; if some workers of the group fail, all workers are stopped and restarted. A node leaving or joining causes a membership change that stops all existing workers, forms a new worker group, and restarts each worker with new RANK and WORLD_SIZE. The recommended rendezvous backend is c10d, multi-node jobs need --rdzv-id, --rdzv-backend=c10d and --rdzv-endpoint (default port 29400 when none is given), and elasticity is expressed as --nnodes=min:max. The launcher exposes TORCHELASTIC_RESTART_COUNT and TORCHELASTIC_MAX_RESTARTS in the environment, and the docs advise checkpointing regularly and never hard-coding rank or world-size assumptions (torchrun documentation).
What it does not do is preserve any state. Surviving workers are killed immediately, so nothing in process memory survives a restart. Elasticity also does not make the math elastic: a different world size changes the global batch size unless you compensate with gradient accumulation, which changes the learning-rate schedule’s meaning. This is a design decision, not a flag.
A launch command for a 64-node job that tolerates shrinking to 56 nodes might look like this:
torchrun \
--nnodes=56:64 \
--nproc-per-node=8 \
--max-restarts=50 \
--rdzv-id=llm-sft-0421 \
--rdzv-backend=c10d \
--rdzv-endpoint=trainer-rdzv.train.svc:29400 \
train.py --config cfg.yaml --resume auto
Keep --max-restarts generous for spot, since each interruption consumes one. Count it against the job-level interruption model: at M of about 0.6 hours, a week-long run sees roughly 270 interruptions by that model, so a limit of 50 would end the job in about a day. Either set it very high and rely on an external alert for crash loops, or let the scheduler own restarts (next section) and set --max-restarts=0 so failures bubble up to a single policy.
Making the training script restart-safe
The entry point should follow one rule: the same command line works for a fresh start and for the fiftieth restart.
def main():
dist.init_process_group("nccl")
torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))
model, optim, sched, loader = build(world_size=dist.get_world_size())
state = TrainState(model, optim, sched, loader.state)
ckpt = latest_committed()
if ckpt:
dcp.load({"app": state}, checkpoint_id=ckpt)
log(f"resumed step={state.step} restart="
f"{os.environ.get('TORCHELASTIC_RESTART_COUNT', '0')}")
train_loop(state, ...)
Install a SIGTERM handler that sets a flag; the loop checks it once per step, finishes the current optimizer step, waits for any in-flight checkpoint future, and exits with a distinctive code so the scheduler can tell “preempted” from “crashed”. Never write checkpoints from inside the signal handler itself, since collectives inside a handler are a deadlock waiting to happen.

Figure 3: Sequence of one interruption handled by TorchElastic: the notice reaches the agent, workers flush and exit, the survivors rendezvous, and the restarted workers load the latest committed checkpoint and restore RNG and data position.
Kubernetes: JobSet and Kueue
On Kubernetes, the natural unit for a multi-node training run is a JobSet, a Kubernetes SIG API that manages a group of Jobs as one workload. Its failurePolicy has a maxRestarts count and an ordered list of rules; the first matching rule decides the action. The documented actions include FailJobSet, RestartJobSet (recreate all child Jobs, counted against maxRestarts) and RestartJobSetAndIgnoreMaxRestarts (recreate without counting), and rules can be scoped by targetReplicatedJobs and onJobFailureReasons such as BackoffLimitExceeded and PodFailurePolicy (JobSet failure policy). The API version shown in that documentation is jobset.x-k8s.io/v1alpha2, so pin and test against your cluster’s installed version.
That structure maps cleanly onto spot. Treat a node loss as a failure that restarts the whole JobSet without consuming the restart budget, and treat genuine code errors as budgeted failures. The mechanism that tells them apart is a Kubernetes podFailurePolicy on the child Job, which can classify a pod disruption (node drained or preempted) differently from an application exit code.
apiVersion: jobset.x-k8s.io/v1alpha2
kind: JobSet
metadata:
name: llm-sft
annotations:
kueue.x-k8s.io/queue-name: training-queue # admission through Kueue
spec:
failurePolicy:
maxRestarts: 5 # budget for real failures
rules:
- action: RestartJobSetAndIgnoreMaxRestarts
onJobFailureReasons: ["PodFailurePolicy"] # disruption, not a bug
replicatedJobs:
- name: workers
replicas: 1
template:
spec:
parallelism: 64
completions: 64
completionMode: Indexed
backoffLimit: 0
podFailurePolicy:
rules:
- action: FailJob # application error: stop and surface
onExitCodes:
containerName: trainer
operator: NotIn
values: [0, 143] # 143 = SIGTERM exit
- action: FailJob
onPodConditions:
- type: DisruptionTarget # node preempted or drained
template:
spec:
terminationGracePeriodSeconds: 120
tolerations:
- key: cloud.google.com/gke-spot
operator: Equal
value: "true"
effect: NoSchedule
containers:
- name: trainer
image: registry.example.com/trainer:2026.10
command: ["torchrun", "--nnodes=64", "--nproc-per-node=8",
"--max-restarts=0", "--rdzv-backend=c10d",
"--rdzv-endpoint=$(JOBSET_NAME)-workers-0-0.llm-sft:29400",
"train.py", "--resume", "auto"]
Treat this as a pattern sketch, not a copy-paste manifest. The exact failure-reason that surfaces for a disruption-triggered FailJob rule, the toleration key for your cloud’s spot taint, and the rendezvous address convention all depend on your versions and cluster. Test by actually deleting a node while the job runs; that single experiment finds more bugs than any review.
Kueue sits in front as the admission layer. It owns quota in a ClusterQueue, decides when a whole JobSet fits, and can queue work rather than letting half a gang of pods sit in Pending holding GPUs. For spot, the useful pattern is two ResourceFlavors, one for spot nodes and one for on-demand, with the queue preferring spot and falling back to on-demand when spot is exhausted or the job is close to a deadline. Kueue’s support for JobSet and its flavor-fungibility settings change between releases, so read the version you run before relying on specifics.

Figure 4: Kubernetes layout for fault-tolerant spot GPU training: Kueue admits the gang, JobSet restarts the whole set when a node is lost, and all pods share checkpoint storage so any replacement can resume.
The restart latency you pay on Kubernetes has several components worth measuring separately: detection (node-controller timeouts, unless a termination handler cordons the node on the notice), replacement node provisioning, container image pull, pod scheduling, NCCL initialization, and checkpoint load. Image pull of a multi-gigabyte CUDA image often dominates; pre-pulling on every node, or using a lazy-loading image format, is the single cheapest restart-time win.
Resume Correctness: The Silent Failures
A resumed run that does not crash can still be wrong. These are the bugs that burn weeks because the loss curve looks plausible.
Dataloader position and sample replay
If the loader restarts from the beginning of the epoch on each resume, a job interrupted every few hours trains repeatedly on the first slice of the data and rarely sees the tail. With 270 interruptions in a week, the head of the dataset could be seen hundreds of times more than the tail. The fix is a deterministic, resumable data order: derive the sample order from (seed, epoch), record (epoch, global_sample_offset) in the checkpoint, and on resume skip forward by index rather than by iterating. Streaming datasets and libraries that support checkpointable loaders make this cheaper; whichever you use, test it by comparing the sequence of sample IDs seen by an uninterrupted run against an interrupted-and-resumed run. They must match.
RNG state and dropout
Dropout masks, data augmentation and stochastic rounding use RNG streams. Restoring per-rank RNG state makes a fixed-size resume reproducible. For a different world size there is no exact mapping, so accept statistical rather than bitwise equivalence and instead re-seed from hash(seed, step, rank) after resume. Most large LLM runs do not need bitwise determinism, but they do need the data order and the learning-rate step to be right.
Scheduler, step counters and loss scaling
The learning-rate schedule must be a pure function of the global step, restored from the checkpoint. A scheduler that restarts at step zero re-warms the learning rate and shows up as a loss spike after each resume. Mixed-precision loss scalers and gradient-clipping statistics (if you track adaptive clipping) belong in the checkpoint as well. Log the restored step prominently on every restart.
Changing world size on resume
DCP stores shards with metadata describing the global tensors, and its loader is designed to reshard on load, so a checkpoint saved with 64 nodes can be loaded by 56. What does not reshard automatically is anything you stored per rank: RNG states, per-rank dataloader offsets. Store the data position as a global sample index, not a per-rank count, and the same checkpoint works at any size. If you change the data-parallel degree, you also change the effective batch size: either hold it fixed with gradient accumulation (56 nodes need a different accumulation than 64 for a fixed global batch, and non-divisible ratios need care) or accept and log the change.
A resume regression test
Add a CI test that trains 20 steps on a tiny model, checkpoints at step 10, kills, resumes, and asserts that the step-20 loss matches an uninterrupted run within a tolerance (exactly, if you use deterministic kernels). Run it for the real sharding strategy you use in production, since FSDP, tensor-parallel and pipeline layouts each stress the save path differently. The same discipline applies if your pipeline combines sequence parallelism; our write-up on ring attention and context parallelism for long-context training shows how those layouts partition state, which is exactly what your checkpoint must capture.
Trade-offs, Gotchas, and What Goes Wrong
Silent data corruption and torn checkpoints. A partially written shard with a valid-looking directory is the classic. Defend with the commit marker, with per-file checksums recorded in the marker, and with a periodic load test of the newest checkpoint in a separate small job. A checkpoint you have never loaded is a hope, not a backup. Corruption can also originate upstream: a GPU with a faulty memory or arithmetic unit can produce plausible but wrong gradients without crashing, and that corrupts weights that you then faithfully checkpoint. Loss-spike detectors, gradient-norm monitors and cross-rank checksums of replicated parameters catch some of it; keep several checkpoint generations so you can roll back past the onset.
Stragglers. Synchronous training runs at the speed of the slowest rank. A node with a throttled GPU, a degraded NIC or a noisy neighbour drags every step. Spot fleets amplify this because replacement nodes may land on different hardware generations or racks. Track per-rank step time and per-collective time; evict a persistent outlier proactively by treating it like an interruption, which your restart machinery already handles.
Checkpoint stalls and I/O contention. All ranks writing at the same instant produces a bandwidth spike that can slow shared storage for other tenants and extend your own staging window. Stagger writers by node group where the storage is the bottleneck, and monitor the gap between staging-complete and upload-complete; if upload regularly exceeds the checkpoint interval, the next save blocks on the previous future and your async design has silently degraded to a synchronous one.
Host memory pressure. Pinned staging buffers hold a full copy of the local shard in CPU RAM. Persisting that buffer across checkpoints, as the documented cache_staged_state_dict option does, removes allocation cost but keeps the memory reserved permanently. Size host RAM for the shard plus data-loader workers plus page cache, and set pod memory limits accordingly, or the OOM killer will take out a healthy rank mid-checkpoint.
Restart storms and crash loops. A bug that crashes on load will consume a restart budget in minutes, while a cloud capacity shortage can look identical from the scheduler’s side. Differentiate by exit code and failure reason, as the JobSet rules above do, and alert on “restarts per hour” rather than on the absolute count.
Correlated reclaim. If the whole pool is reclaimed, no elasticity helps: nothing is left to rendezvous with. Spread across two or three instance types and zones (accepting cross-zone bandwidth costs for data-parallel traffic only if your parallelism tolerates it), and keep a small on-demand base so the job can limp forward rather than stall. Tightly coupled training generally wants a single placement group for network performance, so this is a genuine tension, not a free lunch.
What spot is bad at. Jobs whose restart cost R approaches their M. If recovery takes ten minutes and nodes are interrupted every ten minutes, you make no progress. Very large synchronous jobs on shallow pools are in that regime; the right answer there is on-demand or reserved capacity for the core and spot for evaluation, data preparation and hyperparameter sweeps, which parallelize without coupling.
Practical Recommendations
Start with measurement, not configuration. Run a two-hour pilot at your intended scale on the intended pool, record interruption count and time to recover, and compute M and R from it. Plug them into the Young/Daly formula, then choose the checkpoint interval as a round number near the result. Re-measure monthly, since spot market conditions drift.
Build in this order: a correct resume first (the regression test above), then a durable commit protocol, then asynchronous sharded saves, and only then elasticity and scheduler policy. Each layer is testable on its own, and the order means that when something breaks you know which layer to blame.
Keep the economics honest. Report cost per productive GPU-hour, not per paid GPU-hour, and track waste as a first-class metric: checkpoint overhead, rework and restart time, each as a percentage. If waste climbs above roughly a third of paid time, the discount is mostly gone and you should change the design (smaller blocking time, more pools, a mixed on-demand base) rather than the interval.
A checklist to close the loop:
- Define the checkpoint as weights, optimizer, scheduler, scaler, RNG, data position and step. Test that a resume matches an uninterrupted run.
- Use DCP with
async_save, one request in flight, and a two-phase commit with a marker written after all ranks finish. - Compute T from the measured C and M; revisit after every cluster size change.
- Handle
SIGTERMby finishing the step and the in-flight save; do not rely on the early rebalance signal. - Set
terminationGracePeriodSecondsto the real notice window; classify disruptions separately from crashes in the JobSet failure policy. - Pre-pull images, keep two or three committed checkpoints, and load-test the newest one on a schedule.
- Keep an on-demand fallback queue and alert on restarts per hour, step-time outliers and checkpoint age.
Frequently Asked Questions
How much can spot GPU training really save?
Providers advertise discounts as high as the high-80s to low-90s percent for some machine types, but the realized saving on a synchronous multi-node job is lower because you pay for lost work, restarts and queue time. In the illustrative model above, a 60 percent price discount became roughly a 45 to 55 percent saving depending on checkpoint quality. Measure waste on your own pool before committing, and treat the headline discount as an upper bound.
How much notice do I get before a spot GPU is reclaimed?
On AWS EC2 the interruption notice arrives roughly two minutes before reclaim, via instance metadata and EventBridge. An earlier rebalance recommendation exists but is best effort and may arrive with the notice. On Google Cloud the shutdown window is much shorter, commonly cited as about 30 seconds, so confirm it in current documentation. Plan for little or no notice and make periodic checkpoints, not emergency saves, the foundation.
What is the best checkpoint interval for LLM training?
Use the Young/Daly approximation T = sqrt(2 × C × M), with C as the training time lost per checkpoint and M as the job’s mean time between interruptions. For a 64-node spot job with an 8-second blocking async save and M of about 37 minutes, that is roughly three minutes. Larger clusters have smaller M and therefore shorter intervals. Measure C and M rather than guessing, and recompute when node count changes.
Does torchrun handle spot interruptions by itself?
Partly. torchrun restarts all workers when one fails or membership changes, re-forming the group through rendezvous up to --max-restarts, and assigns new rank and world-size values. It does not save or restore any state, detect cloud interruption notices, or provision replacement nodes. You still need checkpointing in your script, a notice watcher, and a scheduler such as Kubernetes JobSet or Slurm to bring capacity back.
Can I resume a checkpoint on a different number of GPUs?
Often yes. PyTorch Distributed Checkpoint saves shard metadata and can reshard when loading, so a checkpoint from 64 nodes can load on 56. The catches are anything stored per rank, such as RNG state and dataloader offsets, and the change in global batch size, which affects the learning-rate schedule. Store data position as a global index and decide explicitly whether to hold batch size constant using gradient accumulation.
What is the difference between a rebalance recommendation and an interruption notice?
A rebalance recommendation is an early, best-effort signal that a Spot Instance is at elevated risk of interruption; an interruption notice is the actual two-minute warning that EC2 is reclaiming it. AWS notes the recommendation can arrive together with the notice, so you cannot depend on extra lead time. Use the recommendation to start replacement capacity early, and the notice to trigger graceful shutdown.
Further Reading
- Karpenter vs Cluster Autoscaler for GPU nodes: how replacement GPU capacity is provisioned after an interruption.
- Kubernetes cost optimization and GPU rightsizing: the wider cost model that spot training plugs into.
- Ring attention and context parallelism for long-context training: how parallelism layouts partition the state your checkpoint must capture.
- NVIDIA Dynamo and LMCache disaggregated serving: the serving-side counterpart, where interruption tolerance also matters.
- PyTorch Distributed Checkpoint async recipe and torchrun documentation.
- AWS EC2 rebalance recommendations and JobSet failure policy.
By Riju — about
