End-to-End Distributed Training

A capstone checks that the pieces fit together. Lesson 76 implemented collectives. Lesson 77 wrapped them into DDP. Lesson 78 sharded optimiser state with reduce_scatter. Lesson 79 analysed pipeline parallelism. Lesson 80 saved a sharded checkpoint. Each one had its own test. A real training run uses all of these primitives at once, and a wrong composition shows up as a diverging loss, a checkpoint that refuses to resume, or per-rank memory that grows when it should shrink.

The end-to-end demo is built around four invariants, but it only enforces one of them. (a) The loss table prints rank 0's loss for each of the 20 steps for you to read. Nothing checks that it falls, and it should not: the corpus is uniform random tokens, so there is nothing to learn and the demo tests wiring, not convergence. (b) Every rank should hold the same parameter norm. The script measures it once, after the last step, and prints the per-rank values with their drift. (c) Per-rank optimiser memory should follow the ZeRO-1 formula of 12P/N bytes (fp32 master copy, m and v, with P padded to a multiple of N). The script prints the measured shard bytes next to that figure and does not assert they match. (d) The checkpoint at step 10 should reload byte-equal. This is the only gate: if the reloaded optimiser master shards differ from the snapshot taken at save time, the script prints RESUME FAILED and returns 1. The demo is a single command that runs 20 steps and exits 0 when the resume check passes.

The mini GPT

The model is 2 transformer blocks with embed dim 32, 4 attention heads, vocab 64, sequence length 16 and batch 4. It has 29,824 parameters (the script prints the count). That is large enough to exercise the wiring decisions that matter: multi-head attention runs the standard masked path, LayerNorm has weights to sync, and the LM head is a separate linear projection back to the vocab. It is small enough that 20 steps on 4 CPU ranks finish in well under a minute (about 36 seconds end to end here, including the torch import and process spawn).

The training data is torch.randint over the 64-token vocab, generated once from a seeded generator and sliced so each rank sees a different block at every step. The tokens are uniform random, so the next token cannot be predicted from the context. With random tokens there is nothing to learn, so expect the loss to stay near the uniform-guess value of ln 64 (about 4.16). In our run rank 0 printed values between 4.11 and 4.46 with no downward trend.

The composition rules

The loop knows nothing about reduce_scatter or rendezvous files. The ZeRO and checkpoint modules expose narrow interfaces, and the loop composes them.

Why a tiny GPT and not just an MLP

The MLP from lesson 77 was enough to verify gradient sync. A tiny GPT adds three things. The LM head over the vocab is separate (untied here for clarity, where a full GPT typically ties it to the token embedding). The loss is softmax plus cross-entropy, which has more numerical edge cases than MSE. And the forward pass is asymmetric: embeddings, then attention, then an MLP per layer. An MLP capstone would leave open whether the composition handles LayerNorm and the embedding layer's grad shape correctly.

A self-terminating run

The loop runs a fixed 20 steps and exits. There is no while True, no human intervention, and no resume from external state. You can leave the demo running unattended and find a complete log when it finishes, which is a decent sign the system is wired correctly. If a piece deadlocks, the demo never returns and the test rig catches it.

code/main.py implements:

• MiniGPT: 2-layer transformer with masked self-attention and a separate LM head.
• make_corpus(seed, total_tokens): seeded uniform random token ids from torch.randint, sliced into next-token input and target pairs. There is no pattern to learn, so the loss exercises the wiring and not the model.
• _train_worker: spawned per rank. It broadcasts the init params, runs the loop, calls the ZeRO step, and writes the sharded checkpoint at step 10.
• verify_resume: after the main run, reloads the step-10 checkpoint in-process and returns a bool: whether each rank's saved optimiser master shard equals the in-memory snapshot (torch.equal). It does not compare the model state or the m and v shards. main prints the result.
• main: runs the whole demo and prints the loss table, the memory profile, and the verification result, and returns 1 if the check fails.

Run it:

python3 code/main.py

Output: a 20-row loss table (rank 0 only), the final parameter norm of each rank with the drift between them, a 4-row per-rank optimiser shard size next to the expected size, the checkpoint directory path (not the manifest contents), and a "RESUME VERIFIED" line on success.

Three patterns carry the composition over to real runs.

Checkpoint every K minutes, not every K steps. Step time varies with seq length and microbatch count, so a 10-minute cadence catches the same amount of compute regardless of model size. The demo checkpoints by step count for simplicity. Production runs use wall-clock time.

Detect divergence early. Production runs add a NaN guard after backward and a loss-spike detector. If loss jumps by more than 2x in one step, they roll back to the previous checkpoint instead of letting the optimiser march into a degenerate state. The demo has no such guard. Its random-token loss should stay near ln 64, so there is nothing for a detector to catch. Exercise 5 adds one.

Aggregate the memory profile across ranks. Per-rank memory differs in real runs (the rank holding the largest pipeline stage carries more activations). Production logs the max across ranks plus the mean. The demo prints per-rank optimiser shard bytes next to the expected figure so you can compare them with the formula.

Where this shape shows up in production:

• DeepSpeed. Combines DDP, ZeRO, pipeline and activation checkpointing under one config. The composition here is the DeepSpeed shape in miniature.
• PyTorch FSDP. The native equivalent. FullyShardedDataParallel with ShardingStrategy.SHARD_GRAD_OP is ZeRO-2.
• NeMo and Megatron-LM. Add tensor parallel for the very largest models. The rest of the composition has the same shape.

The track ends here. The 6 lessons together make up the distributed-training subsystem a team would build before adopting DeepSpeed. The abstraction has been run against gloo, and the resume check guards the checkpoint path. Phase 17 (infrastructure and production) is where you take this to a real cluster.

1. Add a tensor-parallel split of the attention head and verify the loss matches the single-rank baseline. Use two ranks: half the heads per rank, with an allreduce of the attention output.
2. Add gradient accumulation across 4 microbatches and show that the gradient equals the gradient of one big batch.
3. Add a resume-from-step-10 path that continues training to step 20 and produces the same final loss as the original run.
4. Add a metrics export (loss, grad norm, step time) to JSONL so the run can be visualised after the fact.
5. Add a NaN guard that rolls back to the previous checkpoint on a loss spike, and force a spike with a one-step LR multiplier to exercise the rollback.

• DeepSpeed end-to-end training tutorial
• PyTorch FSDP advanced tutorial
• Megatron-LM training script reference
• Phase 19 Lessons 76-80 - each piece this lesson composes
• Phase 17 - moving the composition to a real cluster