The usual concern about spot instances sounds familiar: you are 60 hours into a 72-hour fine-tuning run when the VM is reclaimed, all progress disappears, and finishing on on-demand costs more than the original savings. The concern is reasonable, but it blames the interruption instead of the missing recovery path. Training code without checkpoint and resume support is the real failure.
Teams running lengthy fine-tuning and pretraining jobs on spot do not receive special interruption rates. They make interruptions inexpensive by saving and restoring checkpoints. With a checkpoint every 15 minutes, a possible 60-hour loss becomes a 15-minute loss. That changes the economics.
This post explains workable production patterns for PyTorch and JAX, provider timing limits to account for, and how the orchestration layer keeps training code focused.
The Preemption Notice Window: What You Actually Have
Each provider offers a different warning period before reclaiming an instance. That difference should shape your shutdown handler.
GCP Preemptible (and Spot) VMs issue an ACPI shutdown signal about 30 seconds before reclamation, and your process receives SIGTERM. In practice, roughly 25-28 seconds remain after delivery. A SIGTERM handler therefore needs to start a checkpoint and flush it to persistent storage in under 25 seconds. On NVMe-backed persistent disk, models up to 7-13 billion parameters usually fit. Larger models may need faster storage, pre-staged buffers, or a fallback to the last scheduled checkpoint instead of the current step.
AWS Spot Instances expose the EC2 Instance Metadata Service at http://169.254.169.254/latest/meta-data/spot/termination-time. It returns a value 2 minutes before termination. Poll it from a background thread at a 5-second interval, then start an orderly checkpoint when the value appears. This window is much longer than GCP's and is generally enough for large checkpoints to S3.
Azure Spot VMs use the Scheduled Events API and provide a 30-second preemption warning. Actual timing differs somewhat between Deallocate and Delete eviction policies. Delete can sometimes allow less than 30 seconds.
Use a 25-second effective shutdown budget and the design will work across AWS, GCP, and Azure. A handler built around AWS's 2-minute window can fail on GCP.
PyTorch: Minimum Viable Checkpoint/Resume
A usable PyTorch checkpoint contains model weights, optimizer state, learning rate scheduler state, and the current step counter. Do not perform tensor I/O inside SIGTERM handling, since signal-context writes can disrupt active operations. Set a flag in the handler, then inspect it at the next safe point in the training loop.
import torch, signal, os, threading
_shutdown_requested = threading.Event()
def _sigterm_handler(signum, frame):
_shutdown_requested.set()
signal.signal(signal.SIGTERM, _sigterm_handler)
def save_checkpoint(model, optimizer, scheduler, step, path):
tmp_path = path + ".tmp"
torch.save({
'step': step,
'model': model.state_dict(),
'optimizer': optimizer.state_dict(),
'scheduler': scheduler.state_dict() if scheduler else None,
}, tmp_path)
os.replace(tmp_path, path) # atomic rename
for step, batch in enumerate(dataloader, start=start_step):
loss = train_step(model, batch, optimizer)
scheduler.step()
if step % checkpoint_interval == 0 or _shutdown_requested.is_set():
save_checkpoint(model, optimizer, scheduler, step, ckpt_path)
if _shutdown_requested.is_set():
break
Use os.replace() for an atomic rename. A direct write can leave a partial file if the VM dies mid-write. Write to a temporary file first, then rename it, so the checkpoint path always references either a complete new checkpoint or the prior one.
On resume, load the state and set start_step before creating the dataloader, allowing processed batches to be skipped:
start_step = 0
if os.path.exists(ckpt_path):
ckpt = torch.load(ckpt_path, map_location='cpu')
model.load_state_dict(ckpt['model'])
optimizer.load_state_dict(ckpt['optimizer'])
if scheduler and ckpt['scheduler']:
scheduler.load_state_dict(ckpt['scheduler'])
start_step = ckpt['step'] + 1
Distributed Training: Multi-Node Checkpoint Coordination
One-node checkpointing is simple. Multi-node jobs require every rank to save before any rank exits. If rank 0 gets SIGTERM and saves while rank 1 does not, the resulting checkpoint cannot be resumed cleanly.
The common approach uses a distributed barrier. When rank 0 sees preemption, it shares a shutdown flag through a distributed key-value store such as etcd, or with a direct torch.distributed broadcast. Each rank finishes its current batch, meets at the barrier, and writes its checkpoint in parallel. Rank 0 exits only after confirming every write.
In PyTorch DDP, SIGTERM reaches rank 0, the process SkyPilot monitors, while other ranks may run elsewhere. The shutdown signal must reach every rank, not only the head node. SkyPilot coordinates this at the job level: when any node faces preemption, all ranks complete shutdown coordination before any node releases its reservation.
JAX: Orbax and Async Checkpointing
JAX keeps state explicit rather than hiding it in module objects, which simplifies checkpointing. The complete training state is already a clean pytree. Orbax is the recommended library, and its async writes matter when shutdown windows are short.
Orbax's AsyncCheckpointer moves the write to a background thread while the foreground thread continues or begins an orderly shutdown. The key spot recovery pattern is:
from orbax import checkpoint as ocp
checkpointer = ocp.AsyncCheckpointer(ocp.StandardCheckpointHandler())
def on_sigterm(signum, frame):
# Trigger checkpoint write without blocking signal handler
checkpointer.save(ckpt_dir / f"step_{current_step}", state)
# Wait for write to complete, with timeout
checkpointer.wait_until_finished(timeout=20)
sys.exit(0)
signal.signal(signal.SIGTERM, on_sigterm)
The 20-second limit is intentional. On GCP, waiting beyond 20 seconds risks termination without a checkpoint. If the write is unfinished, exit and resume from the previous scheduled checkpoint.
Checkpoint Interval: The Overhead Tradeoff
Saving every step is unnecessary. Depending on model size and storage throughput, a large checkpoint can require 30 to 90 seconds of I/O. Excessive saves can turn a training job into an I/O job.
Use a maximum recomputation target of 20 minutes. At the current throughput, find the number of steps completed in 20 minutes, call it N, and save every N steps. This caps interruption loss at 20 minutes of GPU time without making checkpointing an excessive compute cost.
At 100 steps per minute, saving every 2,000 steps, or 20 minutes, is reasonable. At 2 steps per minute, saving every 40 steps gives the same time limit with much less I/O. A time-based rule transfers better across hardware than a fixed step interval.
Post-Resume Validation
Validate the checkpoint immediately after resuming, before training continues. Load it, run a forward pass on a fixed-seed batch, and compare the loss with a reference saved beside the checkpoint. A matching value within tolerance confirms validity. A mismatch indicates corruption, so fall back to the previous checkpoint.
Partial writes can occur, especially when storage is busy during rapid shutdown. Atomic rename lowers the risk but cannot remove it. With two-checkpoint rotation, retain the previous file until the new one passes validation, then delete it. This uses 2x checkpoint storage.
Spot instances trade lower cost for the possibility of reclamation. The teams that gain the most treat interruptions as routine events handled by code, not failures to avoid. Once checkpoint and resume works, restart location becomes an orchestration concern.