Designing GPU workloads for preemption¶

FLAME's GPUs are shared across research groups. To keep them busy, the cluster lets your group borrow GPUs that are allocated to other groups, but which they aren't currently using. So you can often use more resources than your group is strictly guaranteed. The catch is that borrowed GPUs can be reclaimed at any time: when the owning group needs them back, the cluster stops one of your running workloads to free them up. This is called preemption.¹

Preemption is not a failure or an error — it's normal, expected, and can happen to any workload that's borrowing. The point of this guide is to make preemption a minor interruption rather than a disaster: if your workload saves its progress as it goes, it picks up where it left off and you lose only a little work; if it doesn't, you lose everything it had computed.

Designing your workloads to save their state as they go is also a general research computing best practice. It guards against other kinds of sporadic issues (e.g. hardware failures) in addition to pre-emption.

What happens during pre-emption¶

When the cluster preempts your workload, it doesn't kill it instantly. It gives you a chance to wrap up first via the following process:

1. Your workload recieves a SIGTERM signal (the standard Unix "please shut down" request).
2. It waits for a grace period to let your workload save its state and exit cleanly.
3. If your workload is still running when the grace period ends, it's forcibly terminated.

After that, the cluster automatically puts the workload back in the queue and restarts it once GPUs are free again. You don't need to re-run anything manually.

Checkpoint and resume¶

You are, however, responsible for saving and reloading your workload's progress so that when it restarts after being pre-empted, it can pick up where it left off. The standard approach is:

1. Save your progress periodically to persistent storage while the workload runs — for example, write a checkpoint of your training run every N steps. Persistent storage survives the workload being stopped and restarted; anything kept only in the pod's memory or local disk does not.
2. Handle SIGTERM: when your workload receives it, save a final checkpoint and exit.
3. Resume from the last checkpoint on startup: when the cluster restarts your workload, have it look for an existing checkpoint and continue from there instead of starting over from scratch.

With this in place, a preemption costs you at most the work done since your last checkpoint (typically seconds to minutes) instead of the entire run.

You usually don't have to build this from scratch — most popular frameworks have tooling and guides for this checkpoint-and-resume pattern:

If your framework isn't listed, search its docs for "checkpointing," "resuming," or "fault tolerance."

A workload bigger than your guarantee¶

There's a surprising case worth calling out. If a single workload asks for more GPUs of a type than your group is allocated, the entire workload is treated as borrowing — not just the GPUs beyond your guarantee. That means the entire workload can be preempted as a unit, even the part that would have fit inside your guaranteed capacity.

For example, imagine your group has an allocation guaranteeing 4 H200s. You being training a model using 5 H200s. Because the workload as a whole exceeds your guarantee, the entire workload is considered to be borrowing. If another group reclaims their guaranteed capacity, the cluster can preempt your entire workload, not just the 1 extra GPU.

If you need your work not to be pre-empted, work entirely within your group's guaranteed capacity.