Warmup Guide — Distributed Training & Model Parallelism

Zero-to-senior primer for Phase 10. We start from a single hard fact — one GPU cannot hold a frontier model (Phase 00: training costs ~16 bytes/param, so 70B ≈ 1.1 TB) — and end with the full mental model of a training cluster: the collective-communication primitives and the bandwidth-optimal ring all-reduce; the four ways to cut a model (data, tensor, pipeline, ZeRO/FSDP); the communication each one pays; the pipeline bubble; and how they compose into 3D parallelism mapped onto a real topology. By the end you can size a training run on a cluster the way Phase 00 sized inference on a single card.

Table of Contents

• Chapter 1: Why One GPU Isn't Enough
• Chapter 2: Collective Communication & the Ring All-Reduce
• Chapter 3: Data Parallelism — Replicate, All-Reduce, Average
• Chapter 4: Tensor (Model) Parallelism — Sharding a Layer
• Chapter 5: Pipeline Parallelism & the Bubble
• Chapter 6: ZeRO / FSDP — Sharding the State, Not the Compute
• Chapter 7: 3D Parallelism & Mapping to a Cluster
• Chapter 8: Overlap, Activation Checkpointing & the Frameworks
• Lab Walkthrough Guidance
• Success Criteria
• Interview Q&A
• References

Chapter 1: Why One GPU Isn't Enough

From zero. Recall the Phase 00 training-memory rule: to train a model with N parameters in mixed precision with Adam you hold, per parameter, an fp16 weight (2 B), an fp16 gradient (2 B), an fp32 master copy (4 B), and Adam's two fp32 moments (4 + 4 B) — about

$$ \text{train memory} \approx 16 \times N \quad \text{bytes (before activations).} $$

A 7B model is ~112 GB; a 70B model is ~1.1 TB. The biggest single accelerator today has 80–192 GB of HBM. The model state alone does not fit, and that is before the activations saved for the backward pass, which scale with batch × sequence length. So the question that defines this entire phase is mechanical, not philosophical: the bytes are bigger than the box — where do you put the rest?

Three places the bytes can go. There are only four answers, and they map to what you split:

The senior's reframe. A cluster is not a faster GPU; it is a network of GPUs, and the moment you split anything, the bottleneck moves from FLOPs (Phase 00) to communication bandwidth between the cards. NVLink inside a node is ~900 GB/s (H100); InfiniBand between nodes is ~50–400 GB/s; this hierarchy decides which split you can afford where. Internalize that the interconnect is a budget exactly the way HBM bandwidth was a budget for decode.

Common misconception. "Add more GPUs and training gets proportionally faster." Only until the communication to keep them in sync dominates. A DP job can scale to 8 GPUs beautifully and then go slower at 64 because the all-reduce now costs more than the compute it overlaps with. Scaling is a communication problem dressed as a compute problem.

Chapter 2: Collective Communication & the Ring All-Reduce

What a collective is. When N GPUs ("ranks") must agree on a value, they call a collective operation — every rank participates, no central server. The four you must know:

• broadcast — one rank's tensor copied to all ranks.
• all-gather — each rank has a chunk; afterward every rank has all chunks concatenated.
• reduce-scatter — each rank contributes a full tensor; afterward each rank owns the reduced (e.g. summed) value of one chunk.
• all-reduce — each rank contributes a full tensor; afterward every rank holds the element-wise reduction (sum) of all of them. This is the one data parallelism lives on.

