Checkpointing and Resuming Pipelines

Model training jobs can be interrupted for any number of reasons: there could be a hardware or networking failure, the job could be preempted by a higher priority run that needed its resources, or the available time on the training cluster might expire or become too expensive to continue training right now. For all of these reasons (and for all of the other reasons that I can’t think of right now), model training systems all include a mechanism for checkpointing the state of a training run, saving it away, and then restoring that state later so as to be able to pick up the run from exactly where it left off. For the data pipeline, this means keeping track of which samples the training loop has seen and determining where to pick things up again when the training run is restored.

We will now go over how to checkpoint and restore the state of a Pipeline during a training run. We will also explain how Zephon supports elastic resumption, which allows you to continue the same data curriculum on a different number of data-parallel workers.

Checkpointing A Pipeline. The Pipeline.checkpoint() method returns a dictionary that describes the current progress of a Pipeline object. Your training code is responsible for persisting that dictionary with the rest of the state for your training run (the model weights, the state of the optimizer, etc.) Checkpoints should be taken at a completed training step boundary once the batches associated with that step have been incorporated into the model:

examples/guide/pipelines/checkpoint_pipeline.py

"""Save the Pipeline's progress with the rest of the training state."""

from zephon import Pipeline

pipeline = (
    Pipeline(work_source)
    .tokenize(tokenizer_id="gpt2", field="text")
    .pack_flat(max_length=4097, algorithm="wrap", pad_token_id=0)
    .batch(microbatch_size=32)
)

for step, sample_batch in enumerate(pipeline):
    train_step(model, sample_batch.to_training(return_labels=True))
    if step % 1000 == 0:
        # One object, written together: a pipeline state that describes a
        # different step than the model's will silently skip or repeat data.
        save(
            {
                "model": model.state_dict(),
                "optimizer": optimizer.state_dict(),
                "data": pipeline.checkpoint(),
            }
        )

If the training loop accumulates gradients across several microbatches, save the checkpoint after a complete optimizer step, once all batches returned by the Pipeline have been incorporated into that step. The pipeline, model, and optimizer states must describe the same point in training; otherwise, resuming can skip data or train on it twice.

You do not need to consider Zephon’s internal prefetch buffers for checkpointing. Zephon tracks which records it has delivered to the training loop, separately from the records it has prepared ahead of time. However, it cannot tell whether your training code has actually used a batch it returned. If your code reads ahead from the Pipeline, make sure those batches have been incorporated into the completed step before checkpointing.

Aligning Optimizer Steps With Lanes. To preserve the same global training batches when resuming on a different topology (elastic determinism), there is one more requirement: every optimizer step must consume the same number of records from every lane. Recall from Distributed Training that each data-parallel group takes turns reading batches from the lanes it owns. An optimizer step therefore needs to include a whole number of these rounds.

For a run with canonical_replicas=L lanes and dp_degree=G data-parallel groups, this means:

• L must be divisible by G, so every group owns exactly L / G lanes.

• The number of training microbatches consumed by each group per optimizer step must be a multiple of L / G, so every owned lane gets equally many turns.

You can change the training microbatch size between runs, but the new configuration must still satisfy these rules. Equivalently, if the global batch contains N records, each lane contributes N / L records per step, and the microbatch size must divide that contribution evenly.

You can check these conditions statically. For example, if you check out our TorchTitan integration, you can see how it checks the divisibility conditions during training initialization and rejects configurations that violate them before training begins.

Restoring A Training Run. To resume a training run, use the same Dataset, WorkSource, and Pipeline definitions that were used for the original training run, load the saved training state that was generated by the Pipeline.checkpoint() method, and then call Pipeline.restore() before you begin iterating the Pipeline object:

examples/guide/pipelines/restore_pipeline.py

"""Resume a run from a saved checkpoint."""

from zephon import Pipeline

# The same datasets, work source and operators as the original run, down to
# the tokenizer and every transform parameter.
pipeline = (
    Pipeline(work_source)
    .tokenize(tokenizer_id="gpt2", field="text")
    .pack_flat(max_length=4097, algorithm="wrap", pad_token_id=0)
    .batch(microbatch_size=32)
)

