{
 "cells": [
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "---\n",
    "execute:\n",
    "  enabled: false\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=multi-gpu-practice-md-1 fw=all -->\n",
    "\n",
    "# Multi-GPU in Practice\n",
    "\n",
    "The hand-rolled loop of that section taught the mechanism and\n",
    "then lost the race: on a tiny model over a host-staged wire, the second\n",
    "GPU made things slower. Two things were wrong, and only one was the\n",
    "hardware. This section fixes the other — the *software* — by replacing our\n",
    "loop with the production machinery, and in doing so meets the chapter's\n",
    "sharpest framework contrast: PyTorch makes you launch processes and the\n",
    "collectives are explicit; JAX runs in one process and you merely *annotate\n",
    "the layout*, letting the compiler write the collectives for you. We\n",
    "measure real scaling on 2–4 GPUs, sketch how the same ideas shard a model\n",
    "too big to replicate (FSDP), and stop at the edge of the single node.\n",
    "\n",
    "*Prerequisites: the from-scratch data-parallel loop, the ring-allreduce\n",
    "identity, and the cost model of* that section*; the memory\n",
    "anatomy of* that section*. The multi-process idiom below\n",
    "was verified to run under this book's notebook build; why this box's\n",
    "fabric is slow is the topology story of* that section*,\n",
    "and what a collective costs on it is* that section*'s\n",
    "measurement.*"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "multi-gpu-practice-multi-gpu-in-practice"
   },
   "source": [
    "%matplotlib inline\n",
    "from d2l import jax as d2l\n",
    "from flax import nnx\n",
    "import jax\n",
    "from jax import numpy as jnp\n",
    "import numpy as np\n",
    "import optax\n",
    "import os\n",
    "import re\n",
    "import time\n",
    "\n",
    "# XLA's NCCL buffer registration fails harmlessly on this P2P-less box;\n",
    "# opt out rather than let every collective print the warning.\n",
    "os.environ.setdefault('NCCL_LOCAL_REGISTER', '0')"
   ],
   "execution_count": 1,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "'0'"
     },
     "metadata": {},
     "execution_count": 1
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=multi-gpu-practice-md-2 fw=all -->\n",
    "\n",
    "## What Our Hand-Rolled Loop Lacked\n",
    "\n",
    "Our that section implementation had three deficits, and modern\n",
    "data parallelism repairs each.\n",
    "\n",
    "* **No overlap.** Our loop finished the *entire* backward pass, then called\n",
    "  `allreduce`. But gradients become available layer by layer, back to\n",
    "  front, so the last layers' gradients could be communicated while the\n",
    "  earlier layers are still computing. Serializing compute-then-communicate\n",
    "  wastes exactly the time the fabric is idle during backward.\n",
    "* **One Python process.** A single interpreter drove all $k$ GPUs, so one\n",
    "  GIL-bound thread dispatched every kernel — the overhead regime of\n",
    "  that section, multiplied by $k$.\n",
    "* **A star topology.** Our `allreduce` funneled everything through device\n",
    "  0; that section showed the ring moves a constant per device\n",
    "  instead.\n",
    "\n",
    "PyTorch's `DistributedDataParallel` (DDP) fixes all three\n",
    "[@Li.Zhao.Varma.ea.2020]: one **process per GPU** (no shared GIL),\n",
    "NCCL's **ring/tree collectives** (no hub), and — the headline —\n",
    "gradient **bucketing that overlaps communication with the backward pass**\n",
    "(the figure). As each bucket of gradients fills, DDP kicks\n",
    "off its allreduce immediately, so by the time the backward pass reaches the\n",
    "first layer, the last layers' gradients are already summed. This is\n",
    "compute–communication overlap — independent work scheduled onto the fabric\n",
    "while the GPUs keep computing — finally shown where it pays.\n",
    "\n",
    "![What DDP buys over the hand-rolled loop: gradient bucketing lets each\n",
    "bucket's allreduce overlap the rest of the backward pass, instead of\n",
    "waiting for all of it.](../img/mdl-perf-ddp-overlap.svg)\n",
    "\n",
    "## DDP, Really Run\n",
    "\n",
    "DDP needs multiple processes, and a notebook is one process — so we launch\n",
    "the extra ones. The idiom, verified to work under this book's build: write\n",
    "the training script to a sidecar file, then launch it with `torchrun`,\n",
    "which spawns one process per GPU, sets up the rendezvous, and runs the\n",
    "script under `init_process_group`. Each rank writes its result to a scratch\n",
    "directory the notebook reads back (cleared before every launch, so a\n",
    "crashed run can never serve stale numbers). We keep the script minimal; the\n",
    "*loop body* is unchanged from single-GPU training, and the scaffolding\n",
    "around it is what the launcher's multi-process world requires:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "multi-gpu-practice-ddp-really-run-1"
   },
   "source": [
    "\n",
    "class ResNet18(nnx.Module):\n",
    "    \"\"\"A slightly modified ResNet-18 (small stem, no max-pool).\"\"\"\n",
    "    def __init__(self, num_classes=10, rngs=None):\n",
    "        rngs = nnx.Rngs(d2l.get_key()) if rngs is None else rngs\n",
    "        self.net = nnx.Sequential(\n",
    "            nnx.Conv(1, 64, (3, 3), (1, 1), padding='same', rngs=rngs),\n",
    "            nnx.BatchNorm(64, rngs=rngs), nnx.relu,\n",
    "            d2l.Residual(64, in_channels=64, rngs=rngs),\n",
    "            d2l.Residual(64, in_channels=64, rngs=rngs),\n",
    "            d2l.Residual(128, use_1x1conv=True, strides=(2, 2),\n",
    "                         in_channels=64, rngs=rngs),\n",
    "            d2l.Residual(128, in_channels=128, rngs=rngs),\n",
    "            d2l.Residual(256, use_1x1conv=True, strides=(2, 2),\n",
    "                         in_channels=128, rngs=rngs),\n",
    "            d2l.Residual(256, in_channels=256, rngs=rngs),\n",
    "            d2l.Residual(512, use_1x1conv=True, strides=(2, 2),\n",
    "                         in_channels=256, rngs=rngs),\n",
    "            d2l.Residual(512, in_channels=512, rngs=rngs),\n",
    "            lambda x: x.mean(axis=(1, 2)),\n",
    "            nnx.Linear(512, num_classes, rngs=rngs))\n",
    "\n",
    "    def __call__(self, x):\n",
    "        return self.net(x)"
   ],
   "execution_count": 2,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=multi-gpu-practice-md-3 fw=all -->\n",
    "\n",
    "The DDP training script, written to disk from a cell. The two lines that\n",
    "make training data-parallel are `init_process_group(\"nccl\")` once per\n",
    "process and the `DDP(model)` wrap — after those, *the loop body is\n",
    "unchanged from single-GPU*. The rest is the launcher's housekeeping, shown\n",
    "in full so none of it is magic: a rank-local device, a\n",
    "`DistributedSampler` that hands each rank a disjoint shard of the data\n",
    "(re-shuffled per epoch by `set_epoch`), and process-group teardown. One\n",
    "practicality earns its own line: the parent notebook downloads\n",
    "Fashion-MNIST once, quietly, *before* any rank exists, so the ranks never\n",
    "race on the download:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=multi-gpu-practice-md-4 fw=all -->\n",
    "\n",
    "One word before the numbers, because every published speedup quietly picks\n",
    "a convention (that section). This sweep holds the\n",
    "*per-rank* batch fixed at 256, so the global batch grows to $256k$ —\n",
    "**weak scaling**: the efficiencies it prints are throughput per GPU\n",
    "relative to one GPU. Remember that a grown global batch also changes the\n",
    "optimization trajectory itself (that section); the\n",
    "strong-scaling question — same global batch, finishing sooner — comes\n",
    "right after."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=multi-gpu-practice-md-5 fw=all -->\n",
    "\n",
    "On our four-GPU box, ResNet-18 on Fashion-MNIST-64 (11.2M parameters)\n",
    "scales the way the accounting of that section says a\n",
    "*compute-dense* model should: about 1.8× at two GPUs and about 3.3× at\n",
    "four — the cell prints 88% and 82% weak-scaling efficiency. The efficiency\n",
    "sags gently — no cliff — because each step's compute is large enough to\n",
    "hide most of the communication, unlike LeNet. Strong scaling asks a harder\n",
    "question of the same hardware: hold the global batch at 512 and split it\n",
    "ever thinner, so the per-rank batch shrinks as $512/k$ while the allreduce\n",
    "stays the same size:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=multi-gpu-practice-md-6 fw=all -->\n",
    "\n",
    "At two GPUs the two conventions nearly coincide on this model — a 4090 is\n",
    "about equally saturated by a per-rank batch of 256 or 512, so the $k=1$\n",
    "baselines barely differ — but they part company as $k$ grows: by $k=4$ the\n",
    "strong-scaling run gives each GPU only 128 examples per step, and any gap\n",
    "that opens between the strong and weak efficiencies is pure\n",
    "underutilization, the mechanism that sank LeNet, now in miniature. That is\n",
    "the equation read twice: weak scaling holds $t_{\\text{compute}}$\n",
    "per device constant and asks the fabric to keep up; strong scaling shrinks\n",
    "$t_{\\text{compute}}(B/k)$ toward the fixed $t_{\\text{comm}}$ floor.\n",
    "\n",
    "Now confront the sweep with the cost model by measurement rather than by\n",
    "assertion. the equation prices a step's communication at\n",
    "$2N/\\beta$, and DDP itself provides the instrument to check it:\n",
    "`no_sync()` runs backward with gradient synchronization turned off, so the\n",
    "difference between synced and unsynced step times estimates what\n",
    "communication really costs — an estimate, not a truth, since turning sync\n",
    "off also changes what overlaps what:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=multi-gpu-practice-md-7 fw=all -->\n",
    "\n",
    "Prediction and measurement land within tens of percent of each other, and\n",
    "**that agreement is the result**: a scaling curve you can price before you\n",
    "buy, not a marketing \"N× faster\" claim — and it holds at two GPUs as\n",
    "clearly as at four. One line on what a datacenter box changes: an NVLink\n",
    "fabric shrinks $t_{\\text{comm}}$ by roughly two orders of magnitude\n",
    "(the table), so the same accounting predicts near-linear\n",
    "scaling — same model, different constant. (The legacy `nn.DataParallel` is\n",
    "single-process and GIL-bound; use DDP even on one node, as PyTorch's own\n",
    "docs advise.)\n",
    "\n",
    "One loose end from that section deserves to be closed by\n",
    "measurement rather than left as an assertion. The diagnosis there was\n",
    "that NCCL's P2P-less fallback moves bytes with a latency-bound GPU-kernel\n",
    "copy — in effect a performance bug in how the library's default transport\n",
    "interacts with this particular platform — and that one documented switch,\n",
    "`NCCL_SHM_USE_CUDA_MEMCPY=1`, re-routes the same transfer over the DMA\n",
    "copy engines. Measure the bare collective both ways:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=multi-gpu-practice-md-8 fw=all -->\n",
    "\n",
    "Roughly five-fold, from configuration alone. First the general lesson: a\n",
    "collective library's configuration — which transport it picks, what\n",
    "topology it assumes — can move communication performance by *factors*,\n",
    "not percent, so measure yours against what the wire demonstrably carries\n",
    "(that section's raw copy) before trusting it. Then the\n",
    "specific one, which is why every training run above still uses the\n",
    "library's defaults: this workaround wins the microbenchmark and loses\n",
    "the workload. On our box, with this NCCL build, the copy-engine mode\n",
    "deadlocks DDP's overlapped training path within seconds — the very\n",
    "collectives that just ran flawlessly in isolation wedge inside the\n",
    "training loop — so adopting it here would trade a five-fold bandwidth\n",
    "win for a hung notebook. An escape hatch is platform-specific twice over:\n",
    "whether you need it, and whether it survives your workload, are both\n",
    "measurements (the exercises have you reproduce both halves). The cost\n",
    "model is indifferent either way — the equation simply takes\n",
    "whatever $\\beta$ your fabric, as configured, sustains.\n",
    "\n",
    "## Sharding the Redundant: the FSDP Idea\n",
    "\n",
    "DDP replicates *everything* on every rank: $k$ identical copies of the\n",
    "parameters, the gradients, and the optimizer states. For the $16P$-byte\n",
    "training footprint of that section, that is $k-1$ copies\n",
    "of everything, wasted — and it caps the model size at what one GPU holds,\n",
    "the limitation that section flagged. **Fully Sharded Data\n",
    "Parallel** (FSDP) removes the redundancy by *sharding* those tensors across\n",
    "ranks, each rank owning $1/k$ of each, and materializing a full layer only\n",
    "for the moment it is needed [@Zhao.Gu.Varma.ea.2023]. The idea is the\n",
    "ZeRO ladder [@Rajbhandari.Rasley.Ruwase.ea.2020]: shard the\n",
    "optimizer states first (they are the biggest, $8P$), then the gradients,\n",
    "then the parameters — each rung cutting memory toward $1/k$ at the cost of\n",
    "more communication.\n",
    "\n",
    "The mechanism is the that section identity, cashed in. Recall\n",
    "that allreduce = reduce-scatter + all-gather. FSDP simply *keeps the two\n",
    "halves separate*: an **all-gather** reconstructs a layer's full parameters\n",
    "just before it computes, and frees them just after; a **reduce-scatter**\n",
    "sums each layer's gradients but leaves each rank holding only its own\n",
    "shard (the figure). No tensor's full replica ever lives\n",
    "longer than the layer that needs it.\n",
    "\n",
    "![The FSDP lifecycle of one block, under one simplified\n",
    "reshard-after-forward policy: parameters live sharded ($P/k$ per rank); an\n",
    "all-gather materializes the full block only while it computes, then frees\n",
    "it; a reduce-scatter leaves each rank with its gradient\n",
    "shard.](../img/mdl-perf-fsdp-lifecycle.svg)\n",
    "\n",
    "That completes the small family of collectives this chapter needs — worth\n",
    "one table, since the rest of the book will name them without ceremony:\n",
    "\n",
    "| collective | what every rank ends with | where it appears |\n",
    "|---|---|---|\n",
    "| allreduce | the full sum | DDP's gradient buckets; that section |\n",
    "| reduce-scatter | one shard of the sum | FSDP gradients |\n",
    "| all-gather | every shard, concatenated | FSDP parameters, just-in-time |\n",
    "| all-to-all | a different shard from each peer | expert parallelism (that section) |\n",
    "\n",
    "FSDP's payoff — fitting a model that does not fit — is invisible on our\n",
    "11.2M-parameter demo, which occupies a few hundred MB of a 24 GB card, so\n",
    "we show the *shape* of the code rather than run it. The modern API is\n",
    "`fully_shard` over a `DeviceMesh`; the original `FullyShardedDataParallel`\n",
    "wrapper class still imports at our pin, but it is the deprecated legacy\n",
    "path:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "multi-gpu-practice-sharding-the-redundant-the-fsdp-idea"
   },
   "source": [
    "# In JAX the same sharding is a PartitionSpec, not a wrapper — see below.\n",
    "print('JAX shards by annotation; the next subsection is the demo')"
   ],
   "execution_count": 3,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "JAX shards by annotation; the next subsection is the demo\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=multi-gpu-practice-md-9 fw=all -->\n",
    "\n",
    "You reach for FSDP when the training state at your precision — the\n",
    "parameters, gradients, and optimizer states of\n",
    "that section's anatomy, plus activations — no longer\n",
    "fits on one GPU, or when that redundancy is worth trading away for\n",
    "communication; for this card class the threshold arrives at a few billion\n",
    "parameters. The production distributed-training map, and how to combine\n",
    "FSDP with the other parallelism axes, lives in\n",
    "that section.\n",
    "\n",
    "## JAX: Annotate the Layout, the Compiler Writes the Collectives\n",
    "\n",
    "Everything above was PyTorch's world: multiple processes, explicit\n",
    "collectives, a launcher. JAX offers a different deal, and it is the\n",
    "chapter's cleanest framework contrast. One process sees all the GPUs\n",
    "(that section); you describe *how the data is laid out*\n",
    "across them with a `Mesh` and a `NamedSharding`, `device_put` the arrays\n",
    "onto that layout — parameters replicated (`P()`), the batch sharded along\n",
    "its leading axis (`P('data')`) — and `jit` the **unchanged** single-device\n",
    "training step. XLA's automatic-parallelization pass (GSPMD\n",
    "[@Xu.Lee.Chen.ea.2021]) partitions the computation to match the data\n",
    "layout and *inserts the very allreduce* that that section wrote\n",
    "by hand — you never write a collective:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "multi-gpu-practice-jax-annotate-the-layout-the-compiler-writes-the-collectives-1"
   },
   "source": [
    "def make_mesh(k):\n",
    "    return jax.make_mesh((k,), ('data',))\n",
    "\n",
    "def shard_batch(X, y, mesh):\n",
    "    P = jax.sharding.PartitionSpec('data')      # leading axis across devices\n",
    "    sh = jax.sharding.NamedSharding(mesh, P)\n",
    "    return jax.device_put(X, sh), jax.device_put(y, sh)\n",
    "\n",
    "def replicate(module, mesh):\n",
    "    \"\"\"Give every device of the mesh a full copy of a module's state.\"\"\"\n",
    "    sh = jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec())\n",
    "    nnx.update(module, jax.device_put(nnx.state(module), sh))\n",
    "\n",
    "# The training step is written for one device; jit + sharded inputs make it\n",
    "# data-parallel with no change to the body.\n",
    "@nnx.jit\n",
    "def train_step(model, opt, X, y):\n",
    "    def loss_fn(m):\n",
    "        return optax.softmax_cross_entropy_with_integer_labels(\n",
    "            m(X), y).mean()\n",
    "    loss, grads = nnx.value_and_grad(loss_fn)(model)\n",
    "    opt.update(model, grads)                     # XLA inserts the allreduce\n",
    "    return loss"
   ],
   "execution_count": 4,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=multi-gpu-practice-md-10 fw=all -->\n",
    "\n",
    "Now run it — a real `ResNet18`, a real optimizer, a real Fashion-MNIST\n",
    "batch (resized once to 64×64 and kept as host arrays). The reveal is\n",
    "`visualize_array_sharding`, which draws where a tensor actually lives: the\n",
    "batch split across the mesh before the step, and a weight *after* the step\n",
    "still replicated everywhere — the compiler's allreduce is what kept every\n",
    "copy identical:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "multi-gpu-practice-jax-annotate-the-layout-the-compiler-writes-the-collectives-2"
   },
   "source": [
    "imgs, labels = d2l.FashionMNIST().train         # raw arrays; resize once\n",
    "X_all = np.asarray(jax.image.resize(\n",
    "    jnp.float32(imgs[:, :, :, None]) / 255, (len(imgs), 64, 64, 1),\n",
    "    'bilinear'))\n",
    "y_all = np.asarray(labels, np.int32)\n",
    "\n",
    "k = min(4, jax.local_device_count())\n",
    "mesh = make_mesh(k)\n",
    "model = ResNet18(rngs=nnx.Rngs(0))\n",
    "opt = nnx.Optimizer(model, optax.sgd(0.1), wrt=nnx.Param)\n",
    "replicate(model, mesh); replicate(opt, mesh)\n",
    "Xs, ys = shard_batch(X_all[:256], y_all[:256], mesh)\n",
    "loss = train_step(model, opt, Xs, ys)           # one data-parallel step\n",
    "print(f'one sharded step: loss {float(loss):.2f}')\n",
    "print('the batch, split along its leading axis:')\n",
    "jax.debug.visualize_array_sharding(Xs.reshape(len(Xs), -1))\n",
    "print('a weight after the step, still replicated:')\n",
    "jax.debug.visualize_array_sharding(model.net.layers[-1].kernel[...])"
   ],
   "execution_count": 5,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "one sharded step: loss 2.47\nthe batch, split along its leading axis:\n"
     ]
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "\u001b[38;2;255;255;255;48;2;57;59;121m                                                                                \u001b[0m\n\u001b[38;2;255;255;255;48;2;57;59;121m                                     \u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121mGPU 0\u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121m                                      \u001b[0m\n\u001b[38;2;255;255;255;48;2;57;59;121m                                                                                \u001b[0m\n\u001b[38;2;255;255;255;48;2;222;158;214m                                                                                \u001b[0m\n\u001b[38;2;255;255;255;48;2;222;158;214m                                     \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214mGPU 1\u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m                                      \u001b[0m\n\u001b[38;2;255;255;255;48;2;222;158;214m                                                                                \u001b[0m\n\u001b[38;2;255;255;255;48;2;173;73;74m                                                                                \u001b[0m\n\u001b[38;2;255;255;255;48;2;173;73;74m                                     \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74mGPU 2\u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m                                      \u001b[0m\n\u001b[38;2;255;255;255;48;2;173;73;74m                                                                                \u001b[0m\n\u001b[38;2;0;0;0;48;2;181;207;107m                                                                                \u001b[0m\n\u001b[38;2;0;0;0;48;2;181;207;107m                                     \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107mGPU 3\u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m                                      \u001b[0m\n\u001b[38;2;0;0;0;48;2;181;207;107m                                                                                \u001b[0m\n"
     },
     "metadata": {}
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "a weight after the step, still replicated:\n"
     ]
    },
    {
     "output_type": "display_data",
     "data": {
      "text/plain": "\u001b[38;2;255;255;255;48;2;57;59;121m           \u001b[0m\n\u001b[38;2;255;255;255;48;2;57;59;121m           \u001b[0m\n\u001b[38;2;255;255;255;48;2;57;59;121m           \u001b[0m\n\u001b[38;2;255;255;255;48;2;57;59;121m           \u001b[0m\n\u001b[38;2;255;255;255;48;2;57;59;121m           \u001b[0m\n\u001b[38;2;255;255;255;48;2;57;59;121mGPU 0,1,2,3\u001b[0m\n\u001b[38;2;255;255;255;48;2;57;59;121m           \u001b[0m\n\u001b[38;2;255;255;255;48;2;57;59;121m           \u001b[0m\n\u001b[38;2;255;255;255;48;2;57;59;121m           \u001b[0m\n\u001b[38;2;255;255;255;48;2;57;59;121m           \u001b[0m\n\u001b[38;2;255;255;255;48;2;57;59;121m           \u001b[0m\n"
     },
     "metadata": {}
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=multi-gpu-practice-md-11 fw=all -->\n",
    "\n",
    "Take the receipt. that section wrote its `pmean` by hand; here\n",
    "nobody wrote a collective — so where is it? In the compiled program. Lower\n",
    "and compile the step, and search the XLA text for what GSPMD inserted:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "multi-gpu-practice-jax-annotate-the-layout-the-compiler-writes-the-collectives-3"
   },
   "source": [
    "hlo = train_step.lower(model, opt, Xs, ys).compile().as_text()\n",
    "ops = [line.strip() for line in hlo.splitlines()\n",
    "       if ' all-reduce(' in line or ' all-reduce-start(' in line]\n",
    "sizes = [int(re.search(r'\\[(\\d+)\\]', op).group(1)) for op in ops]\n",
    "print(f'{len(ops)} all-reduce ops in the compiled step; the largest '\n",
    "      f'sums {max(sizes) / 1e6:.1f}M floats:')\n",
    "print(ops[sizes.index(max(sizes))][:88] + ' ...')"
   ],
   "execution_count": 6,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "36 all-reduce ops in the compiled step; the largest sums 6.4M floats:\n%all-reduce-start.34 = f32[6447680]{0} all-reduce-start(%wrapped_concatenate.34), channe ...\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=multi-gpu-practice-md-12 fw=all -->\n",
    "\n",
    "Dozens of allreduces, not one — and reading them teaches. The two largest\n",
    "together carry the ~11M gradient floats (XLA bucketed them into two fused\n",
    "sums, much as DDP buckets); the many small ones are batch-norm statistics,\n",
    "because under `jit` even normalization is computed over the *global*\n",
    "batch — the compiler preserves single-device semantics exactly, where DDP\n",
    "leaves each replica its own batch-norm statistics unless you ask for\n",
    "`SyncBatchNorm`. (The exact count is a compiler artifact; the two big ones\n",
    "are the point.)\n",
    "\n",
    "Measured, under the same weak-scaling convention as the DDP sweep\n",
    "(per-device batch 256) — one process, no launcher, no sidecar files:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "multi-gpu-practice-jax-annotate-the-layout-the-compiler-writes-the-collectives-4"
   },
   "source": [
    "def batches(global_batch):\n",
    "    n = (len(X_all) // global_batch) * global_batch\n",
    "    for i in range(0, n, global_batch):\n",
    "        yield X_all[i:i + global_batch], y_all[i:i + global_batch]\n",
    "\n",
    "def jax_throughput(k, per_device_batch=256):\n",
    "    mesh = make_mesh(k)\n",
    "    model = ResNet18(rngs=nnx.Rngs(0))\n",
    "    opt = nnx.Optimizer(model, optax.sgd(0.1), wrt=nnx.Param)\n",
    "    replicate(model, mesh); replicate(opt, mesh)\n",
    "    B = per_device_batch * k\n",
    "    for epoch in range(2):                 # epoch 0 warms up and compiles\n",
    "        t0 = time.time(); n = 0\n",
    "        for X, y in batches(B):\n",
    "            loss = train_step(model, opt, *shard_batch(X, y, mesh))\n",
    "            n += X.shape[0]\n",
    "        loss.block_until_ready()           # completion timing\n",
    "        dt = time.time() - t0\n",
    "    return n / dt\n",
    "\n",
    "ks = [k for k in (1, 2, 4) if k <= jax.local_device_count()]\n",
    "tput = [jax_throughput(k) for k in ks]\n",
    "for k, t in zip(ks, tput):\n",
    "    print(f'{k} GPU(s): {t:.0f} samples/s, {t / tput[0]:.2f}x, '\n",
    "          f'{100 * t / tput[0] / k:.0f}% weak-scaling efficiency')"
   ],
   "execution_count": 7,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "1 GPU(s): 2670 samples/s, 1.00x, 100% weak-scaling efficiency\n2 GPU(s): 4384 samples/s, 1.64x, 82% weak-scaling efficiency\n4 GPU(s): 8001 samples/s, 3.00x, 75% weak-scaling efficiency\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=multi-gpu-practice-md-13 fw=all -->\n",
    "\n",
    "The shape matches the DDP sweep — real gains, efficiency sagging as $k$\n",
    "grows, no cliff — with one instructive difference: the compiled XLA step\n",
    "is faster per device, so the same allreduce is a *larger fraction* of each\n",
    "step and the efficiency reads a few points lower (82% already at two GPUs\n",
    "in our runs). That is the equation again: speed up the compute and\n",
    "the fabric gets *relatively* slower. And note what these numbers are not:\n",
    "this loop feeds pre-staged host arrays while the DDP script pays a real\n",
    "`DataLoader`, so the absolute samples/s are not a framework shoot-out —\n",
    "the scaling curve is the comparison.\n",
    "\n",
    "One more check, because we are trusting a collective nobody wrote.\n",
    "that section ended its hand-built loop with a one-step equality\n",
    "test, and the declarative version deserves the same scrutiny — same\n",
    "initialization, same 256-example batch, one step on one device versus one\n",
    "step sharded across two:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "multi-gpu-practice-jax-annotate-the-layout-the-compiler-writes-the-collectives-5"
   },
   "source": [
    "def one_step_delta(k):\n",
    "    mesh = make_mesh(k)\n",
    "    model = ResNet18(rngs=nnx.Rngs(0))\n",
    "    opt = nnx.Optimizer(model, optax.sgd(0.1), wrt=nnx.Param)\n",
    "    replicate(model, mesh); replicate(opt, mesh)\n",
    "    before = jax.tree.map(np.asarray, nnx.state(model, nnx.Param))\n",
    "    train_step(model, opt, *shard_batch(X_all[:256], y_all[:256], mesh))\n",
    "    after = jax.tree.map(np.asarray, nnx.state(model, nnx.Param))\n",
    "    return jax.tree.map(lambda a, b: a - b, after, before)\n",
    "\n",
    "d1 = one_step_delta(1)\n",
    "d2 = one_step_delta(min(2, jax.local_device_count()))\n",
    "gaps = jax.tree.map(lambda a, b: np.abs(a - b).max(), d1, d2)\n",
    "print('max update difference, k=1 vs k=2: '\n",
    "      f'{max(jax.tree.leaves(gaps)):.1e}')"
   ],
   "execution_count": 8,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "max update difference, k=1 vs k=2: 1.5e-04\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=multi-gpu-practice-md-14 fw=all -->\n",
    "\n",
    "The two runs agree to about $10^{-4}$ against update magnitudes near\n",
    "$10^{-2}$ — not the $10^{-9}$ of that section's LeNet check,\n",
    "and the gap is itself informative: these are two *different compiled\n",
    "programs*, whose tf32 convolutions and batch-norm reductions associate\n",
    "their arithmetic differently, so the residue is rounding. What the check\n",
    "exists to catch — a summed-instead-of-averaged gradient, an accidental\n",
    "$k\\times$ learning rate — would announce itself at the scale of the update\n",
    "itself.\n",
    "\n",
    "And here is the punchline. To move between *data* parallelism, *tensor*\n",
    "parallelism, and *FSDP*-style sharding in JAX, you change the\n",
    "`PartitionSpec` — not the model code. Sharding the batch axis gives data\n",
    "parallelism (above); sharding a weight's feature axis gives tensor\n",
    "parallelism; sharding the parameters and letting XLA all-gather them\n",
    "just-in-time gives the FSDP pattern. **One sharding vocabulary — annotate\n",
    "the layout, the compiler writes the collectives — spans what PyTorch\n",
    "exposes as three different APIs** (DDP, tensor-parallel wrappers, FSDP) —\n",
    "though an effective sharding plan is still model-aware: someone has to\n",
    "choose the layouts and the constraint points. The manual end of the same\n",
    "spectrum is the `jax.shard_map` + `lax.psum` of that section,\n",
    "where you write the collective yourself; `jit` + sharding is the\n",
    "automatic end.\n",
    "\n",
    "| | PyTorch | JAX |\n",
    "|---|---|---|\n",
    "| processes | one per GPU (`torchrun`) | one, sees all GPUs |\n",
    "| collectives | explicit (NCCL, or DDP's buckets) | inserted by XLA/GSPMD |\n",
    "| data parallel | `DistributedDataParallel` | shard the batch axis |\n",
    "| tensor parallel | separate wrappers / DTensor | change the `PartitionSpec` |\n",
    "| sharded (FSDP) | `fully_shard` | change the `PartitionSpec` |\n",
    "| control | imperative, visible | declarative, compiler-driven |\n",
    "\n",
    "Neither deal is strictly better: PyTorch's explicitness makes the\n",
    "communication legible and debuggable; JAX's declarativeness makes the same\n",
    "code span parallelism strategies by editing an annotation. Knowing both is\n",
    "knowing the design space.\n",
    "\n",
    "## When One Node Is Not Enough\n",
    "\n",
    "This chapter stops at the boundary of a single machine, and it is worth\n",
    "naming what lies past it. When a model is too large for even sharded data\n",
    "parallelism on one node — when the parameters, or the batch, or the\n",
    "sequence length outgrow what $k$ local GPUs can hold — the answer is to\n",
    "shard across *machines*, combining data parallelism with tensor, pipeline,\n",
    "and expert parallelism into the \"3D parallelism\" of frontier training. The\n",
    "collectives then run over a network fabric measured in tens of GB/s\n",
    "between nodes rather than an NVLink domain within one, and the cost model\n",
    "of that section acquires a second, slower bandwidth term. That\n",
    "is the province of the Language Models part, which has models and datasets\n",
    "large enough to warrant it; the production library map — Megatron, the\n",
    "FSDP/DTensor stack, DeepSpeed, and how to launch and checkpoint them across\n",
    "a cluster — is that section. From here the communication\n",
    "*algebra* stays the same; what multi-node adds is engineering on top of\n",
    "it — a hierarchy of fabrics (NVLink inside a node, a network between\n",
    "nodes), rendezvous and elastic restart when machines fail, and stragglers\n",
    "that turn a synchronous step into a queueing problem.\n",
    "\n",
    "## Summary\n",
    "\n",
    "* Production data parallelism (DDP) fixes the three deficits of a\n",
    "  hand-rolled loop: one process per GPU (no shared GIL), ring/tree\n",
    "  collectives (no hub), and gradient bucketing that **overlaps\n",
    "  communication with the backward pass**.\n",
    "* Launched from a notebook via `torchrun` on a sidecar script, DDP scales\n",
    "  a compute-dense ResNet-18 by roughly 1.8× on two GPUs and 3.3× on four\n",
    "  on our host-staged box (88% and 82% *weak-scaling* efficiency; a\n",
    "  fixed-global-batch strong-scaling sweep asks the harder question), and a\n",
    "  `no_sync()` measurement confirms the equation's communication\n",
    "  price to within tens of percent. An NVLink box changes only the\n",
    "  constant.\n",
    "* The constant is also configuration-sensitive: one documented NCCL\n",
    "  switch quintuples the bare-collective bandwidth on this box — yet\n",
    "  deadlocks DDP's overlapped training path, so the notebooks keep the\n",
    "  library's defaults and teach the pair of measurements instead.\n",
    "  Collective-library configuration moves communication by factors, not\n",
    "  percent; a workaround is validated per platform *and* per workload.\n",
    "* FSDP shards the $16P$-byte training state across ranks — the ZeRO ladder\n",
    "  — by splitting allreduce back into its reduce-scatter and all-gather\n",
    "  halves and materializing each layer just-in-time. It is for models whose\n",
    "  training state outgrows one GPU — a few billion parameters on this card\n",
    "  class.\n",
    "* JAX inverts the model: one process, annotate the data layout with a\n",
    "  `Mesh` and `PartitionSpec`, `jit` the unchanged step, and XLA inserts\n",
    "  the collectives — visible in the compiled program's `all-reduce` ops,\n",
    "  measured in a k-sweep, and verified by a k=1-versus-k=2 equality check.\n",
    "  Changing the `PartitionSpec` — not the code — moves between data,\n",
    "  tensor, and FSDP-style sharding, spanning what PyTorch exposes as three\n",
    "  APIs.\n",
    "* Multi-node 3D parallelism and the production library map are the\n",
    "  Language Models part and that section.\n",
    "\n",
    "## Exercises\n",
    "\n",
    "1. Vary DDP's `bucket_cap_mb` (the gradient bucket size) and measure\n",
    "   throughput at $k = 2$. Why is there an optimum — what does a too-small\n",
    "   bucket cost, and a too-large one?\n",
    "1. Extend the `no_sync()` measurement to $k = 4$. How much does the\n",
    "   per-step communication time grow from $k=2$? Compare against\n",
    "   the equation's $2(k-1)/k$ growth and the flat $2N/\\beta$\n",
    "   of the equation — which fits better, and what does that tell\n",
    "   you about how NCCL schedules the transfer on this P2P-less box?\n",
    "1. Reproduce both halves of the fabric-configuration story. First extend\n",
    "   the bare-collective comparison across payloads from 1 MB to 256 MB:\n",
    "   where does each transport's effective bandwidth saturate, and does\n",
    "   the five-fold gap persist? Then set `NCCL_SHM_USE_CUDA_MEMCPY=1`\n",
    "   inside `train_ddp.py` itself and relaunch the $k=2$ sweep — on our\n",
    "   box the run wedges within seconds (be ready to kill the launcher).\n",
    "   What does the pair teach about trusting a library's defaults on a\n",
    "   topology it was not tuned for — and about trusting a workaround\n",
    "   anywhere you have not re-validated it, workload and all?\n",
    "1. Write the `PartitionSpec` that shards a weight matrix's *output*\n",
    "   features across the mesh (tensor parallelism) and\n",
    "   `visualize_array_sharding` the result. How does the communication\n",
    "   pattern differ from the batch-sharded (data-parallel) case?\n",
    "1. Size ZeRO stage 3 (parameters, gradients, and optimizer states all\n",
    "   sharded) for a 7-billion-parameter model on 8 GPUs with 80 GB each:\n",
    "   what is the per-GPU training footprint, and does it fit? Redo the\n",
    "   arithmetic for DDP (no sharding) and explain the difference.\n",
    "1. The weak-scaling efficiency in the DDP sweep fell from 88% at $k=2$ to\n",
    "   82% at $k=4$. Using the equation, predict the efficiency at\n",
    "   $k=8$ on the same fabric, and state the assumption your prediction\n",
    "   makes about how $t_{\\text{comm}}$ grows."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "name": "d2l-jax",
   "display_name": "d2l (JAX)",
   "language": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