The naive way (and why it's bad). You could send every rank's gradient to rank 0, sum there, and broadcast back. Rank 0 then receives (N-1)·tensor bytes — a hotspot that gets worse with more ranks. Bandwidth wasted, and rank 0 is a bottleneck.

The ring all-reduce. Arrange the ranks in a logical ring (0 → 1 → 2 → … → N-1 → 0) and split each rank's tensor into N equal chunks. The op is two phases:

1. Reduce-scatter (N-1 steps). At step k, rank r sends one chunk to its neighbor r+1, which adds it into its own copy. The chunk indices are staggered so that after N-1 steps each rank owns the fully-summed value of exactly one chunk.
2. All-gather (N-1 steps). Now rotate those finished chunks around the ring so every rank ends up with every summed chunk — reassemble and you have the full all-reduce.

  Reduce-scatter (N-1 steps)            All-gather (N-1 steps)
  each rank ends owning 1 summed chunk  rotate finished chunks to everyone
  rank0 ──c0──▶ rank1 ──c1──▶ rank2     rank0 ◀──Σc──▶ rank1 ◀──Σc──▶ rank2
     ▲                          │          every rank now holds all Σ chunks
     └────────── ring ──────────┘          → reassemble → full sum on all ranks

Why it is bandwidth-optimal. Across both phases each rank sends and receives N-1 chunks of size tensor/N, so the bytes moved per rank are

$$ \text{bytes/rank} = 2 \times \frac{\text{tensor}}{N} \times (N-1) = 2,\text{tensor}\cdot\frac{N-1}{N}, $$

and the total step count is

$$ \boxed{;\text{steps} = 2(N-1);} $$

The remarkable part: as N → ∞, bytes/rank → 2·tensor — constant, not growing with the cluster. No hotspot, no O(N) blowup. This is the single reason data parallelism scales to thousands of GPUs, and it is exactly what NCCL's ring algorithm implements behind every loss.backward().

The soul of the lab. Whatever the schedule, the result of all-reduce is the plain element-wise sum of every rank's tensor. The lab's ring_all_reduce must reproduce that sum exactly (the soul test) while moving the data the way the ring does — proving you understand both the answer and the mechanism.

Common misconception. "All-reduce is O(N) so it doesn't scale." The latency term grows with N (more hops), but the bandwidth term — the bytes each link carries — is constant. For the large gradient tensors of training, you are bandwidth-bound, so the ring is near-optimal; NCCL switches to a tree algorithm only for small, latency-bound messages.

Chapter 3: Data Parallelism — Replicate, All-Reduce, Average

The mechanism. Data parallelism is the simplest and most-used strategy. Put a full copy of the model on each of N GPUs, give each a different slice ("micro-batch") of the global batch, and:

1. each rank runs forward + backward on its own data → its own gradient g_r;
2. all-reduce the gradients (sum across ranks);
3. divide by N → the mean gradient;
4. every rank applies the same mean gradient → the replicas stay identical.

$$ g_{\text{update}} = \frac{1}{N}\sum_{r=1}^{N} g_r. $$

Why the average is the whole point. Step 3 is not cosmetic. The mean gradient over N micro-batches is mathematically identical to the gradient of one big batch made of all of them (for a mean loss). That equivalence is the entire correctness guarantee of data parallelism: N GPUs behave exactly like a single GPU with an N× larger batch. The lab's data_parallel_gradient_average asserts this against the serial mean — drop or double-count one rank's gradient and the equivalence breaks, the replicas diverge, and your training is silently wrong.

The cost. Every step you all-reduce the full gradient, which is the same size as the model (2N bytes in fp16). For a 7B model that's ~14 GB all-reduced every step. The ring keeps it affordable (each rank moves ~28 GB total, constant in cluster size), and real frameworks overlap that communication with the backward pass: as soon as a layer's gradient is ready, its all-reduce launches while the next layer is still computing (PyTorch DDP buckets gradients precisely to enable this — Chapter 8).

When DP runs out. Data parallelism replicates the model, so it does nothing for the memory problem of Chapter 1 — if the model + optimizer state doesn't fit on one GPU, DP can't help. That is the wall that motivates tensor parallelism, pipeline parallelism, and ZeRO.

Common misconception. "Bigger global batch is free quality." A larger effective batch (more DP ranks) changes the optimization dynamics — you often must scale the learning rate (linear scaling rule, warmup) and may hit a "large-batch generalization" wall. DP gives you throughput; it does not give you a better optimizer for free.

Chapter 4: Tensor (Model) Parallelism — Sharding a Layer

The problem it solves. Sometimes a single layer is too big for one GPU — a d_model=12288 MLP weight in a 175B model is itself tens of GB. Tensor parallelism (TP), introduced by Megatron-LM, splits the matrices within a layer across ranks so each holds a slice. The two ways to slice a matmul Y = W X are the heart of it.

Column-parallel. Split W by its output rows (columns of Wᵀ): rank r holds rows W_r and computes Y_r = W_r X — a disjoint slice of the output. No rank needs another's data, so the per-rank outputs simply concatenate into the full Y. No communication inside the op.

  W = [ W0 ]   →  rank0: Y0 = W0·X
      [ W1 ]      rank1: Y1 = W1·X      Y = concat(Y0, Y1)   (no comm)

Row-parallel. Split W by its input columns and X by its rows: rank r holds a column-slice W^{(r)} and the matching input slice X^{(r)}, and computes a partial sum Y_r = W^{(r)} X^{(r)}. The true Y is the sum of these partials, so you must all-reduce them:

$$ Y = \sum_{r} W^{(r)} X^{(r)} \quad\Rightarrow\quad \text{all-reduce the partial } Y_r. $$

  W = [ W^0 | W^1 ]   rank0: Y0 = W^0·X^0   ┐ all-reduce
  X = [ X^0 ; X^1 ]   rank1: Y1 = W^1·X^1   ┘ → Y = Y0 + Y1

The Megatron trick (why this pairing is genius). Megatron makes the MLP block's first matmul column-parallel and the second row-parallel. The column-parallel output is already sharded exactly the way the row-parallel second layer wants its input — so the entire two-matmul block needs only one all-reduce, at the very end (the g operator), with a conjugate identity/copy on the forward (f). Same for attention. That is how TP keeps communication to a minimum.

Why TP lives inside one node. That all-reduce happens inside every layer, on the critical path, many times per step — it cannot be overlapped away. So TP demands the fattest interconnect you have: NVLink within a single 8-GPU node (~900 GB/s), never the slower InfiniBand between nodes. Rule of thumb: tensor_model_parallel_size ≤ GPUs per node. Cross that boundary and the in-layer all-reduce over InfiniBand will dominate your step time.

The soul test. Both shardings must reconstruct the exact unsharded W X — column by concatenation, row by all-reduce. The lab's tensor_parallel_matmul proves it for both axes; if your sharded result doesn't equal the dense matmul, your model is computing a different function.

Common misconception. "Tensor parallelism is just data parallelism for big layers." No — DP replicates and syncs between steps; TP shards a single matrix and syncs within a layer, every forward and backward, which is why its communication is far more latency-sensitive and node-local.

Chapter 5: Pipeline Parallelism & the Bubble

The mechanism. Pipeline parallelism (PP) splits the model by depth: stage 0 holds layers 1–k, stage 1 holds k+1–2k, and so on across P GPUs. A batch flows forward through the stages like an assembly line, then gradients flow backward. The catch is the assembly-line startup: while stage 0 processes the first input, stages 1…P-1 sit idle waiting for work to reach them; symmetrically they drain at the end. That idle time is the bubble.

Micro-batches. To keep the pipe busy you split the batch into M micro-batches and feed them in a stream, so multiple stages work on different micro-batches at once. The fraction of total time still wasted on fill + drain is the classic GPipe result:

$$ \boxed{;\text{bubble} = \frac{P-1}{M+P-1};} $$

  GPipe, P=4 stages, M=4 microbatches (F = forward, B = backward, . = idle/bubble)
  stage0: F1 F2 F3 F4 . . . B4 B3 B2 B1
  stage1: .  F1 F2 F3 F4 . B4 B3 B2 B1 .
  stage2: .  .  F1 F2 F3 F4 B4 B3 B2 B1 . .
  stage3: .  .  .  F1 F2 F3 F4 B4 B3 B2 B1 . . .   ← bubble = the dots

Reading the formula. With P=4, M=4 the bubble is 3/7 ≈ 43% — nearly half the cluster idle. Crank M to 64 and it falls to 3/67 ≈ 4.5%; as M → ∞ it → 0. The fix for a pipeline bubble is more micro-batches, not more stages — adding stages raises P and makes the bubble worse.

GPipe vs 1F1B. Vanilla GPipe runs all forwards, then all backwards, so it must hold the activations of every in-flight micro-batch — memory grows with M. 1F1B (one-forward-one- backward, from PipeDream-Flush) interleaves: once the pipe is warm, each stage does a forward then immediately a backward, so it holds only ~P micro-batches of activations — same bubble, far less memory. Interleaved 1F1B (Megatron) gives each GPU several non-contiguous "virtual" stages, dividing the bubble by the number of virtual stages at the cost of more communication.

The lab's pipeline_schedule_1f1b models the makespan as 2M useful forward+backward unit-steps plus 2(P-1) bubble steps (fill + drain), and you can check that bubble_steps / total_steps equals (P-1)/(M+P-1) — the formula and the schedule agree.

Common misconception. "Pipeline parallelism is for speed." It's primarily for memory/fit — it lets a model deeper than one GPU's memory train at all, across nodes (PP communication is just point-to-point activations between adjacent stages, so it tolerates the slower InfiniBand far better than TP's all-reduce). The bubble is the price you pay for that fit.

Chapter 6: ZeRO / FSDP — Sharding the State, Not the Compute

The waste DP leaves on the table. Chapter 3's data parallelism replicates the entire 16N training state on every GPU — the parameters (2N), gradients (2N), and the big one, the optimizer state (12N: fp32 master 4N + Adam moments 8N). With 64 DP ranks you store the same 12N optimizer bytes 64 times. That redundancy is pure waste, and ZeRO (Zero Redundancy Optimizer, Rajbhandari et al.) eliminates it by sharding the state across the DP ranks instead of replicating it. PyTorch's FSDP (Fully Sharded Data Parallel) is the same idea, built into core.

The three stages. Each stage shards one more bucket of the 16N state; the rest stays replicated:

$$ \text{stage-3 per-GPU} = \frac{2N + 2N + 12N}{N_{\text{ranks}}} = \frac{16N}{N_{\text{ranks}}} \xrightarrow[N_{\text{ranks}}\to\infty]{} 0. $$

Worked numbers (7B, 8 GPUs). Stage 0 = 112 GB/GPU (won't fit an 80 GB card). Stage 1 = 38.5 GB. Stage 2 = 26.2 GB. Stage 3 = 14 GB. The optimizer state — the dominant 84 GB — is the first and biggest win, which is why stage 1 alone often makes a previously-impossible run fit, and stage 3 is what lets a 70B model train on commodity cards.

What stage 3 costs. If each GPU only holds 1/N of the params, it can't run a layer's matmul alone. So FSDP all-gathers the full parameters of each layer just in time for its forward (and again for backward), uses them, then frees them. You trade extra communication (an all-gather per layer) for a massive memory saving — and you overlap that all-gather with the previous layer's compute so the wall-clock hit is small. This is why FSDP/ZeRO-3 is the default for large models: the memory is the binding constraint, and the comm overlaps.

The lab. zero_memory_per_gpu reproduces this exact ladder — and the soul test asserts stage3 < stage2 < stage1 < stage0 and that stage 3 → 16N/N as ranks grow. That monotone drop is the reason the technique exists.

Common misconception. "ZeRO is a different parallelism axis from DP." No — ZeRO/FSDP is data parallelism with the redundant state sharded away. You still split the batch and behave like one big batch; you just stop storing 64 copies of the optimizer. (TP and PP, by contrast, change what compute each GPU does.)

Chapter 7: 3D Parallelism & Mapping to a Cluster

The composition. Real frontier runs don't pick one strategy — they compose them. 3D parallelism is \text{DP} \times \text{TP} \times \text{PP} (with ZeRO often layered onto the DP axis), and the total GPU count is the product:

$$ N_{\text{GPUs}} = N_{\text{DP}} \times N_{\text{TP}} \times N_{\text{PP}}. $$

Mapping to the topology is where seniority shows, because each axis has different communication and must land on the matching hardware tier:

  Cluster: 8 nodes × 8 GPUs = 64 GPUs                      Communication tier
  ┌──────────────────────────────────────────────┐
  │  TP=8   → inside ONE node (NVLink ~900 GB/s)   │   in-layer all-reduce, every layer
  │  PP=8   → ACROSS nodes (InfiniBand)            │   point-to-point activations, tolerant
  │  DP=1   → (or split GPUs left over)            │   gradient all-reduce, overlappable
  └──────────────────────────────────────────────┘
  Rule: TP ≤ GPUs/node;  PP across nodes;  DP/ZeRO outermost

The placement rules (memorize these — they are the answer to the cluster-mapping interview question):

1. Tensor parallelism stays inside a node. Its all-reduce is on the per-layer critical path, so it needs NVLink. TP ≤ GPUs per node (usually 8).
2. Pipeline parallelism spans nodes. It only sends activations point-to-point between adjacent stages, so it tolerates InfiniBand — put the pipeline across the slow links.
3. Data parallelism / ZeRO is outermost. Its gradient all-reduce overlaps with backward and is the most latency-tolerant, so it absorbs whatever GPUs are left and scales the cluster wider.

Why the order matters. Get it wrong — say, TP across nodes — and the in-layer all-reduce runs over InfiniBand instead of NVLink, and your 64-GPU job runs slower than a 8-GPU one. The mapping, not the raw GPU count, determines whether the cluster trains at all.

Common misconception. "More parallelism axes = faster." Each axis adds communication; the art is using the minimum parallelism that makes the model fit, then spending the rest on DP for throughput. TP and PP are fit mechanisms (use only as much as memory forces); DP is the throughput mechanism (scale it as wide as the network allows).

Chapter 8: Overlap, Activation Checkpointing & the Frameworks

Communication–computation overlap. The single most important systems trick: launch the communication for one piece of work while the GPU computes the next, so the network cost hides behind useful FLOPs. PyTorch DDP buckets gradients and all-reduces each bucket the instant it's ready during the backward pass, overlapping comm with the rest of backward. FSDP prefetches the next layer's all-gather while the current layer computes. A distributed run with no overlap is leaving most of its speedup on the table — the second thing to check after "is the mapping right."

Activation checkpointing (gradient checkpointing). Chapter 1's memory accounting ignored activations — the forward-pass intermediates saved for the backward pass, which scale with batch × sequence × layers and can rival the model state. Activation checkpointing stores only a few activations and recomputes the rest during backward: you trade ~33% more compute for a large activation-memory cut. It composes with every parallelism strategy and is essential for long context or large micro-batches. (It's the same compute-for-memory trade the Phase 00 KV-cache chapter hinted at, now on the training side.)

The frameworks (where the algorithms live).

The senior's debugging order when a distributed job is slow: (1) is the mapping right (TP in node, PP across)? (2) is communication overlapped with compute? (3) is the collective algorithm appropriate for the message size (ring vs tree)? (4) is activation checkpointing on where it should be? Almost every "scaling fell over at 64 GPUs" ticket is one of these four.

Common misconception. "The framework handles it, I don't need the math." The frameworks expose exactly these knobs — tensor_model_parallel_size, pipeline_model_parallel_size, zero_optimization.stage, activation_checkpointing, gradient_as_bucket_view — and choosing them is the math in this warmup. The framework runs the plan; you still have to write the plan.

Lab Walkthrough Guidance

The lab (lab-01-parallelism-engine) turns this warmup into code. Suggested order (matches the file):

1. ring_all_reduce — Chapter 2. The trick is the schedule: chunk each vector into N pieces, run N-1 reduce-scatter steps (neighbor adds the incoming chunk), then N-1 all-gather steps (neighbor replaces). The result must equal the naive element-wise sum — the soul test. Handle the remainder when length % N ≠ 0 (last chunk absorbs it).
2. all_reduce_cost — Chapter 2. steps = 2(N-1), bytes = 2·tensor·(N-1)/N; N=1 is the zero-cost edge case.
3. data_parallel_gradient_average — Chapter 3. All-reduce, then divide by N; assert it equals the serial mean.
4. tensor_parallel_matmul — Chapter 4. Column = concatenate per-rank dot products; row = validate the column-slices tile x, compute partial outputs, all-reduce them. Both must equal the dense W·x — the TP soul test.
5. pipeline_bubble_fraction / pipeline_schedule_1f1b — Chapter 5. The fraction is (P-1)/(M+P-1); the schedule's bubble_steps/total_steps must match it.
6. zero_memory_per_gpu — Chapter 6. Shard optimizer at stage ≥ 1, grads at ≥ 2, params at ≥ 3; the soul test is stage3 < stage2 < stage1 < stage0 and stage 3 → 16N/N.

Run red → green: pytest test_lab.py -v, then LAB_MODULE=solution pytest test_lab.py -v, then python solution.py for the worked numbers.

Success Criteria

• All tests pass against your lab.py and against LAB_MODULE=solution.
• You can derive the ring all-reduce's 2(N-1) steps and 2·tensor·(N-1)/N bytes/rank from the reduce-scatter + all-gather decomposition, and explain why bytes/rank is ~constant in N.
• You can explain why data-parallel gradient averaging is exactly one big batch, and what breaks if a rank is dropped.
• You can state which tensor-parallel sharding (column/row) needs an all-reduce and why TP must stay inside an NVLink node.
• You can compute pipeline_bubble_fraction(4, 4) = 3/7 in your head and say the fix is more micro-batches, plus the GPipe-vs-1F1B memory difference.
• You can walk the ZeRO ladder (112 → 38.5 → 26.2 → 14 GB for 7B on 8 GPUs) and name what stage 3 pays (a per-layer all-gather).
• You can map a DP × TP × PP job onto an 8-node × 8-GPU cluster and justify each placement.

Interview Q&A

• "Why is the ring all-reduce bandwidth-optimal, and how many steps?" — 2(N-1) steps (reduce-scatter + all-gather, each N-1); each rank moves 2·tensor·(N-1)/N bytes → ~2·tensor regardless of N, with no hotspot. That's why DP scales to thousands of GPUs.
• "What does an all-reduce compute, and why does data parallelism need it?" — the element-wise sum of every rank's gradient; divided by N it's the mean, which equals a single big-batch gradient — the correctness guarantee of DP.
• "Column vs row tensor parallelism — which communicates?" — column-parallel concatenates (no comm); row-parallel all-reduces partial sums. Megatron pairs column→row so a block needs one all-reduce; that all-reduce is on the critical path, so TP stays inside NVLink (TP ≤ GPUs/node).
• "What's the pipeline bubble and how do you shrink it?" — (P-1)/(M+P-1) idle fraction from fill+drain; shrink with more micro-batches M (→0 as M→∞), not more stages. 1F1B gets the same bubble as GPipe with ~P (not M) micro-batches of activation memory.
• "Walk me up the ZeRO stages." — 1 shards optimizer (12N, the biggest win), 2 +gradients, 3 +params (FSDP FULL_SHARD); per-GPU 16N → 16N/N. Stage 3 pays a per-layer all-gather of params, overlapped with compute.
• "Map a 175B run onto 8 nodes × 8 GPUs." — TP=8 inside each node (NVLink, in-layer all-reduce), PP=8 across nodes (point-to-point activations tolerate InfiniBand), DP/ZeRO outermost for the rest; minimum TP/PP to fit, maximum DP for throughput.
• "DDP vs FSDP for a 13B model on 80 GB cards?" — DDP replicates 16N ≈ 208 GB per GPU → won't fit; FSDP FULL_SHARD shards to 208/N GB → fits, at the cost of per-layer param all-gathers. Lead with the memory math.
• "My job scaled to 8 GPUs but is slow at 64 — why?" — communication now dominates: check the mapping (TP crossing nodes?), comm/compute overlap (bucketing, prefetch), the collective algorithm for the message size, and activation checkpointing. It's a network problem, not a FLOPs problem.
• "What is activation checkpointing and when do you reach for it?" — recompute activations in backward instead of storing them: ~33% more compute for a big activation-memory cut; essential for long context / large micro-batches and composes with every parallelism axis.
• "Ring vs tree all-reduce?" — ring is bandwidth-optimal for large (gradient-sized) tensors; tree/double-binary-tree wins for small, latency-bound messages. NCCL chooses by message size.

References

• Sergeev & Del Balso, Horovod (2018) — the ring all-reduce brought to deep learning; the clearest description of the 2(N-1)-step reduce-scatter + all-gather schedule.
• Shoeybi et al., Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism (2019) — column/row tensor parallelism and the f/g all-reduce placement.
• Rajbhandari et al., ZeRO: Memory Optimizations Toward Training Trillion Parameter Models (2020) — the stage 1/2/3 state-sharding ladder; the 16 bytes/param accounting.
• Huang et al., GPipe: Easy Scaling with Micro-Batch Pipeline Parallelism (2019) — the bubble and micro-batching.
• Narayanan et al., PipeDream / PipeDream-Flush (1F1B) (2019/2021) and Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM (2021) — 1F1B, interleaved schedules, and 3D-parallelism cluster mapping.
• Zhao et al., PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel (2023) — the FSDP design and the param all-gather/overlap.
• Rasley et al., DeepSpeed (2020) and the DeepSpeed / PyTorch DDP / FSDP / NCCL docs — the knobs (zero_optimization.stage, tensor_model_parallel_size, gradient_as_bucket_view) and the ring all-reduce write-ups.