pipeline.restore(checkpoint["data"])  # before the first iteration, not after

for sample_batch in pipeline:
    train_step(model, sample_batch.to_training(return_labels=True))

It is critical that you keep the underlying data and operators identical for the resumed training run, including settings like the choice of tokenizer or the parameters for any transformations, shuffles, mixtures, or packing operations. Custom transformation functions defined via a map_transform operator must also be deterministic; it’s obviously not possible to guarantee that a Pipeline can exactly reproduce a non-deterministic decision that is made during pipeline execution.

The state dictionary returned by Pipeline.checkpoint records which work chunks the run had finished and which were still in progress, rather than the internal state of each operator. When you call Pipeline.restore with the state information, Zephon rebuilds the state of the pipeline, including any operators that buffer or re-order records, by re-processing a limited number of recent work chunks instead of the entire training curriculum. For pipelines made up entirely of order-preserving operators, that means only the chunks that were in flight when the checkpoint was taken. For pipelines that re-order or regroup records, like those that need to shuffle or pack, it means the chunks that were processed since the most recent flush, so the flush_every_k_chunks setting also limits how much work a restore operation has to redo. The How Pipelines Run page describes this mechanism and the tradeoffs involved with it in more detail.

Checkpointing In Distributed Training Jobs. In distributed training, each rank sees only a portion of the overall job’s data loading progress, so Zephon combines the progress of every participating Pipeline into a single checkpoint. For this to work, each Pipeline needs a coordination identity along with a shared aggregate_dir and run_id, as described in Distributed Training. The aggregation directory is only used while a checkpoint is being assembled, and you don’t need to keep the directory around in order to resume a run from the checkpoint later.

Every Pipeline that is counted in the world_size must call checkpoint() at the same training step, even though only one of them (the Pipeline that has global_rank=0) will save the result. If a single rank calls the checkpoint() method by itself, it will wait for the others until the aggregation timeout expires, and then fail. Training Integrations shows how to set things up correctly for the TorchTitan and Megatron-LM distributed training frameworks.

Elastic Resumption. Suppose that a training run starts on eight data-parallel ranks, but after preemption, only four of the ranks are available. If the run was planned for elasticity (see the discussion on this in Planning for Elasticity), you can resume it on the remaining four ranks without losing your place in the curriculum. Build the Pipeline with the new topology (i.e., the new dp_degree, dp_group_id, world_size, and global_rank values) but with the same canonical_replicas value that the original run used, and then restore the checkpoint as usual:

examples/guide/pipelines/elastic_resume.py

"""Resume on four data-parallel groups a run that started on eight."""

from torch.distributed.device_mesh import init_device_mesh

from zephon import Pipeline

mesh = init_device_mesh("cuda", (4, 2), mesh_dim_names=("dp", "tp"))
dp_mesh = mesh["dp"]

# The topology is new; canonical_replicas is not. Zephon refuses a checkpoint
# whose lane count changed, and hands each of the four groups two lanes.
pipeline = (
    Pipeline(work_source)
    .batch(microbatch_size=32)
    .options(
        canonical_replicas=8,
        dp_degree=dp_mesh.size(),
        dp_group_id=dp_mesh.get_local_rank(),
        world_size=dp_mesh.size(),
        global_rank=dp_mesh.get_local_rank(),
        aggregate_dir="s3://my-bucket/runs/aggregate",
        run_id="pretrain-7b-2026-03-14",
    )
)

pipeline.restore(checkpoint["data"])

Pipeline.restore() reassigns the lanes of logical data to the new ranks. With canonical_replicas=8, each of the four ranks now processes two lanes apiece. No samples are skipped or processed twice, and if the run follows the batches-per-step rule described in Planning for Elasticity, each global training step contains exactly the same samples it would have on the original eight ranks. Zephon refuses to restore a checkpoint if the canonical_replicas value has changed.