Distributed Training Fundamentals: Multi-GPU Essentials

Picture an engineer at an Indian AI lab training a Hindi-first language model — the kind of work teams like Sarvam AI build — on eight GPUs sitting in one server. The training batch is split into eight shards, one per GPU. Each GPU runs its own forward and backward pass on its shard and produces its own gradient vector — a list of numbers telling every parameter which direction to move. Here is the problem: GPU 0's gradient was computed from different sentences than GPU 7's gradient. If each GPU just applies its own gradient to its own copy of the model, the optimizer step on GPU 0 nudges the weights one way and the optimizer step on GPU 7 nudges them another way. After one step, there are eight different models. After a thousand steps, eight completely unrelated models. Data-parallel training only works because, before any optimizer step, every GPU's gradient is combined into a single agreed-upon gradient and every GPU receives an identical copy of that combined result. That operation — everyone contributes a value, everyone gets the same combined answer back — is called an all-reduce, and how you implement it is the difference between a training run that scales cleanly to dozens of GPUs and one that stalls out at four.

The naive fix: a parameter server

The obvious design is a central switchboard. Designate one machine — call it the parameter server (PS) — and have every worker GPU send it their local gradient. The PS sums the incoming gradients and sends the sum back to every worker. This was exactly how early distributed TensorFlow (2016–17) was commonly deployed, and it works correctly. It just doesn't scale, and the reason is a pure bandwidth argument, not a correctness one.

Let S be the size in bytes of one worker's full gradient (the flattened concatenation of every parameter's gradient tensor), and let there be N worker GPUs plus one dedicated PS. Every worker sends S bytes in and receives S bytes back — 2S of traffic per worker, which sounds fine. But look at the PS: it receives N separate S-byte messages and sends N separate S-byte messages, all through its own network interface. Total traffic through the PS is 2N·S. At N = 4 and S = 800 MB, that's 2 × 4 × 800 MB = 6.4 GB flowing through one NIC. Double the GPU count to N = 8 and the PS traffic doubles to 12.8 GB, while each individual worker's own traffic (2S) never changed. The PS's network link is now doing eight times the work of any single worker's link, and every worker sits idle waiting for it to catch up. This is not a hardware failure — it is the topology itself concentrating all N workers' traffic onto one node's bandwidth budget.

Ring all-reduce: spreading the load

The fix, formalized by Patarasuk and Yuan (Bandwidth Optimal All-Reduce Algorithms for Clusters of Workstations, 2009) and brought into deep learning practice by Baidu's Silicon Valley AI Lab in 2017 and popularized further by Sergeev and Del Balso's Horovod (2018), removes the central node entirely. Arrange the N GPUs in a logical ring: GPU i only ever talks to its two ring neighbors, GPU (i−1) mod N and GPU (i+1) mod N. Split each GPU's gradient vector into N equal chunks. The algorithm runs in two phases, each taking exactly N−1 steps.

Scatter-reduce phase. At each step, every GPU simultaneously sends one chunk to its right neighbor and receives one chunk from its left neighbor, adding the received chunk into its own matching chunk. The chunk that started at GPU i keeps moving one position further around the ring at every step, picking up a new addend at each hop. After N−1 hops, that chunk has visited every other GPU exactly once and now holds the true sum across all N GPUs — but it lives on only one GPU.

All-gather phase. The algorithm runs N−1 more steps, but now GPUs simply relay the already-fully-summed chunks around the ring without adding anything further, so that by the end every GPU holds every chunk's final sum. Total steps: 2(N−1). Total data moved per step, per GPU: exactly one chunk, of size S/N.

Why this is bandwidth-optimal — precisely

Here is the sharper version of the claim, the one worth getting exactly right rather than hand-waving. Sum the total bytes moved across the whole ring system: N GPUs, each active for 2(N−1) steps, each step moving S/N bytes, gives a system-wide total of N · 2(N−1) · (S/N) = 2(N−1)·S. Compare that to the parameter-server system-wide total of 2N·S. Ring already moves fewer total bytes system-wide — but that is not the headline result. The headline result is what happens per node. In the ring, each individual GPU's own send-plus-receive traffic over the entire algorithm is 2(N−1)/N · S. As N grows, that ratio climbs toward 2 and never exceeds it — at N = 4 it's 1.5S, at N = 100 it's 1.98S. No GPU, no matter how large the cluster gets, ever handles more than roughly 2S of traffic. Meanwhile the PS's traffic (2N·S) grows without bound. That is what "bandwidth-optimal" precisely means here: not that ring moves zero bytes, but that no single link or node is ever the bottleneck, so the busiest node's load stays flat as the cluster scales.

Bandwidth-optimal is not the same as latency-optimal, and it's worth naming the gap explicitly. The ring still needs 2(N−1) sequential steps, and each step costs a fixed round-trip latency before any bytes even start moving. For a 512-GPU cluster that's 1,022 hops of fixed overhead, which is why NVIDIA's NCCL library — the collective-communication engine most training frameworks call under the hood — doesn't use a pure ring for every message size. For large gradient buffers, where bandwidth dominates, it uses ring all-reduce exactly as described here. For small buffers, where the fixed per-hop latency dominates instead, it switches to a tree-shaped algorithm that needs only O(log N) hops instead of O(N). Ring wins the bandwidth game; trees win the latency game; production systems use both depending on message size.

