# 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: ```{literalinclude} ../../../examples/guide/pipelines/checkpoint_pipeline.py :language: python :caption: examples/guide/pipelines/checkpoint_pipeline.py ``` 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](distributed_training.md) 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](../training_integrations.md), 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: ```{literalinclude} ../../../examples/guide/pipelines/restore_pipeline.py :language: python :caption: examples/guide/pipelines/restore_pipeline.py ``` 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](how_pipelines_run.md) 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](distributed_training.md). 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](../training_integrations.md) shows how to set things up correctly for the TorchTitan and Megatron-LM distributed training frameworks. (elastic-resumption)= **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 {ref}`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: ```{literalinclude} ../../../examples/guide/pipelines/elastic_resume.py :language: python :caption: examples/guide/pipelines/elastic_resume.py ``` `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 {ref}`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.