When the machine takes it back

Every rank runs the same script, which chapter eight establishes, and every rank therefore reaches the load line on its own clock. Rank zero is usually the one that writes. If rank one reads before that write has landed, and the previous checkpoint is still sitting on disk, it loads the previous checkpoint.

Two ranks, two step numbers, no exception anywhere. Rank one resumes at step 10 and rank zero at step 20, and from that point they are training different runs while believing they are training one. Delete the old file and the failure changes shape rather than going away: rank one gets a FileNotFoundError on the first attempt, which is louder but no more correct.

dist.barrier() after the save is the fix, and the second half of the capture shows both ranks at step 20. The same barrier belongs on the other side of a save as well, before rank zero starts writing, if any rank might still be reading the file it is about to replace. Combine that with the temp-and-rename from lesson one and a reader can never observe a partial file at all.

import time

import torch
import torch.distributed as dist
import torch.multiprocessing as mp


def worker(rank, world, wait):
    dist.init_process_group("gloo", init_method="tcp://127.0.0.1:29530",
                            rank=rank, world_size=world)
    if rank == 0:
        time.sleep(0.5)                      # writing a real checkpoint is not instant
        torch.save({"step": 20}, "shared.pt")
    if wait:
        dist.barrier()
    print(f"rank {rank} resumes at step {torch.load('shared.pt')['step']}")
    dist.destroy_process_group()


if __name__ == "__main__":
    for wait in (False, True):
        torch.save({"step": 10}, "shared.pt")   # the previous checkpoint, still on disk
        mp.spawn(worker, args=(2, wait), nprocs=2, join=True)

# rank 1 resumes at step 10
# rank 0 resumes at step 20
# rank 1 resumes at step 20
# rank 0 resumes at step 20

A collective is agreement about shapes and dtypes, not about meaning. Two ranks that resumed from different checkpoints all-reduce their step numbers, 10 and 20, and get 30 back. The operation is correct. Nothing in c10d has any opinion about whether the two numbers should have been the same.

The check has to be yours, and it costs two collectives at startup. Reduce the restored step with MIN, reduce it again with MAX, and refuse to run when the two differ; here they come back 10 and 20 on both ranks, which is the disagreement made visible on the line before the training loop. Do the same for anything else the ranks have to share, the epoch counter among them.

One mechanism unifies part of the state on its own, which is worth knowing because of what it leaves behind. Constructing DistributedDataParallel broadcasts module state from rank zero, and the distributed unit measures that constructor broadcast directly. For a resume it means the weights end up consistent whatever each rank loaded, while the optimizer state, which DDP never touches, stays exactly as inconsistent as the files were.

The weights agree because DDP made them agree. The step counters and the optimizer moments agree only if you checked.

import torch
import torch.distributed as dist
import torch.multiprocessing as mp


def worker(rank, world):
    dist.init_process_group("gloo", init_method="tcp://127.0.0.1:29532",
                            rank=rank, world_size=world)
    step = torch.tensor([10.0 if rank == 0 else 20.0])   # two ranks, two checkpoints

    total = step.clone()
    dist.all_reduce(total)                               # the collective is happy
    print(f"rank {rank}: all_reduce returned {total.item()}")

    lo, hi = step.clone(), step.clone()                  # the check that is not free
    dist.all_reduce(lo, op=dist.ReduceOp.MIN)
    dist.all_reduce(hi, op=dist.ReduceOp.MAX)
    print(f"rank {rank}: step min {lo.item()} max {hi.item()} agreed {lo.item() == hi.item()}")
    dist.destroy_process_group()


if __name__ == "__main__":
    mp.spawn(worker, args=(2,), nprocs=2, join=True)

# rank 1: all_reduce returned 30.0
# rank 1: step min 10.0 max 20.0 agreed False
# rank 0: all_reduce returned 30.0
# rank 0: step min 10.0 max 20.0 agreed False

Preemption is not a clean exit. A rank disappears in the middle of a collective the others have already entered, and those others do not fail, they wait. The distributed unit captures what that looks like when the timeout finally fires, a transport-layer error naming your own operation and the milliseconds it waited, and it is the same error whether the peer was preempted, killed by the memory manager, or stuck.

What belongs to a resumable loop is the number in front of it. default_pg_timeout is thirty minutes, so an untuned job spends half an hour dead before anything says so, and that half hour is charged. Pass a timeout to init_process_group that matches how long you are willing to lose, and treat the error it eventually raises as the signal to restart rather than as a bug to fix.

The training script is the wrong place to handle it either way. The error names your operation, not the missing peer, and the surviving ranks cannot form a new group by themselves. Restarting is the launcher's job, which is what the next section is about.

The elastic launcher's contract is short and blunt, and reading it in its own words settles most design questions about a preemptible loop. For a job with n workers, if k of them fail, all workers are stopped and restarted, up to max_restarts. Not the failed ones. All of them.

The training-script guidance follows from that. Workers are restarted with the same program arguments, so progress is lost back to the most recent checkpoint, and the script is expected to carry load_checkpoint and save_checkpoint of its own. Nothing in the launcher preserves anything held in a process's memory across a restart, because there is no process left to preserve it in.

Elastic membership makes the shape of the job part of the state. --nnodes=1:4 gives a minimum and a maximum, and when a node departs or arrives the docs describe the same sequence: existing workers stopped, a new worker group formed, all workers started with a new RANK and WORLD_SIZE. The launcher's own warning is that RANK is not stable between restarts and that no assumption about the correlation between RANK and LOCAL_RANK should be hard-coded.

For a checkpoint that means two rules. Nothing in the file may be keyed by rank position if the run can reshard, and anything that scales with world size, a per-rank batch count or a data shard assignment, has to be recomputed after the restart rather than restored from the file.

A single-file torch.save from rank zero works while the whole model fits in one process. Shard the parameters across ranks, with FSDP or a tensor-parallel mesh, and there is no rank holding the whole thing to write it.

torch.distributed.checkpoint is the answer in the box. It writes a directory rather than a file, with at least one file per rank, saving and loading from every rank in parallel; the load is in-place, so the model allocates its own storage and the checkpoint reads into it. The property that matters for a preemptible job is load-time resharding: a checkpoint saved from one cluster topology can be loaded into another, which is what an elastic restart with a different world size needs.

The pairing that makes it usable across strategies is get_state_dict and set_state_dict, which normalize a model and optimizer state dict into fully qualified names rather than the positional ids chapter four's first lesson takes apart. That normalization is what lets a checkpoint written under one parallelism be read under another.

None of this is measured here. The environment on this machine cannot run the coordinating collective the saver needs, so the sharded save and the resharding load are Colab-pending rather than captured, and everything above is the upstream documentation's claim rather than this page's measurement.

The capstone above this arc asks for a run on real hardware, a checkpoint holding the whole state, a deliberate kill, a resume, and an overlay proving the curve continued rather than restarted. Every piece of that has now been taken apart on a machine with no accelerator in it, which is the point: none of the difficulty was ever in the chip.

What the chip adds is the device RNG from lesson one, the bridge's own save path, and a scheduler that can take the machine back mid-epoch. LAB·P4 runs the whole sequence over torch_xla and publishes its reference numbers; this arc is what those numbers mean and how to know when yours are lying.

The half-page provenance report is the last piece and the one people skip. Which chip, which dtype, which shapes, which bridge, and the comparison you actually ran, stated as the comparison rather than as a claim about it. A resumed loss list equal to the uninterrupted one is a result. A curve that looks right is a screenshot.