Warmup Guide — Distributed Training & Data

Zero-to-expert primer for Phase 10: the two scaling problems pretraining forces — parallelizing the compute (DDP, ZeRO, TP/PP) and building the data pipeline (dedup, filtering, mixing) whose quality silently bounds everything.

Table of Contents

Chapter 1: Why Distribute — The Memory and Time Walls

Two independent walls force distribution:

  • Memory: training state per parameter in mixed precision ≈ 16 bytes (FP16 weights 2 + FP16 grads 2 + FP32 master weights 4 + Adam m/v 8 — Phase 05 Ch. 2's fact, itemized). A 7B model: ~112 GB before activations — no single GPU holds it.
  • Time: Chinchilla-optimal 7B training ≈ 6 × 7B × 140B tokens ≈ 6e21 FLOPs — months on one GPU, days on hundreds.

The taxonomy that organizes everything: data parallelism replicates the model and splits the batch; model parallelism splits the model (tensor-wise or layer-wise); ZeRO is data parallelism with the redundancy sharded away. Real runs compose all of them ("3D parallelism"), but each is understandable alone — and the lab needs only DDP-level understanding plus the vocabulary for the rest.

Chapter 2: Data Parallelism and All-Reduce

DDP, mechanically: N ranks each hold a full model replica; each step, each rank runs forward/backward on its own micro-batch; gradients are averaged across ranks before the optimizer step — so all replicas stay bit-identical (same init, same averaged grads, same step).

The communication primitive: ring all-reduce — each of N ranks passes chunks around a ring; total traffic per rank ≈ 2 × (gradient bytes) × (N−1)/N, independent of N — the reason data parallelism scales to thousands of GPUs. The latency-hiding trick that makes DDP fast in practice: overlap — gradients for late layers are ready while early layers still backprop, so all-reduce streams in buckets concurrently with the backward pass (PyTorch DDP's bucketing). When someone's training "doesn't scale," the first suspects are tiny per-rank batches (communication can't hide behind too-little compute) and unbucketed/synchronous reduction.

Effective batch = micro_batch × accumulation × N — the tokens-per-step bookkeeping of Phase 05 Ch. 4, now with a cluster dimension, and the same LR-scaling interaction.

Chapter 3: ZeRO — Sharding the Redundancy

DDP's waste: N replicas of identical optimizer state, gradients, and weights. ZeRO shards them progressively across the data-parallel group:

  • Stage 1: shard optimizer states (the FP32 master + m/v — 12 of the 16 bytes/param). Each rank updates only its shard; updated params are all-gathered.
  • Stage 2: + shard gradients (reduce-scatter instead of all-reduce — each rank keeps only its gradient shard).
  • Stage 3 / FSDP: + shard the parameters themselves — each layer's weights are all-gathered just-in-time for its forward/backward, then freed. Memory per rank approaches (total state)/N; communication rises ~1.5× vs DDP (the extra all-gathers).

The mental model: ZeRO keeps data parallelism's programming model (every rank sees the whole model logically) while paying memory like model parallelism. FSDP is PyTorch's native Stage-3; it's how 7B–70B fine-tuning happens on commodity nodes — and its interaction with LoRA (tiny trainable set → Stage 1–2 suffices) explains why PEFT changed infrastructure requirements, not just science (Phase 06).

Chapter 4: Tensor and Pipeline Parallelism — The 3D Recipe

When a single layer outgrows a GPU, or latency demands splitting within an op:

  • Tensor parallelism (Megatron-style): split weight matrices — column-parallel $W_1$ then row-parallel $W_2$ in the FFN means one all-reduce per block instead of per matmul; attention splits naturally by heads. Communication is per-layer and latency-sensitive → TP lives within a node on NVLink (TP=2–8).
  • Pipeline parallelism: assign contiguous layer ranges to stages; micro-batches stream through. The bubble (stages idle during fill/drain) shrinks with more micro-batches per step (1F1B scheduling); PP spans nodes happily (communication is small boundary activations).
  • The 3D rule of thumb (Megatron/LLaMA-style): TP within node, PP across nodes, ZeRO/DP across the remainder — and sequence/context parallelism joins at very long context. You won't run 3D in this track's labs; you need the layout to read modern training reports and to reason about where a given failure (slow step, OOM, divergence on one rank) localizes.

Chapter 5: Activation Memory and Checkpointing

The other memory consumer — often the binding one at long sequence lengths: activations stored for backward scale ≈ batch × seq × hidden × layers (with the attention term worse pre-FlashAttention). Activation (gradient) checkpointing: store only block-boundary activations; recompute the interior during backward — ~√-layers memory at ~33% extra forward compute. The recompute-vs-store trade is the same one FlashAttention's backward (model-accuracy Phase 08) and Mamba's scan make: compute is cheaper than memory bandwidth/capacity is the era's recurring exchange rate. Combined with mixed precision (Phase 05 Ch. 5) and ZeRO, this is the standard "fit the run" toolkit, in the order you should reach for it: precision → ZeRO stage → checkpointing → parallelism redesign.

Chapter 6: The Data Pipeline — Where Model Quality Is Decided

The unglamorous half that determines more eval variance than most architecture choices. The standard pretraining pipeline (your lab builds a small one end-to-end):

  1. Acquisition: web crawls (Common Crawl WARC/WET), code (licensing-filtered), curated sources (books, wiki, papers).
  2. Extraction: HTML → text (boilerplate removal — trafilatura-class tooling); quality is extraction-dependent before any filtering.
  3. Language ID (fastText-class), then quality filtering: heuristic rules (Gopher rules: symbol ratios, repetition, doc length), model-based scoring (perplexity against a clean reference, or a quality classifier — beware: classifiers encode taste, and FineWeb-class work shows the choice moves downstream evals substantially).
  4. Deduplication and decontamination (Ch. 7).
  5. Tokenize and pack (Phase 01's tokenizer): documents concatenated with EOS separators into fixed-length sequences; shuffled at document level (sequence-level shuffling after packing leaks cross-doc context); sharded for the dataloader with deterministic resumability (a crashed 3-week run must resume mid-epoch exactly).

The engineering character: this is a data engineering system — content-addressed shards, manifest files, versioned configs (the model-accuracy capstone's ledger ethic) — because "which data trained this checkpoint" must be answerable forever.

Chapter 7: Deduplication and Decontamination

  • Why dedup matters: duplicated text wastes compute, amplifies memorization (verbatim regurgitation rates track duplication counts), and skews evals. Web crawls are massively duplicated (boilerplate, mirrors, spam).
  • Exact dedup: hash documents (or shingled spans for near-verbatim). Cheap, catches mirrors.
  • Near-dedup — MinHash + LSH (the lab's algorithmic centerpiece): shingle each doc into n-grams; a MinHash signature of k permutation-minima estimates Jaccard similarity (P[minhash collision] = J(A,B) — the elegant identity at the core); LSH banding (b bands × r rows: match probability $1-(1-J^r)^b$, an S-curve you tune to a similarity threshold) finds candidate pairs without O(N²) comparison; candidates verify exactly, then cluster and keep one representative.
  • Decontamination: the same machinery aimed at your eval sets (Phase 08 Ch. 3's contamination, prevented at the source): n-gram overlap scans of training shards against benchmark items, with the honest caveat that paraphrase contamination survives n-gram screens.

Chapter 8: Mixing, Curricula, and Epochs

  • Mixture weights: pretraining corpora are blends (web/code/books/reference) with weights set by ablation, not principle alone — code fractions notably affect reasoning; multilingual fractions trade English headroom for coverage. Mixtures are the lever frontier labs actually tune (DoReMi-style methods learn weights).
  • Epochs: the modern default is ~1 epoch over a huge corpus; repeating data has measurable diminishing returns (4+ epochs ≈ noticeably degraded vs fresh tokens), with high-quality sources tolerating a few repeats (the data-constrained scaling work) — relevant now that frontier runs approach the supply of good tokens.
  • Curriculum: ordering effects are mostly weak at pretraining scale, with one robust exception — long-context extension as a final phase (train short, finish with long sequences at adjusted RoPE) and quality-upweighted "annealing" data at the end of training (several modern recipes), both cheap to know about and to look for in tech reports.

Lab Walkthrough Guidance

Lab 02 — Pretraining Data Pipeline:

  1. Build stage-by-stage with counters at every stage (docs in / docs out / bytes / reasons-rejected) — the funnel report is the deliverable's spine; a pipeline without it is undebuggable.
  2. Implement Gopher-style heuristic filters as data (a rules list), then a perplexity filter; inspect random samples of what each filter rejects — every filter encodes a bias, and looking at rejects is how you find filters eating poetry or code.
  3. MinHash: verify the collision-probability identity empirically on synthetic pairs of known Jaccard before scaling; then tune LSH bands to your threshold using the S-curve; report precision/recall of near-dup detection on a labeled sample.
  4. Decontaminate against a small benchmark slice; report hits found (planted ones — seed your corpus with eval items to test the screen).
  5. Tokenize and pack with document shuffling + deterministic resumable iteration; prove resumability (kill at step k, resume, identical batch sequence).
  6. Train your Phase 05 nanoGPT on filtered vs unfiltered slices at matched token budget — the eval delta is the phase's thesis, produced by your own hands.

Success Criteria

You are ready for Phase 11 when you can, from memory:

  1. Itemize the 16 bytes/param and compute total training-state memory for any model.
  2. Explain ring all-reduce's N-independence and DDP's overlap trick; name the two "doesn't scale" suspects.
  3. State what each ZeRO stage shards, the memory/communication trade, and why FSDP+LoRA need less.
  4. Place TP within-node and PP across-node with reasons; define the pipeline bubble.
  5. Derive the checkpointing trade (√-layers memory, +33% compute) and the era's compute-for-memory exchange-rate theme.
  6. Walk the data pipeline's six stages, MinHash/LSH's collision identity and S-curve, and the dedup→memorization link.

Interview Q&A

Q: 8 GPUs train only 5× faster than 1. Diagnose. Step-time decomposition first: if compute per rank is unchanged but steps are slower — communication: per-rank batch too small for overlap to hide all-reduce (raise micro-batch or accumulation), interconnect (PCIe vs NVLink — measure bus bandwidth), unbucketed sync, or a straggler rank (one slow GPU/thermals gates the ring — per-rank timing exposes it). If compute changed too: dataloader contention feeding 8 ranks (shared storage), or accidental Python-level serialization. The skill: separate scale-up losses into communication, straggler, and input-pipeline buckets with per-rank timers before changing anything.

Q: Why does ZeRO-3/FSDP exist when tensor parallelism already splits the model? Different problems: TP splits compute within ops to cut per-token latency and fit huge layers — but demands NVLink-class interconnect, code-intrusive sharding, and synchronizes per layer. ZeRO-3 shards state while keeping pure-data-parallel semantics — works over ordinary interconnects, no model rewrite, scales with the DP group. Training a 13B on 8 commodity GPUs: FSDP. Serving a 70B at low latency or training where single layers don't fit: TP (then both, composed, at frontier scale).

Q: How would you prove a model's regurgitation risk came from data duplication? Correlate: index the training corpus (suffix array / n-gram index), sample model generations, find verbatim training matches, and plot memorization rate against each matched sequence's duplication count in the corpus — the published result (and the one you'd reproduce) shows monotonic scaling. Then the fix: near-dedup the corpus, retrain or continue-train, re-measure. The methodology — measure, intervene, re-measure, with the index doing the work — is the answer.

References

  • Rajbhandari et al., ZeRO: Memory Optimizations Toward Training Trillion Parameter Models (2019) — arXiv:1910.02054
  • Shoeybi et al., Megatron-LM: Tensor Parallelism (2019) — arXiv:1909.08053
  • Huang et al., GPipe (2018) — arXiv:1811.06965 and Narayanan et al., Efficient Large-Scale Training (1F1B/3D) (2021) — arXiv:2104.04473
  • PyTorch FSDP docs and DDP internals notes
  • Rae et al., Gopher (2021) — arXiv:2112.11446 — appendix A: the filtering rules
  • Lee et al., Deduplicating Training Data Makes Language Models Better (2021) — arXiv:2107.06499
  • Carlini et al., Quantifying Memorization Across Neural Language Models (2022) — arXiv:2202.07646
  • Penedo et al., FineWeb (2024) — arXiv:2406.17557 — the modern open data-pipeline writeup
  • Muennighoff et al., Scaling Data-Constrained Language Models (2023) — arXiv:2305.16264 — the epochs question
  • Broder, On the Resemblance and Containment of Documents (1997) — MinHash's origin