Worked example: one full ring all-reduce round on four GPUs

Take N = 4 GPUs, ring order G0 → G1 → G2 → G3 → G0. Each GPU's local gradient is a 4-element vector, split into chunks a, b, c, d (index 0–3):

G0 = [a0,b0,c0,d0] = [2, 5, 1, 8]
G1 = [a1,b1,c1,d1] = [3, 0, 4, 2]
G2 = [a2,b2,c2,d2] = [1, 2, 6, 0]
G3 = [a3,b3,c3,d3] = [4, 1, 0, 3]

The target: every GPU should end with the column sums [10, 8, 11, 13] (a: 2+3+1+4=10, b: 5+0+2+1=8, c: 1+4+6+0=11, d: 8+2+0+3=13). Each step, every GPU sends its most recently updated chunk to its right neighbor and adds whatever it receives from its left neighbor into the matching chunk slot.

By step 3 (end of scatter-reduce), each GPU holds exactly one fully-summed chunk — G0 has the true b-sum (8), G1 has the true c-sum (11), G2 has the true d-sum (13), G3 has the true a-sum (10) — bolded values above. Steps 4–6 (all-gather) just relay those four finished sums around the ring without any further addition, and by step 6 every GPU independently holds the identical [10, 8, 11, 13]. Six steps, four GPUs, and no GPU ever transmitted more than one quarter of its gradient vector at a time.

Gradient bucketing: paying the per-message tax once

Ring all-reduce is bandwidth-efficient per byte, but a real model doesn't hand you one gradient vector — it hands you hundreds of separate tensors, one per layer parameter, and many of them are tiny (a LayerNorm gain vector might be a few hundred bytes). Every all-reduce call, regardless of size, pays a fixed overhead: kernel launch on the GPU, a network round-trip to coordinate the ring. If that fixed cost is α milliseconds and the transfer cost is β milliseconds per kilobyte, calling all-reduce separately on 300 small tensors means paying 300·α before a single useful byte moves. PyTorch's DistributedDataParallel avoids this by bucketing: it groups parameter gradients into buffers of roughly a fixed size (documented default around 25 MB) and issues one all-reduce per bucket instead of one per tensor — the mechanism described in Li et al.'s 2020 systems paper on PyTorch's distributed training internals. The simulation below makes the saving concrete with illustrative constants:

def allreduce_time(sizes_kb, alpha_ms=0.3, beta_ms_per_kb=0.001):
    return sum(alpha_ms + beta_ms_per_kb * s for s in sizes_kb)

def bucketed_time(buckets_kb, alpha_ms=0.3, beta_ms_per_kb=0.001):
    return sum(alpha_ms + beta_ms_per_kb * sum(b) for b in buckets_kb)

# 6 parameter tensors: 4 small norm/bias vectors + 2 large weight matrices (KB)
tensor_sizes_kb = [4, 4, 2, 1024, 1024, 8]
unbucketed = allreduce_time(tensor_sizes_kb)

# DDP-style grouping into 2 buckets
buckets = [[4, 4, 2, 8], [1024, 1024]]
bucketed = bucketed_time(buckets)

print(f"Unbucketed: {unbucketed:.3f} ms across {len(tensor_sizes_kb)} calls")
print(f"Bucketed: {bucketed:.3f} ms across {len(buckets)} calls")
print(f"Speedup: {unbucketed / bucketed:.2f}x")

Tracing it: unbucketed pays 6 × 0.3 ms = 1.8 ms of fixed overhead plus 0.001 × 2066 KB = 2.066 ms of transfer, totaling 3.866 ms across 6 calls. Bucketed pays only 2 × 0.3 ms = 0.6 ms of overhead (the transfer cost is unchanged at 2.066 ms, since the same bytes still have to move), totaling 2.666 ms across 2 calls. The printed output is Unbucketed: 3.866 ms across 6 calls, Bucketed: 2.666 ms across 2 calls, Speedup: 1.45x. Notice all of the saving comes from the four small tensors — bucketing the two already-large weight matrices together barely changes anything, because their cost was bandwidth-bound to begin with.

Bucketing also enables overlap with backward computation. Autograd fires each parameter's gradient-ready hook the moment that layer's backward computation finishes, which happens in reverse order — output layers first, input layers last. DDP attaches a hook to every parameter that marks it ready inside its bucket; once a bucket is fully ready, it launches that bucket's all-reduce asynchronously while backward is still computing gradients for earlier layers:

# Illustrative pseudocode for DDP's ready-hook trigger (not runnable as-is)
def make_hook(bucket, param_index):
    def hook(grad):
        bucket.mark_ready(param_index, grad)
        if bucket.all_ready():
            async_allreduce(bucket.flatten())  # overlaps with earlier layers' backward
    return hook

This is why bucket size is a real tuning knob, not just a batching convenience: too small and per-message overhead creeps back in; too large and the bucket doesn't fill (and doesn't launch) until backward has progressed further than necessary, shrinking the compute/communication overlap window.

