AI Engineering/Phase 19/Lesson 80

Sharded Checkpoint and Atomic Resume

A 70B-parameter training job is paused by a node failure every few hours. The checkpoint format decides whether you lose 30 minutes or 30 hours. A sharded checkpoint writes every rank's shard in parallel and records ownership in a manifest. Resume loads each rank's shard from its own file, reconstructs the state on the same world size, and the optimiser steps as if nothing happened. Atomic write keeps a half-finished checkpoint from poisoning the next resume.

Build 1.5 hr Python

Prerequisites: Phase 19 Track C lessons 42-49

#Learning Objectives

Save a multi-rank checkpoint as a per-rank shard file plus a manifest that records which rank owns what.
Use the atomic write pattern (write to a temp path then rename) so a crash mid-write never produces a half-finished checkpoint.
Resume from the manifest, verifying byte-equal state for both fp16 parameters and the ZeRO optimiser state on every rank.
Defend the manifest schema against the three failure modes: world-size change, shard count mismatch, and partial write.

Full Access lesson

Keep going. The rest of this lesson is one step away.

Unlock all 503 lessons, video, code, quizzes and every future update. Pay in your own currency with Paystack.

See plans Sign in

The text and code of this lesson are based on AI Engineering from Scratch by Rohit Ghumare, MIT licensed.