Checkpointing: Surviving Preemption and Node Failures

Why Checkpointing Is Non-Negotiable

A 48-hour training run on 32 GPUs represents enormous compute investment. If a node fails at hour 47, losing all progress is catastrophic. On a shared cluster with preemption enabled, your job can be evicted at any time. Checkpointing — periodically saving model weights and optimizer state to disk — is the practice that makes long training runs survivable.

Slurm does not handle checkpointing for you. It's the responsibility of your training code. However, Slurm provides signals that let your code checkpoint gracefully before being killed.

Signal-Based Graceful Shutdown

Configure Slurm to send a warning signal (say, SIGUSR1, 120 seconds before the time limit) via a single --signal directive.

In your Python training code, catch this signal and trigger a checkpoint save:

import signal
import sys

def checkpoint_and_exit(signum, frame):
    print(f"Received signal {signum}. Saving checkpoint before exit...")
    save_checkpoint(model, optimizer, step, "/scratch/checkpoints/emergency.pt")
    sys.exit(0)

signal.signal(signal.SIGUSR1, checkpoint_and_exit)

Periodic Checkpointing in Training Loops

Beyond signal handling, always checkpoint on a regular interval:

for step in range(start_step, total_steps):
    loss = train_step(batch)

    # Checkpoint every 500 steps
    if step % 500 == 0 and rank == 0:
        torch.save({
            "step": step,
            "model_state_dict": model.state_dict(),
            "optimizer_state_dict": optimizer.state_dict(),
            "loss": loss,
        }, f"/scratch/checkpoints/step_{step:06d}.pt")

    # Keep only the last 3 checkpoints to save disk space
    cleanup_old_checkpoints(keep=3)

Resuming from Checkpoint

Structure your job script to automatically resume if a checkpoint exists: at startup, the script looks for the most recent checkpoint file in its checkpoint directory, and if one is found, passes its path to the training command as a --resume-from argument instead of starting fresh.

This pattern — combined with Slurm's --requeue flag — allows a job to be automatically resubmitted after preemption and pick up exactly where it left off.