Common misconception

Students who first hear "every GPU ends up with everyone else's gradient" often picture each GPU literally transmitting its full S-byte gradient to every other GPU, as if allreduce were N-1 separate full broadcasts. Under that mental model, per-GPU traffic would be (N−1)·S — for N=8 that's 7 times the local gradient size, which would make large clusters hopeless. That is not what happens. The chunking is the entire trick: no chunk larger than S/N ever crosses any single link, and it is only because each chunk is fully reduced by the time it finishes its scatter-reduce lap that the all-gather phase can distribute finished answers instead of raw data. The full gradient reconstitutes on every GPU, but no link ever carries more than one chunk at a time.

Active recall

Attempt each question before reading its answer.

1. Why does a ring GPU's total traffic stay near 2S regardless of N, while the parameter-server node's traffic grows as 2N·S?
2. N = 8 GPUs, S = 800 MB per GPU. Compute total send-plus-receive traffic for one ring GPU, and total traffic through a dedicated parameter server.
3. In the 4-GPU worked example, suppose GPU2's original c-chunk value (6) was actually a corrupted 60. Which final column sum is wrong, by how much, and after which step does the corruption first appear on a GPU other than G2? Does it ever reach all four GPUs?
4. True or false: during ring all-reduce, some link at some point must carry a GPU's complete, unchunked gradient. Justify your answer.
5. A transformer has 400 parameter tensors, most under 2 KB (biases, LayerNorm weights) alongside a handful of large embedding and attention matrices. Why does bucketing help here specifically, and what goes wrong if the bucket size is set far larger than any individual layer's gradient?
6. The bandwidth-optimal bound is 2(N−1)/N · S, approaching 2S as N → ∞. Does "bandwidth-optimal" mean ring all-reduce's wall-clock time is also independent of N? Explain what does still grow with N.

Answers

1. Every ring GPU only ever exchanges S/N-sized chunks with its two immediate neighbors, over 2(N−1) steps, giving 2(N−1)/N · S total — a ratio that is bounded above by 2 for every N. The parameter server, by contrast, is the single point every worker's full S bytes must pass through twice (once inbound, once outbound), so its load is literally N times a single worker's share: 2N·S, unbounded as N grows.
2. Ring: 2(N−1)/N · S = 2×7/8×800 MB = 1,400 MB ≈ 1.4 GB per GPU. Parameter server: 2N·S = 2×8×800 MB = 12,800 MB = 12.8 GB through the one PS node — about 9× the load any single ring GPU ever carries.
3. Only the c-column (index 2) is affected; a, b, d stay correct at 10, 8, 13. Tracing the corrupted value: step 1, G3 receives 60 instead of 6 (first appearance elsewhere, off by +54). Step 2, G0 inherits 61 instead of 7. Step 3, G1's "finished" c-sum becomes 65 instead of 11 — the scatter-reduce phase has already locked in the wrong answer. The all-gather phase then does exactly its job of broadcasting that (wrong) finished value everywhere: G2 gets 65 in step 4, G3 in step 5, G0 in step 6. By the end, all four GPUs hold c = 65. A single bad value on one GPU contaminates one coordinate on every replica — this is exactly why NaN/Inf gradients on a single worker corrupt an entire distributed training run, not just that worker's copy.
4. False. Every chunk transmitted is exactly S/N in size, never the full S-byte gradient. The full gradient exists complete only after all N chunks have separately traveled and been reassembled — no single transmission ever carries more than one chunk.
5. Bucketing helps because the 400 small tensors would otherwise force 400 separate all-reduce calls, each paying the same fixed per-call latency overhead (α) regardless of how little data it carries — that fixed cost dominates when tensors are only a few KB. Grouping them into a handful of buckets amortizes α across many tensors at once. If the bucket size is set far larger than needed, a bucket may not fill until backward computation has progressed through most or all of the layers, delaying when its all-reduce can even start — destroying the overlap between communication and the still-ongoing backward pass, and effectively serializing what should have been hidden behind compute.
6. No — bandwidth-optimal caps the amount of data any single node handles, not the number of communication steps. Wall-clock time still has a term proportional to 2(N−1), the sequential hop count in the ring, each hop paying a fixed latency cost regardless of chunk size. As N grows, the bandwidth term stays flat near 2S, but the latency term (2(N−1)·α) keeps growing — which is precisely why production libraries like NCCL abandon the pure ring for small messages and switch to a tree topology with O(log N) hops instead.

Think About It

Think about this: How would you explain distributed training fundamentals: multi-gpu essentials to a friend who has never seen a computer? What real-world analogy would you use? Imagine you had to build a system using these concepts — what would be your first step? Try this: before moving on, write down three things you learned and one question you still have.

Key Takeaways — Summary and Recap

Let us recap what we covered: the core ideas behind distributed training fundamentals: multi-gpu essentials, how they connect to real-world applications, and why they matter for your journey in computer science. Remember these key points as you move forward. For competitive exam preparation (CBSE, JEE, BITSAT), focus on understanding the WHY behind each concept, not just the WHAT.