AI Engineering/Phase 19/Lesson 48

Distributed Data Parallel and FSDP from Scratch

Multi-rank training is two collectives and one rule. Broadcast the parameters at startup, average the gradients after backward, never let the ranks disagree about what step they are on.

Build 1.5 hr Python

Prerequisites: Phase 19 lessons 42 to 45

#Learning Objectives

Bring up a process group across N ranks with the gloo backend, no special hardware.
Implement a minimal DDP wrapper that broadcasts parameters at construction and all-reduces gradients after backward.
Prove that the all-reduce of per-rank gradients matches a single-process gradient on the concatenated input.
Sketch FSDP parameter sharding: each rank holds a slice, the full tensor is gathered for the forward pass and dropped after.

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.