{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# Module 23: Distributed Training" ], "id": "d3b07a2e-eb4e-4d48-8f28-a015e4162e1b" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [], "id": "1dd51e82-e7ee-4dc3-a7a3-e9b62ff940f7" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "output_type": "display_data", "metadata": {}, "data": {} } ], "source": [], "id": "8f00f4fb-e527-4e40-bc54-a7a43b54bd03" }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Introduction\n", "\n", "Every training module so far ([m02](../m02_autograd/lesson.qmd),\n", "[m07](../m07_training/lesson.qmd)) quietly assumed **one device holds\n", "the whole model**. Frontier LLMs do not fit. A 7-billion-parameter model\n", "needs *120 GB* just to hold its training state — more than any single\n", "GPU — and that is before a single activation. **Distributed training**\n", "is how the same math from m07 runs across a cluster.\n", "\n", "There are two orthogonal problems, and this module builds a from-scratch\n", "answer to each:\n", "\n", "- **Speed** → **data parallelism.** Replicate the model on $N$ devices,\n", " feed each a different slice of the batch, then **average the\n", " gradients** so every replica takes the exact step one big-batch device\n", " would. The averaging is a **ring all-reduce**, which we build as a\n", " real, bandwidth-optimal algorithm.\n", "- **Memory** → **ZeRO.** Plain data parallelism keeps a *full* copy of\n", " the parameters, gradients, and optimizer states on **every** device —\n", " $N$-way redundant. **ZeRO** (Zero Redundancy Optimizer) partitions\n", " them across devices, cutting per-device memory from $16\\Psi$ bytes\n", " toward $16\\Psi/N$.\n", "\n", "Why it matters for LLMs:\n", "\n", "- **Nothing else fits.** Every model past a few billion parameters is\n", " trained distributed; the memory math below is the reason the field\n", " talks in “GPU-hours”.\n", "- **The step is unchanged.** Done right, $N$ devices produce\n", " *bit-for-bit* the same gradient as one — distribution is a systems\n", " trick, not a new algorithm.\n", "- **Memory is a budget you can rebalance.** ZeRO turns “the model\n", " doesn’t fit” into “add devices,” on a precise, linear curve you will\n", " drive yourself.\n", "\n", "### What You’ll Learn\n", "\n", "After this module, you can:\n", "\n", "- Compute the **training memory budget** of any model: the $2+2+12 = 16$\n", " bytes-per-parameter model of mixed-precision Adam.\n", "- Explain **data parallelism** and why averaging gradients equals one\n", " big batch.\n", "- Build a **ring all-reduce** from scratch (reduce-scatter → all-gather)\n", " and prove it equals the naive sum, at $2(N{-}1)/N \\cdot M$ bytes per\n", " device.\n", "- Derive the three **ZeRO** stages ($P_{os}$, $P_{os+g}$, $P_{os+g+p}$)\n", " and their per-device memory, and see the $4\\times$/$8\\times$/$N\\times$\n", " savings.\n", "- Place **tensor parallel**, **pipeline parallel**, **FSDP**, and\n", " **gradient checkpointing** in the same map.\n", "\n", "### Prerequisites\n", "\n", "This module requires familiarity with:\n", "\n", "- [Module 02: Autograd](../m02_autograd/lesson.qmd) — gradients are the\n", " thing we average across devices.\n", "- [Module 07: Training](../m07_training/lesson.qmd) — the optimizer\n", " states (Adam’s momentum and variance) that dominate the memory budget.\n", "- [Module 09: Efficient\n", " Attention](../m09_efficient_attention/lesson.qmd) — the *inference*\n", " memory wall (the KV cache); this module is its *training* twin.\n", "\n", "## Intuition: One GPU Is Not Enough\n", "\n", "m09 hit a memory wall at *inference* — the KV cache grows with the\n", "context. Training hits a wall long before that, and for a different\n", "reason: to take **one Adam step** you must keep several full-model-sized\n", "tensors resident at once.\n", "\n", "Think of the model’s parameters $\\Psi$ as the unit. Modern training runs\n", "in **mixed precision**: the forward and backward passes use fast 16-bit\n", "floats, but the optimizer keeps a high-precision 32-bit copy so tiny\n", "updates don’t vanish. Counting bytes per parameter:\n", "\n", "- an **fp16 copy of the parameters** — used in the forward/backward pass\n", " → 2 bytes\n", "- an **fp16 copy of the gradients** → 2 bytes\n", "- the **fp32 optimizer states**: a master copy of the parameters (4),\n", " Adam’s momentum (4), and Adam’s variance (4) → 12 bytes\n", "\n", "That is $2 + 2 + 12 = \\mathbf{16}$ bytes for *every* parameter, every\n", "step. A 7.5B model therefore needs $7.5 \\times 10^9 \\times 16 = 120$ GB\n", "— and an 80 GB A100 cannot hold it. The parameters you actually compute\n", "with are only $\\tfrac{1}{8}$ of the bill; the optimizer states are\n", "$\\tfrac{3}{4}$ of it. **That imbalance is what ZeRO exploits.**\n", "\n", "> **Key Insight**\n", ">\n", "> The thing that doesn’t fit is rarely the parameters — it’s the\n", "> *optimizer state*. Adam’s momentum and variance, plus the fp32 master\n", "> copy, are $12$ of the $16$ bytes. Halving parameter precision barely\n", "> helps; partitioning the optimizer state is where the memory is.\n", "\n", "## The Math: The Training Memory Budget\n", "\n", "The byte-count above is exactly the memory model from the **ZeRO**\n", "paper. In code it is `training_memory_bytes` in `distributed.py`:" ], "id": "13dddbf2-c71a-445d-b2a2-e1c8ded47440" }, { "cell_type": "code", "execution_count": 1, "metadata": {}, "outputs": [ { "output_type": "stream", "name": "stdout", "text": [ " params: 15.0 GB\n", " grads: 15.0 GB\n", " optimizer: 90.0 GB\n", " total: 120.0 GB" ] } ], "source": [ "import sys\n", "sys.path.insert(0, \".\")\n", "from distributed import training_memory_bytes, zero_memory_per_device\n", "\n", "mem = training_memory_bytes(7_500_000_000) # 7.5B parameters\n", "for k, v in mem.items():\n", " print(f\"{k:>10}: {v/1e9:6.1f} GB\")" ], "id": "ae0d17b0" }, { "cell_type": "markdown", "metadata": {}, "source": [ "Params and grads are 2 bytes each; the optimizer states are 12 (Adam’s\n", "$K=12$). Now the key move — **ZeRO data parallelism partitions that\n", "state across $N$ devices** in three stages, each partitioning one more\n", "kind of state:\n", "\n", "| Stage | Name | What is partitioned | Per-device bytes |\n", "|-------|--------------|------------------------|--------------------|\n", "| 0 | baseline | nothing (full replica) | $16\\Psi$ |\n", "| 1 | $P_{os}$ | optimizer states | $4\\Psi + 12\\Psi/N$ |\n", "| 2 | $P_{os+g}$ | \\+ gradients | $2\\Psi + 14\\Psi/N$ |\n", "| 3 | $P_{os+g+p}$ | \\+ parameters | $16\\Psi/N$ |\n", "\n", "Watch a 7.5B model land on 64 GPUs:" ], "id": "f2a5d38c-3d5f-4e9d-a6e7-083d1571b7bf" }, { "cell_type": "code", "execution_count": 2, "metadata": {}, "outputs": [ { "output_type": "stream", "name": "stdout", "text": [ "stage 0: 120.00 GB/device\n", "stage 1: 31.41 GB/device\n", "stage 2: 16.64 GB/device\n", "stage 3: 1.88 GB/device" ] } ], "source": [ "psi, n = 7_500_000_000, 64\n", "for stage in range(4):\n", " gb = zero_memory_per_device(psi, n, stage) / 1e9\n", " print(f\"stage {stage}: {gb:7.2f} GB/device\")" ], "id": "2b41c13f" }, { "cell_type": "markdown", "metadata": {}, "source": [ "120 GB collapses to under 2 — the same numbers the ZeRO paper reports (a\n", "$4\\times$, $8\\times$, and $N\\times$ reduction). **The model didn’t\n", "shrink; the redundancy did.**\n", "\n", "### Drive the memory budget\n", "\n", "Slide the model size, the device count, and the ZeRO stage. The bar is\n", "per-device memory; the ghost behind it is the $16\\Psi$ baseline you\n", "started from." ], "id": "5847831d-91d8-47c9-ad9d-ea137a13e846" }, { "cell_type": "code", "execution_count": 3, "metadata": {}, "outputs": [], "source": [], "id": "61525752" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "output_type": "display_data", "metadata": {}, "data": {} } ], "source": [], "id": "43996108-f1c0-421c-86ef-3b373d33b938" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "output_type": "display_data", "metadata": {}, "data": {} } ], "source": [], "id": "d42571e1-416e-4d12-a0b1-111320564de2" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "output_type": "display_data", "metadata": {}, "data": {} } ], "source": [], "id": "ad43389f-9161-45d5-99b8-cc42f3341560" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "output_type": "display_data", "metadata": {}, "data": {} } ], "source": [], "id": "f2b56631-916f-49b3-b837-f79f5249f60e" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "output_type": "display_data", "metadata": {}, "data": {} } ], "source": [], "id": "567374b3-3055-46fc-b92d-afe023c95875" }, { "cell_type": "markdown", "metadata": {}, "source": [ "> **Try This**\n", ">\n", "> 1. **Set params to 70B.** Even stage 3 needs many devices to fit an\n", "> 80 GB GPU — this is why 70B models train on hundreds of GPUs.\n", "> 2. **Hold the model fixed and slide $N$.** Stage 1 flattens out (it\n", "> can never beat $4\\Psi$); stage 3 keeps falling as $16\\Psi/N$. That\n", "> gap *is* the difference between partitioning some state and\n", "> partitioning all of it.\n", "> 3. **Compare stage 1 vs stage 3 at $N=8$.** Small clusters still get\n", "> most of the win from partitioning just the optimizer state — the\n", "> cheapest stage to run.\n", "\n", "## Code: Data Parallelism & Ring All-Reduce\n", "\n", "Memory is one half; **speed** is the other. Data parallelism runs $N$\n", "replicas of the model, each on a different **shard** of the batch.\n", "Replica $i$ computes the gradient of its own mini-batch; then all\n", "replicas **average** their gradients so every one applies the *same*\n", "update. Because each shard’s loss is a mean over an equal slice, that\n", "average is *exactly* the gradient of the loss over the whole batch:\n", "\n", "$$\\frac{1}{N}\\sum_{i=1}^{N} \\nabla_i \\;=\\; \\nabla \\Big(\\tfrac{1}{B}\\textstyle\\sum_{b} \\ell_b\\Big)$$\n", "\n", "Let’s prove it on a linear model. `data_parallel_gradients` shards the\n", "batch, computes each shard’s gradient with autograd, sums them with a\n", "ring all-reduce, and divides by $N$:" ], "id": "1dad52b9-4140-4b32-93e5-3a97f0a2ecab" }, { "cell_type": "code", "execution_count": 4, "metadata": {}, "outputs": [ { "output_type": "stream", "name": "stdout", "text": [ "max |g_dp - g_single|: 5.960464477539063e-08\n", "identical step: True" ] } ], "source": [ "import torch\n", "from distributed import data_parallel_gradients, single_device_gradient\n", "\n", "torch.manual_seed(0)\n", "W = torch.randn(3, 5) # a shared (out=3, in=5) weight\n", "x_shards = [torch.randn(4, 5) for _ in range(4)] # 4 devices, 4 samples each\n", "y_shards = [torch.randn(4, 3) for _ in range(4)]\n", "\n", "g_dp = data_parallel_gradients(W, x_shards, y_shards)\n", "g_sd = single_device_gradient(W, torch.cat(x_shards), torch.cat(y_shards))\n", "\n", "print(\"max |g_dp - g_single|:\", (g_dp - g_sd).abs().max().item())\n", "print(\"identical step:\", torch.allclose(g_dp, g_sd, atol=1e-5))" ], "id": "f533efd7" }, { "cell_type": "markdown", "metadata": {}, "source": [ "Four devices, four separate gradients, one all-reduce — and the result\n", "is the single-device gradient over all 16 samples. **Distribution\n", "changed nothing about the math.** The only new machinery is the\n", "collective that did the averaging.\n", "\n", "### The all-reduce, and why it must be a ring\n", "\n", "An **all-reduce** leaves *every* device holding the sum of all devices’\n", "tensors. The obvious way — send everything to device 0, add, broadcast\n", "back — pushes $O(N \\cdot M)$ bytes through one poor GPU’s link while the\n", "rest sit idle. The **ring all-reduce** fixes both: it uses every link at\n", "once and moves only $2(N{-}1)/N \\cdot M$ bytes per device, *independent\n", "of $N$* as it grows.\n", "\n", "The trick is to arrange the devices in a logical ring (each sends only\n", "to its right neighbour) and split every tensor into $N$ chunks, then run\n", "two phases:\n", "\n", "1. **Reduce-scatter** ($N{-}1$ steps): chunks circulate and accumulate\n", " until each device owns the *fully summed* value of exactly one\n", " chunk.\n", "2. **All-gather** ($N{-}1$ steps): those finished chunks circulate\n", " again until every device has all of them.\n", "\n", "This is `ring_all_reduce` in `distributed.py` — a real implementation,\n", "snapshotting each step’s sends before applying receives so it matches\n", "true simultaneous communication. Its correctness anchor is exact:" ], "id": "2bbd47ff-1e1f-4879-ad64-87da29e5d61b" }, { "cell_type": "code", "execution_count": 5, "metadata": {}, "outputs": [ { "output_type": "stream", "name": "stdout", "text": [ "device 0 result: [1111.0, 2222.0, 3333.0, 4444.0]\n", "== naive sum: True\n", "all devices agree: True\n", "bytes/device (M=4KB, N=4): 6000.0" ] } ], "source": [ "from distributed import ring_all_reduce, ring_all_reduce_bytes\n", "\n", "tensors = [\n", " torch.tensor([1., 2., 3., 4.]),\n", " torch.tensor([10., 20., 30., 40.]),\n", " torch.tensor([100., 200., 300., 400.]),\n", " torch.tensor([1000., 2000., 3000., 4000.]),\n", "]\n", "out = ring_all_reduce(tensors)\n", "print(\"device 0 result:\", out[0].tolist())\n", "print(\"== naive sum: \", torch.equal(out[0], torch.stack(tensors).sum(0)))\n", "print(\"all devices agree:\", all(torch.equal(o, out[0]) for o in out))\n", "print(\"bytes/device (M=4KB, N=4):\", ring_all_reduce_bytes(4000, 4))" ], "id": "c616b1b9" }, { "cell_type": "markdown", "metadata": {}, "source": [ "Bit-for-bit equal to the sum, on every device. And the communication\n", "cost $2(N{-}1)/N \\cdot M$ approaches $2M$ — a fixed budget no matter how\n", "many GPUs you add, which is exactly why the ring, not a central reducer,\n", "is what real systems (NCCL, Horovod) run.\n", "\n", "### Step through the ring\n", "\n", "Four devices, each starting with one row of values; watch reduce-scatter\n", "accumulate one finished chunk per device (the ringed cell), then\n", "all-gather spread the finished chunks around. By the last step every\n", "cell equals the column sum." ], "id": "93e45a85-af70-4bc1-8b47-efbcaf951fd5" }, { "cell_type": "code", "execution_count": 6, "metadata": {}, "outputs": [], "source": [], "id": "197899ba" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "output_type": "display_data", "metadata": {}, "data": {} } ], "source": [], "id": "eb982ea3-fea8-48a2-9ffe-631b95ebf71a" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "output_type": "display_data", "metadata": {}, "data": {} } ], "source": [], "id": "52184149-9166-4e75-9d54-1da945718f3a" }, { "cell_type": "markdown", "metadata": {}, "source": [ "> **Try This**\n", ">\n", "> 1. **Step to the end of reduce-scatter (step 3).** Exactly one cell\n", "> per device is ringed — the finished chunk. No device has the whole\n", "> answer yet.\n", "> 2. **Finish the all-gather.** The ringed chunks copy around the ring\n", "> until every row equals the column sums. That second half is why\n", "> the cost is $2\\times$ the scatter, not $1\\times$.\n", "\n", "## ZeRO: Stop Replicating What You Can Partition\n", "\n", "Data parallelism is fast but memory-wasteful: all $N$ devices hold the\n", "*same* 16 bytes/param. **ZeRO** keeps data parallelism’s compute pattern\n", "but removes the redundancy — each device stores only its **slice** of\n", "the state and fetches the rest with collectives (an all-gather for\n", "parameters just before they’re used, a reduce-scatter for gradients)\n", "exactly when needed.\n", "\n", "- **Stage 1 — $P_{os}$:** partition the **optimizer states** (the 12\n", " fp32 bytes). Each device updates only its slice of the master weights,\n", " then all-gathers the updated fp16 params. Per-device:\n", " $4\\Psi + 12\\Psi/N$.\n", "- **Stage 2 — $P_{os+g}$:** also partition the **gradients** — a device\n", " only needs the gradient slice for the optimizer slice it owns.\n", " Per-device: $2\\Psi + 14\\Psi/N$.\n", "- **Stage 3 — $P_{os+g+p}$:** also partition the **parameters**\n", " themselves, all-gathering each layer’s weights just-in-time for its\n", " forward/backward, then discarding them. Per-device: $16\\Psi/N$ — a\n", " true $N\\times$ split. (This is what PyTorch ships as **FSDP**.)\n", "\n", "The step through the stages, with the $N=64$ / 7.5B numbers as the\n", "caption:" ], "id": "0dcc951a-a997-42b5-989f-d1c1ab02fa25" }, { "cell_type": "code", "execution_count": 7, "metadata": {}, "outputs": [], "source": [], "id": "89637aa6" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "output_type": "display_data", "metadata": {}, "data": {} } ], "source": [], "id": "dfee206f-c779-4aa8-b0d7-4608423b5a31" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "output_type": "display_data", "metadata": {}, "data": {} } ], "source": [], "id": "c8a94ada-1bcf-4868-a0ce-208174c8eb9d" }, { "cell_type": "markdown", "metadata": {}, "source": [ "> **Key Insight**\n", ">\n", "> ZeRO is not a different parallelism from data parallelism — it *is*\n", "> data parallelism, with the redundant storage removed and paid back\n", "> with communication. Stage 3 partitions everything and is what\n", "> PyTorch’s **FSDP** implements; you trade a modest amount of extra\n", "> all-gather traffic for an $N\\times$ memory cut.\n", "\n", "## Beyond Data Parallelism\n", "\n", "Data parallelism (and ZeRO) shards the *batch* and, optionally, the\n", "*state* — but every device still runs the whole model graph. When a\n", "single **layer** is too big, or the pipeline of layers is too deep, two\n", "other axes come in. They compose: real systems run **3D parallelism**\n", "(data × tensor × pipeline).\n", "\n", "| Kind | What it splits | The cost it pays |\n", "|------------|-----------------------------|--------------------------------|\n", "| **Data parallel** (this module) | the batch | an all-reduce of gradients per step |\n", "| **Tensor parallel** | each layer’s matrices (split $QKV$/FFN across devices) | an all-reduce *inside* every layer — needs fast intra-node links |\n", "| **Pipeline parallel** | the stack of layers into stages | a “bubble” of idle time; hidden by micro-batching |\n", "| **ZeRO / FSDP** | the optimizer/grad/param state | extra all-gathers to reassemble params just-in-time |\n", "\n", "One more lever is orthogonal to all of these — **gradient\n", "checkpointing** (a.k.a. activation recomputation). The backward pass\n", "needs the activations from the forward pass; storing them all costs\n", "memory that grows with depth × batch × sequence. Checkpointing keeps\n", "only a few and **recomputes** the rest during the backward pass —\n", "trading roughly one extra forward pass (\\$\\$33% more compute) for a\n", "large activation-memory cut. It stacks on top of everything above.\n", "\n", "## Common Pitfalls\n", "\n", "> **Uneven shards break the “one big batch” equality**\n", ">\n", "> The gradient average equals the full-batch gradient *only when the\n", "> shards are the same size*. With ragged shards you’d need a weighted\n", "> average by sample count; `data_parallel_gradients` rejects uneven\n", "> shards rather than silently bias the step.\n", "\n", "> **Scaling $N$ scales the effective batch — retune the LR**\n", ">\n", "> $N$ devices with per-device batch $b$ is one global batch of $Nb$.\n", "> Larger batches usually need a larger learning rate (the linear-scaling\n", "> rule) and a warmup, or the first steps diverge. Distribution changes\n", "> the *batch*, and the batch changes the *optimizer*\n", "> ([m07](../m07_training/lesson.qmd)).\n", "\n", "> **Communication can dominate if you don’t overlap it**\n", ">\n", "> The all-reduce moves $\\sim 2M$ bytes every step. Naively you’d\n", "> compute, then communicate, then wait. Real frameworks **overlap** the\n", "> gradient all-reduce of early layers with the backward pass of later\n", "> ones, so most of the transfer is free. Without overlap, distributed\n", "> training can be *slower* than one device.\n", "\n", "> **ZeRO trades memory for communication, not for free**\n", ">\n", "> Higher ZeRO stages add all-gather traffic to reassemble what they\n", "> partitioned. Stage 3 on a slow interconnect can be\n", "> communication-bound. Pick the lowest stage that makes the model fit —\n", "> stage 1 is often enough and the cheapest.\n", "\n", "## Exercises\n", "\n", "### Exercise 1: All-reduce **mean** from all-reduce **sum**\n", "\n", "`ring_all_reduce` returns the *sum*. Wrap it to return the mean (what\n", "gradient averaging actually needs), and confirm it matches\n", "`torch.stack(...).mean(0)`." ], "id": "d95c3a68-ec8c-4fee-ac81-94c0d135ce99" }, { "cell_type": "code", "execution_count": 8, "metadata": {}, "outputs": [ { "output_type": "stream", "name": "stdout", "text": [ "matches mean: True" ] } ], "source": [ "import torch\n", "from distributed import ring_all_reduce\n", "\n", "def all_reduce_mean(tensors):\n", " # Your implementation here: sum, then divide by the device count.\n", " summed = ring_all_reduce(tensors)\n", " n = len(tensors)\n", " return [s / n for s in summed]\n", "\n", "xs = [torch.randn(6) for _ in range(4)]\n", "got = all_reduce_mean(xs)[0]\n", "print(\"matches mean:\", torch.allclose(got, torch.stack(xs).mean(0), atol=1e-6))" ], "id": "0bab3ea6" }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Exercise 2: Derive the stage-2 memory formula\n", "\n", "Stage 2 partitions the optimizer states (12 bytes) **and** the gradients\n", "(2 bytes) across $N$, but replicates the fp16 params (2 bytes). Write\n", "the per-device byte formula and check it against\n", "`zero_memory_per_device(psi, N, 2)`." ], "id": "c1647fd7-f7c1-4992-b06c-9753a1798f84" }, { "cell_type": "code", "execution_count": 9, "metadata": {}, "outputs": [ { "output_type": "stream", "name": "stdout", "text": [ "mine: 3750000.0\n", "ref: 3750000.0" ] } ], "source": [ "from distributed import zero_memory_per_device\n", "\n", "def stage2_bytes(psi, n):\n", " # Your implementation here: 2 bytes replicated + (2 + 12) bytes partitioned.\n", " return 2 * psi + (2 + 12) * psi / n\n", "\n", "psi, n = 1_000_000, 8\n", "print(\"mine:\", stage2_bytes(psi, n))\n", "print(\"ref: \", zero_memory_per_device(psi, n, 2))" ], "id": "80e7c27e" }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Exercise 3: Gradient accumulation — data parallelism on one device\n", "\n", "Gradient accumulation gets a big effective batch on *one* device by\n", "summing gradients over several micro-batches before stepping. Show it\n", "equals the data-parallel gradient over the same shards." ], "id": "dda6f08c-0374-4c15-80df-f9c5f803fa3d" }, { "cell_type": "code", "execution_count": 10, "metadata": {}, "outputs": [ { "output_type": "stream", "name": "stdout", "text": [ "accumulation == data parallel: True" ] } ], "source": [ "import torch\n", "from distributed import data_parallel_gradients\n", "\n", "torch.manual_seed(1)\n", "W = torch.randn(2, 3)\n", "shards_x = [torch.randn(4, 3) for _ in range(3)]\n", "shards_y = [torch.randn(4, 2) for _ in range(3)]\n", "\n", "def accumulate(W, xs, ys):\n", " # Your implementation here: mean-grad per micro-batch, averaged.\n", " grads = []\n", " for x, y in zip(xs, ys):\n", " w = W.detach().clone().requires_grad_(True)\n", " loss = ((x @ w.T - y) ** 2).mean()\n", " grads.append(torch.autograd.grad(loss, w)[0])\n", " return sum(grads) / len(grads)\n", "\n", "acc = accumulate(W, shards_x, shards_y)\n", "dp = data_parallel_gradients(W, shards_x, shards_y)\n", "print(\"accumulation == data parallel:\", torch.allclose(acc, dp, atol=1e-6))" ], "id": "1a9c3c48" }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Summary\n", "\n", "Key takeaways:\n", "\n", "1. **Training memory is $16$ bytes/param** — under mixed-precision\n", " Adam: 2 (fp16 params) + 2 (fp16 grads) + 12 (fp32 master +\n", " momentum + variance). The optimizer state, not the parameters, is\n", " what doesn’t fit.\n", "2. **Data parallelism replicates the model and averages gradients**,\n", " and that average is *bit-for-bit* the gradient one big-batch device\n", " would compute — a distributed step is not an approximation.\n", "3. **The ring all-reduce** does the averaging in two phases\n", " (reduce-scatter → all-gather), bit-exact to the naive sum, at\n", " $2(N{-}1)/N \\cdot M$ bytes per device — a fixed budget as $N$ grows,\n", " which is why it, not a central reducer, is the standard.\n", "4. **ZeRO removes data parallelism’s redundancy** in three stages —\n", " partitioning the optimizer states ($4\\Psi + 12\\Psi/N$), then\n", " gradients ($2\\Psi + 14\\Psi/N$), then parameters ($16\\Psi/N$) — up to\n", " an $N\\times$ memory cut. Stage 3 is FSDP.\n", "5. **Other axes compose:** tensor parallel splits each layer, pipeline\n", " parallel splits the stack, and gradient checkpointing trades compute\n", " for activation memory. Frontier runs combine all of them.\n", "\n", "## What’s Next\n", "\n", "You can now build the model (m01–m06), train it (m07), align it (m12),\n", "serve it (m08–m09), and **scale its training across a cluster** (this\n", "module). The frontier from here: **tensor & pipeline parallelism** as\n", "runnable code, **FSDP** wired into a real `GPTModel`, and **gradient\n", "checkpointing** built from an autograd hook — all filed on the roadmap.\n", "\n", "### Going Deeper\n", "\n", "**Core Papers:**\n", "\n", "- [ZeRO: Memory Optimizations Toward Training Trillion Parameter\n", " Models](https://arxiv.org/abs/1910.02054) — Rajbhandari et al. (2020):\n", " the $2+2+12$ memory model and the three-stage partitioning built here.\n", "- [Bandwidth Optimal All-reduce Algorithms for Clusters of\n", " Workstations](https://web.cels.anl.gov/~thakur/papers/ring.pdf) —\n", " Patarasuk & Yuan (2009): the ring all-reduce and its\n", " $2(N{-}1)/N \\cdot M$ cost.\n", "- [Horovod: fast and easy distributed deep learning in\n", " TensorFlow](https://arxiv.org/abs/1802.05799) — Sergeev & Del Balso\n", " (2018): ring all-reduce brought to deep learning.\n", "- [PyTorch FSDP: Experiences on Scaling Fully Sharded Data\n", " Parallel](https://arxiv.org/abs/2304.11277) — Zhao et al. (2023):\n", " ZeRO-3 as a production PyTorch API.\n", "- [Training Deep Nets with Sublinear Memory\n", " Cost](https://arxiv.org/abs/1604.06174) — Chen et al. (2016): gradient\n", " checkpointing.\n", "\n", "**Practical Resources:**\n", "\n", "- [DeepSpeed](https://www.deepspeed.ai/) — the library that popularized\n", " ZeRO.\n", "- [Megatron-LM](https://arxiv.org/abs/1909.08053) — Shoeybi et\n", " al. (2019): tensor parallelism for transformers." ], "id": "c6f475b6-685f-4416-9e00-fea4d2351b3d" } ], "nbformat": 4, "nbformat_minor": 5, "metadata": { "kernelspec": { "name": "python3", "display_name": "Python 3 (ipykernel)", "language": "python", "path": "/opt/buildhome/.asdf/installs/python/3.11.14/share/jupyter/kernels/python3" }, "language_info": { "name": "python", "codemirror_mode": { "name": "ipython", "version": "3" }, "file_extension": ".py", "mimetype": "text/x-python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.11.14" } } }