Durable training

The longer a training job runs, the more likely something under it fails: a spot node is reclaimed, a node fails, a batch runs out of memory, or the cluster has no free H100s. Flyte turns each of these into something your code handles, instead of a run you start over.

Resume from a checkpoint after a retry

flyte.ctx().checkpoint is a per-task store in your object store. It survives across a task’s retries. Save to it as training progresses, and load from it at the start of the task:

@env.task(retries=RETRIES)
async def train_linear(epochs: int = 10) -> float:
    checkpoint = flyte.ctx().checkpoint

    model = nn.Linear(4, 1)
    opt = torch.optim.SGD(model.parameters(), lr=0.01)

    # Resume model, optimizer, and epoch from the previous attempt, if any.
    prev = await checkpoint.load()
    if prev:
        blob = torch.load(prev, map_location="cpu", weights_only=False)
        model.load_state_dict(blob["model"])
        opt.load_state_dict(blob["opt"])
        start = int(blob["epoch"]) + 1
    else:
        start = 0

    wpath = pathlib.Path("pytorch_linear") / "training.pt"
    wpath.parent.mkdir(parents=True, exist_ok=True)

    failure_interval = epochs // RETRIES
    for epoch in range(start, epochs):
        x = torch.randn(8, 4)
        y = torch.randn(8, 1)
        loss = torch.nn.functional.mse_loss(model(x), y)
        opt.zero_grad()
        loss.backward()
        opt.step()

        if epoch > start and epoch % failure_interval == 0:
            # Simulate a failure so the next attempt resumes from the checkpoint
            raise RuntimeError(f"Simulated failure at epoch {epoch}")

        # Save model, optimizer, and epoch state to object storage.
        torch.save(
            {"model": model.state_dict(), "opt": opt.state_dict(), "epoch": epoch},
            wpath,
        )
        await checkpoint.save(wpath)

    with torch.no_grad():
        return float(model(torch.ones(1, 4)).squeeze().item())

Save the optimizer state and the step or epoch along with the weights, so the resumed attempt continues the same schedule. To connect a framework’s own checkpointing, such as the Hugging Face Trainer or Lightning’s ModelCheckpoint, use a callback. See Intra-task checkpoints, and the Brain tumor MRI classification tutorial for a Lightning training task that resumes this way.

Use spot capacity

Spot and preemptible GPUs cost a fraction of on-demand ones. With checkpoints in place, a preemption costs only the work since the last save. Set interruptible=True and give the task retries:

import flyte

train_env = flyte.TaskEnvironment(
    name="train",
    resources=flyte.Resources(gpu="A100:8"),
    interruptible=True,
)


@train_env.task(retries=5)
async def train(epochs: int) -> flyte.io.Dir: ...

Keep the driver and evaluation tasks on on-demand capacity with interruptible=False. See Interruptible tasks.

Restart a multi-node job

When one pod of a clustered job fails, the other ranks can’t continue. ClusterFailurePolicy restarts the whole group in place:

import flyte
from flyte.clustered import ClusterFailurePolicy, ClusteredTaskEnvironment, TorchRun

env = ClusteredTaskEnvironment(
    name="pretrain",
    resources=flyte.Resources(gpu="H100:8", shm="16Gi"),
    replicas=4,
    nproc_per_node=8,
    runtime=TorchRun(),
    failure_policy=ClusterFailurePolicy(max_restarts=2, restart_on_host_maintenance=True),
)

• max_restarts is the number of in-place restarts before the attempt fails. They don’t count against the task’s retries.
• restart_on_host_maintenance=True makes a restart caused by a node eviction free, so planned maintenance doesn’t use up max_restarts.
• resume_after_preemption=True keeps the attempt when Kueue preempts the job after it started. The job waits to be readmitted and then resumes, without using a retry. It needs Flyte backend v2.0.54 or later. Without it, a preemption ends the attempt, so set retries on the task.

Pod-local disk is wiped on a restart, so write checkpoints from rank 0 to flyte.ctx().checkpoint. flyte.ctx().restart_attempt counts the in-place restarts, which is useful for logging. See Clustered task environment.

Recover from out-of-memory errors

When a container runs out of host memory, the parent task receives flyte.errors.OOMError and can call the task again with more memory:

import flyte.errors


@driver_env.task
async def main(epochs: int) -> flyte.io.Dir:
    for memory in ["256Gi", "512Gi", "1Ti"]:
        try:
            resources = flyte.Resources(gpu="A100:8", memory=memory)
            return await train.override(resources=resources, retries=0)(epochs=epochs)
        except flyte.errors.OOMError:
            continue
    raise RuntimeError("Out of memory at every size")

override keeps the task’s own retries unless you set them. retries=0 sends an out-of-memory error straight back to the driver, instead of trying the same size again first.

Running out of GPU memory is different: the framework raises its own exception, such as torch.cuda.OutOfMemoryError. See Error handling to catch it and retry with a smaller batch.

Fall back to another GPU type

A job that asks for scarce GPUs can wait in the queue for a long time. Set timeout=flyte.Timeout(max_queued_time=...) on the training task to bound the wait. When it expires, the driver catches flyte.errors.MaxQueuedTimeExceededError and calls the task again with override(resources=...) for the next GPU type on its list. See Falling back when GPU capacity isn’t available for the full example.

Reuse finished stages

A pipeline that fails at evaluation shouldn’t retrain the model. Two features keep finished stages from running again:

• Recovery starts a new run from a failed one and reuses every action that succeeded. See Recover a failed run.
• Caching reuses a task’s output in any later run with the same inputs and code. Set cache="auto" on data prep and training. See Caching.

Related

• Retries and timeouts
• Distributed LLM pretraining: checkpointing and fault tolerance for a multi-node pretraining run.
• Fine-tuning a VLM: multi-node fine-tuning with recovery checkpoints.
• Pre-training
• Post-training