{
 "cells": [
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "---\n",
    "execute:\n",
    "  enabled: false\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=multiple-gpus-md-1 fw=all -->\n",
    "\n",
    "# Multi-GPU from First Principles\n",
    "\n",
    "The ladder of that section ended with the one rung that\n",
    "buys more arithmetic and more *aggregate* memory at once: another GPU.\n",
    "This section\n",
    "adds it — but honestly, and from the ground up. We build data-parallel\n",
    "training by hand, with explicit device-to-device copies and a\n",
    "hand-rolled gradient sum, so that when that section\n",
    "replaces our loop with production machinery, you know exactly what that\n",
    "machinery does. We derive the communication algorithm the professionals\n",
    "use (ring allreduce) and its cost, and — because this book's build box\n",
    "has no fast inter-GPU fabric — we *measure* what communication actually\n",
    "costs and discover the central fact of parallel training: **a second GPU\n",
    "is not free, and whether it pays is an accounting question you can answer\n",
    "before you run.**\n",
    "\n",
    "That last point is why this section is built on a machine with no NVLink\n",
    "and no peer-to-peer transfer (that section). A datacenter\n",
    "fabric shrinks the constant in front of the communication term; it does\n",
    "not repeal the accounting — at frontier scale the same arithmetic decides\n",
    "how thousands of GPUs are spent, and it matters more, not less. Our slow\n",
    "fabric merely makes the cost loud enough to hear, which makes this box\n",
    "the better teacher. Every conclusion here holds at two GPUs as well as\n",
    "four — the number of devices is a variable, never a constant.\n",
    "\n",
    "*Prerequisites: minibatch SGD and the effect of batch size on the gradient\n",
    "estimate (*that section*); LeNet (*that section*); the\n",
    "interconnect measurements of* that section*.*"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "multiple-gpus-multi-gpu-from-first-principles"
   },
   "source": [
    "%matplotlib inline\n",
    "from d2l import jax as d2l\n",
    "import jax\n",
    "from jax import numpy as jnp\n",
    "import numpy as np\n",
    "import optax\n",
    "import os\n",
    "from functools import partial\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=multiple-gpus-md-2 fw=all -->\n",
    "\n",
    "## Three Ways to Split\n",
    "\n",
    "Given more than one GPU and a model to train, there are three ways to\n",
    "divide the work (the figure).\n",
    "\n",
    "![Three ways to split a model across two devices: replicate it and split\n",
    "the batch (data parallel); split the layer stack (pipeline parallel);\n",
    "split every layer's width (tensor parallel).](../img/mdl-perf-splitting.svg)\n",
    "\n",
    "**Data parallelism** replicates the whole model on every GPU and splits the\n",
    "*batch*: each device runs the full network on a different slice of the\n",
    "minibatch, then the devices sum their gradients so every replica takes the\n",
    "same optimizer step. It is the simplest, works for any model that fits on\n",
    "one GPU, and needs communication only once per step — so it is what this\n",
    "chapter builds. **Pipeline parallelism** splits the *layer stack* across\n",
    "devices, each holding a few consecutive layers; it lets a model too deep\n",
    "for one GPU fit, at the cost of tight inter-stage synchronization.\n",
    "**Tensor parallelism** splits *within* each layer — each device holds a\n",
    "slice of every weight matrix — and communicates several times per layer.\n",
    "The latter two matter only at scales this part defers to the Language\n",
    "Models chapters, which have models large enough to warrant them; a single\n",
    "historical note is that the very first of them appeared in 2012, when\n",
    "AlexNet was split across two GPUs simply because its weights did not fit\n",
    "in one card's 3 GB [@Krizhevsky.Sutskever.Hinton.2012]. Data\n",
    "parallelism is our subject, and — a warning that that section\n",
    "already made — it does *not* let you train a bigger model: every GPU still\n",
    "holds a full copy.\n",
    "\n",
    "## Data Parallelism by Hand\n",
    "\n",
    "Data-parallel training on $k$ GPUs is five steps per minibatch: split the\n",
    "batch into $k$ shards; run forward and backward on each shard against that\n",
    "device's copy of the parameters; sum the $k$ gradient sets so all devices\n",
    "agree; and let every device apply the same update\n",
    "(the figure). We build each piece for a small LeNet, then\n",
    "run it — and watch the second GPU buy us nothing.\n",
    "\n",
    "![One data-parallel step on two GPUs. Split the minibatch, run identical\n",
    "forward and backward passes on each replica, sum the gradients with an\n",
    "allreduce, then apply the identical update\n",
    "everywhere.](../img/mdl-perf-data-parallel.svg)\n",
    "\n",
    "We define LeNet from raw tensors (not a module) so that parameters,\n",
    "gradients, and their movement between devices are fully visible:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "multiple-gpus-data-parallelism-by-hand-1"
   },
   "source": [
    "scale = 0.01\n",
    "def new_params(key):\n",
    "    ks = jax.random.split(key, 4)\n",
    "    return [jax.random.normal(ks[0], (20, 1, 3, 3)) * scale, jnp.zeros(20),\n",
    "            jax.random.normal(ks[1], (50, 20, 5, 5)) * scale, jnp.zeros(50),\n",
    "            jax.random.normal(ks[2], (800, 128)) * scale, jnp.zeros(128),\n",
    "            jax.random.normal(ks[3], (128, 10)) * scale, jnp.zeros(10)]\n",
    "\n",
    "def lenet(params, X):\n",
    "    conv = lambda x, W: jax.lax.conv_general_dilated(\n",
    "        x, W, (1, 1), 'VALID', dimension_numbers=('NCHW', 'OIHW', 'NCHW'))\n",
    "    pool = lambda x: jax.lax.reduce_window(\n",
    "        x, 0.0, jax.lax.add, (1, 1, 2, 2), (1, 1, 2, 2), 'VALID') / 4.0\n",
    "    h1 = pool(jax.nn.relu(conv(X, params[0]) + params[1].reshape(1, -1, 1, 1)))\n",
    "    h2 = pool(jax.nn.relu(conv(h1, params[2]) + params[3].reshape(1, -1, 1, 1)))\n",
    "    h2 = h2.reshape(h2.shape[0], -1)\n",
    "    h3 = jax.nn.relu(jnp.dot(h2, params[4]) + params[5])\n",
    "    return jnp.dot(h3, params[6]) + params[7]"
   ],
   "execution_count": 2,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=multiple-gpus-md-3 fw=all -->\n",
    "\n",
    "The two operations that make it parallel are *broadcasting* parameters to\n",
    "each device and *summing* gradients across them. We build the PyTorch\n",
    "version by hand, because the hand-rolled version is the lesson; the\n",
    "`#@save`d `split_batch` (which chops a minibatch into per-device shards)\n",
    "rounds out the toolkit:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "multiple-gpus-data-parallelism-by-hand-2"
   },
   "source": [
    "def get_params(params, num_devices):\n",
    "    \"\"\"Data parallelism replicates the full parameter set on every device;\n",
    "    shard_map's P() in-spec (below) broadcasts one copy to each.\"\"\"\n",
    "    return params\n",
    "\n",
    "def split_batch(X, y, num_devices):\n",
    "    \"\"\"Reshape `X` and `y` onto a leading device axis of size num_devices.\"\"\"\n",
    "    assert X.shape[0] % num_devices == 0\n",
    "    reshape = lambda a: a.reshape(num_devices, -1, *a.shape[1:])\n",
    "    return reshape(X), reshape(y)"
   ],
   "execution_count": 3,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=multiple-gpus-md-4 fw=all -->\n",
    "\n",
    "`allreduce` is the heart of it, and the naive version above is\n",
    "deliberately clumsy: it gathers every device's gradient onto device 0,\n",
    "sums there, and broadcasts back — a *star* topology, with device 0 as a\n",
    "hub that handles all $k-1$ inbound and $k-1$ outbound transfers. Watch it\n",
    "work on two vectors living on two GPUs:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "multiple-gpus-data-parallelism-by-hand-3"
   },
   "source": [
    "# JAX writes the same allreduce as a collective inside shard_map (below,\n",
    "# a pmean: the sum divided by k); here we show the summed result it is\n",
    "# built on.\n",
    "data = jnp.stack([jnp.ones((1, 2)) * (i + 1)\n",
    "                  for i in range(min(2, jax.local_device_count()))])\n",
    "print('before:', data.tolist())\n",
    "print('summed:', jnp.broadcast_to(data.sum(0), data.shape).tolist())"
   ],
   "execution_count": 4,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "before: [[[1.0, 1.0]], [[2.0, 2.0]]]\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "summed: [[[3.0, 3.0]], [[3.0, 3.0]]]\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=multiple-gpus-md-5 fw=all -->\n",
    "\n",
    "The training step assembles the pieces. Each device computes its shard's\n",
    "gradient; we allreduce parameter by parameter; each device applies plain\n",
    "SGD. The whole thing runs in one Python process — nothing here needs\n",
    "multiple processes, because we move tensors explicitly:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "multiple-gpus-data-parallelism-by-hand-4"
   },
   "source": [
    "@partial(jax.jit, static_argnames=('lr', 'mesh'))\n",
    "def train_step(params, X, y, lr, mesh):\n",
    "    \"\"\"One data-parallel step: shard_map makes the pmean collective explicit.\n",
    "    `X`, `y` arrive with the batch sharded across devices (P('data')) and\n",
    "    `params` replicated (P()) -- see `train` below; shard_map hands each device\n",
    "    the full parameter replica and its own batch shard. pcast marks the replica\n",
    "    as this device's own local copy, so the gradient below is the shard's own;\n",
    "    pmean then averages the shard gradients across devices.\"\"\"\n",
    "    P = jax.sharding.PartitionSpec\n",
    "\n",
    "    def per_device(params, X, y):\n",
    "        def loss_fn(p):\n",
    "            logits = lenet(p, X[0])   # X[0]: strip the size-1 sharded axis\n",
    "            return optax.softmax_cross_entropy_with_integer_labels(\n",
    "                logits, y[0]).mean()\n",
    "        local = jax.lax.pcast(params, 'data', to='varying')  # my own copy\n",
    "        grads = jax.grad(loss_fn)(local)        # my shard's mean-loss gradient\n",
    "        grads = jax.lax.pmean(grads, 'data')    # The allreduce, in one line\n",
    "        return jax.tree.map(lambda p, g: p - lr * g, params, grads)\n",
    "\n",
    "    step = jax.shard_map(per_device, mesh=mesh,\n",
    "                         in_specs=(P(), P('data'), P('data')),\n",
    "                         out_specs=P())\n",
    "    return step(params, X, y)"
   ],
   "execution_count": 5,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=multiple-gpus-md-6 fw=all -->\n",
    "\n",
    "The two tabs make the same computation visible in two idioms. PyTorch\n",
    "moves gradients between devices with explicit `.to(device)` copies inside\n",
    "`allreduce`; JAX writes the collective as a single `jax.lax.pmean` inside\n",
    "a `jax.shard_map`, where the `PartitionSpec('data')` annotation tells XLA\n",
    "that the leading axis is sharded across devices. (Note the top-level\n",
    "`jax.shard_map` — the older `jax.pmap` is a compatibility shim as of JAX\n",
    "0.8 and we do not use it.) One subtlety in the JAX tab earns its line:\n",
    "under `shard_map` the parameters arrive *replicated*, and differentiating\n",
    "a replicated tensor makes JAX's autodiff insert a gradient *sum* of its\n",
    "own — the transpose of \"broadcast to every device\" is \"add up every\n",
    "device's contribution\", and a further explicit `psum` on top would then\n",
    "double-count. The `pcast(..., to='varying')` line opts out by declaring\n",
    "each replica device-local, so the code owns the collective: every device\n",
    "takes its shard's mean-loss gradient, and `pmean` averages them into\n",
    "exactly the gradient one device would compute on the whole batch. The\n",
    "collective is *visible in the code* in both tabs — exactly the point of\n",
    "building it by hand.\n",
    "\n",
    "Both tabs also make a strong claim — that $k$ devices take *the same*\n",
    "step one device would — and a claim like that deserves a check, not a\n",
    "promise. From the same initialization, on the same minibatch, one step on\n",
    "two GPUs must move the parameters exactly where one step on one GPU does.\n",
    "The test is cheap, and it is the test that catches normalization bugs —\n",
    "a summed instead of averaged gradient is an accidental $k\\times$ learning\n",
    "rate — that no accuracy curve reliably reveals:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "multiple-gpus-data-parallelism-by-hand-7"
   },
   "source": [
    "P = jax.sharding.PartitionSpec\n",
    "X = jax.random.normal(jax.random.PRNGKey(1), (256, 1, 28, 28))\n",
    "y = jax.random.randint(jax.random.PRNGKey(2), (256,), 0, 10)\n",
    "base = new_params(jax.random.PRNGKey(0))\n",
    "deltas = []\n",
    "for k in (1, min(2, jax.local_device_count())):\n",
    "    mesh = jax.make_mesh((k,), ('data',))\n",
    "    place = lambda a, spec: jax.device_put(\n",
    "        a, jax.sharding.NamedSharding(mesh, spec))\n",
    "    params = jax.tree.map(lambda p: place(p, P()), base)\n",
    "    Xs, ys = split_batch(X, y, k)\n",
    "    new = train_step(params, place(Xs, P('data')), place(ys, P('data')),\n",
    "                     0.2, mesh)\n",
    "    deltas.append([np.asarray(n - b) for n, b in zip(new, base)])\n",
    "print('max update difference, k=1 vs k=2: '\n",
    "      f'{max(np.abs(a - b).max() for a, b in zip(*deltas)):.2e}')"
   ],
   "execution_count": 6,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "E0720 20:17:16.713029   44902 cuda_timer.cc:88] Delay kernel timed out: measured time has sub-optimal accuracy. There may be a missing warmup execution, please investigate in Nsight Systems.\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "max update difference, k=1 vs k=2: 1.86e-09\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=multiple-gpus-md-7 fw=all -->\n",
    "\n",
    "The parameter updates agree to floating-point noise — around $10^{-9}$,\n",
    "against update magnitudes near $10^{-2}$. \"Mathematically identical\" is\n",
    "now a measurement, not a slogan. With correctness established, the\n",
    "training loop that wraps `train_batch`/`train_step` is the multi-GPU\n",
    "cousin of that section's. First, one GPU:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "multiple-gpus-data-parallelism-by-hand-5"
   },
   "source": [
    "def train(num_gpus, batch_size, lr):\n",
    "    P = jax.sharding.PartitionSpec\n",
    "    data = d2l.FashionMNIST(batch_size=batch_size)\n",
    "    train_iter = data.get_dataloader(train=True)\n",
    "    mesh = jax.make_mesh((num_gpus,), ('data',))\n",
    "    place = lambda a, spec: jax.device_put(\n",
    "        a, jax.sharding.NamedSharding(mesh, spec))\n",
    "    params = get_params(new_params(jax.random.PRNGKey(0)), num_gpus)\n",
    "    params = jax.tree.map(lambda p: place(p, P()), params)  # a replica per device\n",
    "    timer, num_epochs = d2l.Timer(), 5\n",
    "    for epoch in range(num_epochs):\n",
    "        timer.start()\n",
    "        for X, y in train_iter:\n",
    "            X = jnp.array(X).transpose(0, 3, 1, 2)\n",
    "            if X.shape[0] % num_gpus:       # drop a ragged final batch\n",
    "                continue\n",
    "            Xs, ys = split_batch(X, jnp.array(y), num_gpus)\n",
    "            Xs, ys = place(Xs, P('data')), place(ys, P('data'))  # shard the batch\n",
    "            params = train_step(params, Xs, ys, lr, mesh)\n",
    "        jax.block_until_ready(params)\n",
    "        timer.stop()\n",
    "    apply_fn, hits, total = jax.jit(lenet), 0, 0\n",
    "    p0 = jax.tree.map(lambda p: jax.device_put(p, jax.devices()[0]), params)\n",
    "    for X, y in data.get_dataloader(train=False):\n",
    "        X = jnp.array(X).transpose(0, 3, 1, 2)\n",
    "        hits += int((apply_fn(p0, X).argmax(axis=1) == jnp.array(y)).sum())\n",
    "        total += len(y)\n",
    "    print(f'test acc {hits / total:.2f}, {timer.avg():.2f} sec/epoch '\n",
    "          f'on {num_gpus} GPU(s)')\n",
    "\n",
    "train(num_gpus=1, batch_size=256, lr=0.2)"
   ],
   "execution_count": 7,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "test acc 0.79, 0.56 sec/epoch on 1 GPU(s)\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=multiple-gpus-md-8 fw=all -->\n",
    "\n",
    "Now two GPUs, same batch size and learning rate. The optimization is\n",
    "mathematically identical, as the one-step check just proved — any\n",
    "accuracy gap between these short runs is run-to-run noise from\n",
    "initialization and rounding, not the parallelism — but the wall-clock is\n",
    "**no better**:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "multiple-gpus-data-parallelism-by-hand-6"
   },
   "source": [
    "train(num_gpus=min(2, jax.local_device_count()), batch_size=256, lr=0.2)"
   ],
   "execution_count": 8,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "test acc 0.81, 1.00 sec/epoch on 2 GPU(s)\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=multiple-gpus-md-9 fw=all -->\n",
    "\n",
    "The second GPU buys **no speedup** — on our box the two-GPU run is at\n",
    "best on par with one GPU, and in JAX outright slower — and this is not a\n",
    "bug, it is the syllabus.\n",
    "The reason is *not* slow communication: LeNet's parameters are tiny — the\n",
    "whole gradient set is about half a megabyte — so its allreduce is\n",
    "negligible, and we measure the transport directly below. The reason is\n",
    "that LeNet is simply too small\n",
    "to parallelize. Splitting a 256-example batch into two 128-example shards\n",
    "underfeeds each GPU — a convolution on half the batch does not run at half\n",
    "the time, because a small batch never filled the device to begin with —\n",
    "and the Python orchestration and per-step dispatch are not amortized\n",
    "by microsecond-scale compute. The technique is not wrong; the *regime* is\n",
    "wrong. This tiny model is the worst case for data parallelism, and\n",
    "diagnosing exactly why — separating the communication cost, which is small\n",
    "here, from the underutilization cost, which is not — is the next two\n",
    "subsections' work.\n",
    "\n",
    "## Doing Better: Ring Allreduce\n",
    "\n",
    "Our star `allreduce` has an obvious flaw: device 0 is a hub through which\n",
    "everything passes, so it moves $(k-1)N$ bytes in and $(k-1)N$ out for a\n",
    "parameter set of $N$ bytes, and the other devices' links sit idle while it\n",
    "works. The algorithm the professionals use — *ring allreduce*\n",
    "[@Patarasuk.Yuan.2009] — removes the hub entirely, and it is worth\n",
    "deriving because the same identity reappears as the seed of FSDP in\n",
    "that section.\n",
    "\n",
    "Arrange the $k$ devices in a ring, each talking only to its neighbor.\n",
    "Split each device's gradient vector into $k$ equal chunks. The allreduce\n",
    "runs in two phases (the figure). In **reduce-scatter**,\n",
    "over $k-1$ steps, device $i$ sends one chunk to its neighbor while\n",
    "receiving and accumulating another; after the phase, each device owns one\n",
    "fully-summed chunk (a different chunk per device). In **all-gather**,\n",
    "another $k-1$ steps pass those finished chunks around the ring until every\n",
    "device has all $k$. At every step, every link carries exactly one chunk of\n",
    "size $N/k$, so each device sends $(k-1) \\cdot N/k$ bytes per phase, and\n",
    "both phases together move\n",
    "\n",
    "$$\n",
    "\\frac{2(k-1)}{k}\\, N \\quad\\text{bytes per device,}\n",
    "$$\n",
    "\n",
    "which is **nearly constant in $k$**: it rises from $N$ at $k=2$ toward\n",
    "its bound of $2N$ and never exceeds it, however many devices join the\n",
    "ring. The star pattern's hub moved $(k-1)N$ — linear in $k$; the ring\n",
    "moves at most $2N$. That is why production collective libraries are built\n",
    "on rings and trees, chosen per message size and topology, alongside\n",
    "hardware-specific schemes; for messages small enough that per-hop\n",
    "*latency* dominates the byte count, trees win — a term this byte-only\n",
    "accounting deliberately ignores.\n",
    "\n",
    "![Ring allreduce on four GPUs, computed step by step. Reduce-scatter\n",
    "accumulates one chunk per device around the ring; all-gather then\n",
    "distributes the finished chunks. Each cell counts how many of the four\n",
    "contributions a chunk holds; a full chunk (4) is\n",
    "complete.](../img/mdl-perf-ring-allreduce.svg)\n",
    "\n",
    "The catch, on our box, is that the elegant $2(k-1)/k$ accounting assumes\n",
    "the links are the bottleneck — and they are, but *which* links? With no\n",
    "peer-to-peer transfer, every \"neighbor to neighbor\" hop is really a\n",
    "round trip through host memory (that section), so the\n",
    "ring's theoretical advantage over the star is largely erased: the\n",
    "transport, not the topology, is the ceiling. This is the theory-versus-\n",
    "practice lesson in miniature — NCCL will still pick a ring or tree per\n",
    "message size, but on this hardware the constant in front of $N$ is what\n",
    "hurts, and no algorithm fixes a slow wire.\n",
    "\n",
    "## The Accounting\n",
    "\n",
    "We can now answer the question data parallelism always poses — *does the\n",
    "next GPU pay?* — with a cost model rather than a guess. One step on $k$\n",
    "GPUs takes roughly\n",
    "\n",
    "$$\n",
    "t_{\\text{step}}(k) \\;\\approx\\;\n",
    "\\underbrace{t_{\\text{compute}}(B/k)}_{\\text{forward+backward on a shard}}\n",
    "\\;+\\;\n",
    "\\underbrace{t_{\\text{comm}}}_{\\approx\\, 2N / \\beta},\n",
    "$$\n",
    "\n",
    "where $B$ is the global batch, $N$ the bytes of gradients to reduce, and\n",
    "$\\beta$ the achievable allreduce bandwidth (the equation gives\n",
    "the $2N$; the per-device traffic is bounded by it for any $k$). The\n",
    "compute term *shrinks* with more GPUs (each does less of the batch); the\n",
    "communication term does not. Parallelism pays exactly when the compute\n",
    "you offload exceeds the communication you take on — models with high\n",
    "compute per byte of gradient traffic, large per-device batches, and fast\n",
    "links all push in your favor; a tiny model on a slow link, like LeNet on\n",
    "our box, is the case where it never pays.\n",
    "\n",
    "Two scaling conventions deserve names here, because every published\n",
    "speedup quietly picks one. the equation holds the global batch $B$\n",
    "fixed while $k$ grows — *strong scaling*: the same work, hopefully\n",
    "finishing sooner. The alternative holds the *per-device* batch fixed, so\n",
    "the global batch grows to $kB_{\\textrm{device}}$ — *weak scaling*: more\n",
    "work per unit time, quoted as throughput. The throughput sweep of\n",
    "that section holds the per-rank batch fixed and must\n",
    "be read under the weak-scaling convention — and remember that a grown\n",
    "global batch changes the optimization trajectory itself\n",
    "(that section).\n",
    "\n",
    "Let's plug in real numbers by measuring $\\beta$ directly. One bookkeeping\n",
    "convention first, so the numbers below are comparable: we quote *bytes\n",
    "moved per device per allreduce* ($\\approx 2N$ for the ring) divided by\n",
    "wall-clock time — an effective bandwidth. NCCL's benchmark \"busbw\"\n",
    "normalizes differently, so its figures are not directly comparable to\n",
    "ours. Here is the effective bandwidth of the allreduce each tab actually\n",
    "runs — the hand-rolled star in PyTorch, XLA's `psum` collective in JAX:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "multiple-gpus-the-accounting"
   },
   "source": [
    "if jax.local_device_count() >= 2:\n",
    "    mesh = jax.make_mesh((2,), ('data',))\n",
    "    P = jax.sharding.PartitionSpec\n",
    "    N = 64 * 1024 * 1024\n",
    "    x = jax.device_put(jnp.ones((2, N)),\n",
    "                       jax.sharding.NamedSharding(mesh, P('data')))\n",
    "    psum = jax.jit(jax.shard_map(\n",
    "        lambda a: jax.lax.psum(a, 'data'), mesh=mesh,\n",
    "        in_specs=P('data'), out_specs=P('data')))\n",
    "    t = d2l.Benchmark(lambda: psum(x), warmup=2, repeats=5).time\n",
    "    print(f'psum {2 * N * 4 / t / 1e9:.2f} GB/s effective over {1000*t:.1f} ms')\n",
    "else:\n",
    "    print('needs 2 GPUs')"
   ],
   "execution_count": 9,
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "psum 6.53 GB/s effective over 82.3 ms\n"
     ]
    }
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=multiple-gpus-md-10 fw=jax -->\n",
    "\n",
    "The measured `psum` sustains well under ten GB/s in our convention —\n",
    "more than two orders of magnitude below an NVLink domain's ~1.8 TB/s per\n",
    "GPU (the table). That number is itself a theory-versus-practice\n",
    "lesson: XLA hands the collective to NCCL, whose ring/tree chunking is\n",
    "tuned for peer-to-peer fabrics, and on our P2P-less box every \"neighbor\"\n",
    "hop stages through host memory — so the tuned library lands well *below*\n",
    "the roughly twenty GB/s a naive one-shot `.to()` copy achieves (the\n",
    "PyTorch tab measures it). The wire is not the culprit — both paths cross\n",
    "the same PCIe links. The gap is in *how* the fallback transport moves\n",
    "bytes (a latency-bound GPU-kernel copy by default, rather than the DMA\n",
    "copy engines), and a single NCCL environment switch recovers most of it\n",
    "in bare collectives (the PyTorch tab takes this apart, and\n",
    "that section measures the workaround — and its\n",
    "limits). The ceiling\n",
    "was one stage of the transport — found, as always in this chapter, by\n",
    "measuring."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- d2l:prose id=multiple-gpus-md-11 fw=all -->\n",
    "\n",
    "Now read the cost model against what we measured. LeNet's\n",
    "parameters are tiny, so $2N/\\beta$ is a fraction of a millisecond — the\n",
    "communication term is *not* what denies the speedup. The culprit is the\n",
    "other term: $t_{\\text{compute}}(B/k)$ does not actually fall like $1/k$\n",
    "for a small model, because halving an already-small batch leaves each GPU\n",
    "underutilized, so $t_{\\text{compute}}(B/2) \\approx t_{\\text{compute}}(B)$\n",
    "and the second GPU does redundant-feeling work for no wall-clock gain. The\n",
    "model pays off only when compute genuinely scales with the batch — a\n",
    "compute-dense network with a large per-device batch, where\n",
    "$t_{\\text{compute}}(B/k) \\approx t_{\\text{compute}}(B)/k$ dominates the\n",
    "small $t_{\\text{comm}}$. That is exactly the regime\n",
    "that section moves to, and where the second GPU finally\n",
    "earns its keep.\n",
    "\n",
    "A closing word of history, because it names the lineage. Before\n",
    "synchronous ring allreduce won for dense training, large-scale learning\n",
    "organized this same communication through *parameter servers*: dedicated\n",
    "storage nodes exposed a `push` (accumulate my gradient) and `pull`\n",
    "(give me the current sum) interface, sharding the parameters across\n",
    "servers so aggregation bandwidth scaled with the fleet\n",
    "[@Li.Andersen.Park.ea.2014]. That push/pull abstraction is the\n",
    "authors' own lineage, and it lives on today in recommender-system\n",
    "embedding tables (the territory of that section), where the\n",
    "parameters are too large and too sparse for every worker to hold a full\n",
    "replica — and, more broadly, wherever sparse or asynchronous state\n",
    "dominates the traffic. For dense training, collectives won; the modern\n",
    "production map is that section.\n",
    "\n",
    "## Summary\n",
    "\n",
    "* Data parallelism replicates the model, splits the batch, sums gradients\n",
    "  with an allreduce, and applies the identical update on every device.\n",
    "  It is the simplest form of multi-GPU training and does not enlarge the\n",
    "  model that fits.\n",
    "* Built by hand, the gradient sum is a star: device 0 is a hub moving\n",
    "  $(k-1)N$ bytes each way. Ring allreduce removes the hub — reduce-scatter\n",
    "  then all-gather — moving $2(k-1)/k \\cdot N$ bytes per device: nearly\n",
    "  constant, bounded by $2N$ however many devices join. That identity\n",
    "  reappears as the basis of FSDP.\n",
    "* The cost model $t_{\\text{step}}(k) \\approx t_{\\text{compute}}(B/k) +\n",
    "  2N/\\beta$ decides whether a GPU pays — but only when compute genuinely\n",
    "  scales with the batch. It is a *strong-scaling* statement (fixed global\n",
    "  batch); *weak scaling* fixes the per-device batch instead, and every\n",
    "  quoted speedup picks one convention. LeNet is the deliberate worst\n",
    "  case: its gradients\n",
    "  are tiny (communication negligible), yet a second GPU buys no speedup\n",
    "  because halving an already-small batch underutilizes each device, so\n",
    "  $t_{\\text{compute}}(B/k)$ does not fall. A one-step k=1-versus-k=2\n",
    "  equality check *proves* the parallel step is the same mathematics before\n",
    "  any timing is read. Inter-GPU bandwidth on this P2P-less box is\n",
    "  PCIe-limited (tens of GB/s for a raw copy; far less through NCCL's\n",
    "  fallback transport at its defaults — a documented switch recovers most\n",
    "  of it in bare collectives, and that section\n",
    "  measures both the workaround and its limits) — two to three orders of\n",
    "  magnitude below an NVLink domain.\n",
    "* Parameter servers (push/pull) organized this communication for the\n",
    "  asynchronous, multi-machine era; synchronous collectives won for dense\n",
    "  training, and the pattern survives in recsys embedding systems and\n",
    "  wherever sparse or asynchronous state dominates.\n",
    "\n",
    "## Exercises\n",
    "\n",
    "1. Extend the hand-rolled version to $k = 4$ (guard with `d2l.num_gpus()`)\n",
    "   and measure `sec/epoch` at $k \\in \\{1, 2, 4\\}$. Does the slowdown grow,\n",
    "   shrink, or hold? Explain using the equation.\n",
    "1. Implement ring allreduce with explicit `.to()` copies (reduce-scatter\n",
    "   then all-gather) and test whether it beats the star `allreduce` on our\n",
    "   host-staged box. It barely can — explain why in terms of\n",
    "   that section.\n",
    "1. Compute the ring's per-device traffic (the equation) for\n",
    "   ResNet-18's ~11M parameters in fp32, and, using your measured $\\beta$,\n",
    "   predict $t_{\\text{comm}}$ per step. that section\n",
    "   will measure the real thing — how close is your prediction?\n",
    "1. A thought experiment: sending gradients in bf16 instead of fp32 halves\n",
    "   $N$ and hence $t_{\\text{comm}}$. What could break, and which term of\n",
    "   the equation does it help — the one that shrinks with $k$ or the\n",
    "   one that does not?\n",
    "1. Scale the batch size with $k$ (from $B$ to $kB$) so each device keeps a\n",
    "   full $B$-sized shard. Now which term of the equation dominates,\n",
    "   and does the second GPU pay? Relate your answer to the batch-size\n",
    "   discussion of that section."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "name": "d2l-jax",
   "display_name": "d2l (JAX)",
   "language": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
