Expert Parallelism at Scale

The previous four sections built DeepSeekMoE on a single device: 256 routed experts, 1 shared expert, top-8 gating. The arithmetic is clean on paper. It is impossible on hardware. A single H800 holds 80 GB; a DeepSeek-V3 MoE layer alone weighs hundreds of gigabytes. The expert pool simply does not fit on one GPU. Expert parallelism — sharding the experts across many GPUs and shuttling tokens to wherever their chosen expert lives — is the engineering trick that turns the beautiful single-device math into a 671B-parameter system that actually trains.

The bet of expert parallelism. Stop trying to fit every expert on every GPU. Put each expert on exactly one GPU, and instead mail the tokens to the expert. You pay for two all-to-all collectives per layer; in exchange you get a model 30× bigger than any single accelerator could hold.

The Memory Wall of 671B Parameters

Start with a concrete accounting. DeepSeek-V3 has 256 routed experts, each a 2-layer FFN with hidden dimension $d_{ff} = 2048$ and model dimension $d = 7168$. That is roughly $2 \cdot d \cdot d_{ff} \approx 2.9 \cdot 10^7$ parameters per expert, or about 58 MB at BF16. One layer's 256 experts therefore weigh $256 \cdot 58\,\text{MB} \approx 15\,\text{GB}$. Multiply across 58 MoE layers and you reach roughly 880 GB of expert parameters alone — before counting attention, embeddings, the optimizer states (which are usually 4× the parameter footprint), gradients, or activations.

An 80 GB H800 cannot hold a single MoE layer's experts in BF16, let alone the whole model with optimizer state. Pure data parallelism is dead on arrival: replicating an 880 GB expert pool on every GPU is not an arithmetic problem, it is a physics problem. Tensor parallelism (the Megatron-style trick of splitting each matmul across devices) helps for a single dense FFN, but it does not exploit the structural sparsity of MoE — it shards each expert across GPUs even though most experts do nothing for most tokens.

Sharding Experts Across the Cluster

Picture a hospital again, but now spread across four cities. Each city has two specialists; eight specialists in total. When a patient arrives at a local clinic, the triage desk decides which specialist they need, and the patient flies to the city where that specialist works. After the appointment they fly home. The clinics are tiny (one router each). The specialists are stationary. The patients travel.

The four cities are GPUs. The eight specialists are experts (two per GPU). The patients are tokens. The flight is the all-to-all collective: in a single coordinated step, every GPU sends its tokens to wherever they need to go and receives the tokens it must process — all at the same time, all over the network. After expert compute, a reverse all-to-all flies the outputs home.

Two pieces of the system stay replicated, not sharded: the router and the shared experts. The router is small (a single $d \times E$ linear) and must run on every device — every device decides where its own tokens go without asking anyone. The shared experts, similarly small relative to the routed pool, live on every device and run locally for every token. The expensive, sharded thing is the routed pool.

The Math: Permutations and All-to-All

Let $W$ be the world size of the expert-parallel group, $E$ the number of routed experts, and assume $E$ is divisible by $W$. Each device $r \in \{0, \dots, W-1\}$ owns the contiguous expert block $[r \cdot E/W,\; (r+1) \cdot E/W)$. Define $\mathrm{owner}(e) = \lfloor e / (E/W) \rfloor$.

On step $t$, device $r$ holds a local token slice of shape $(T_r, d)$. After routing, it has top-$k$ expert ids per token, which we flatten into a list of $T_r \cdot k$ dispatch decisions. Group these by destination device, producing per-destination send counts $c_{r \to r'}$. The dispatch all-to-all is the permutation:

$\text{recv}_{r'} = \bigsqcup_{r=0}^{W-1} \text{send}_{r \to r'}, \quad |\text{send}_{r \to r'}| = c_{r \to r'}$

Each device $r'$ then runs its local experts on the received rows — no communication during compute. The combine all-to-all is the inverse permutation, mailing each output back to its source device. After applying the gates and summing across the $k$ slots, the per-token output is restored.

Communication cost

Each device sends about $T_r \cdot k \cdot d$ elements during dispatch and the same again during combine. The total volume per layer per device is therefore roughly $2 \cdot T_r \cdot k \cdot d \cdot b$ bytes (where $b$ is bytes per element — 2 for BF16). For DeepSeek-V3 with $T_r = 4096, k = 8, d = 7168, b = 2$, that's about 0.94 GB per device per layer per step, and we have 58 such layers. The all-to-all traffic is the single largest bandwidth consumer in MoE training.

The load-balance precondition

The math above assumes $c_{r \to r'}$ is roughly uniform. If one device's experts are popular and the rest are dead, that device's inbox overflows while the others sit idle — compute and communication both stall. Expert parallelism therefore requires the load-balancing mechanisms we will derive in chapter 6 (the auxiliary-loss-free bias terms). Without balance, the whole sharding strategy collapses.

Manual Numerical Walkthrough

Let us trace a tiny dispatch by hand. Four devices, eight experts (two per device), top-1 routing, three tokens per device.

Visualizing the Dispatch

Step through the five stages below. Watch how the colored tokens (color = chosen expert id) leave their home GPU during Dispatch, land next to the matching expert on the destination GPU during Compute, and travel back during Combine. The side panel tracks how many tokens actually had to cross the network — the rest picked a local expert by luck of the router.

Three observations to anchor. First, every GPU sends and receives in the same collective — that is what "all-to-all" means and why it is bandwidth-efficient. Second, the traffic counter exposes the load-balance pressure: in a uniformly balanced batch, every device sends and receives the same number of rows. Third, the compute step happens entirely locally — once tokens land, the experts run normally, identical to the single-device MoE from earlier sections.

Plain Python: A Four-Device Simulator

Before touching torch.distributed, here is the entire mechanism in NumPy. We simulate four devices as four Python lists, route every device's tokens locally, build the send buffers, and assemble each device's inbox by hand. The collective is the dictionary rearrangement on lines 36–43 — that is, literally, what an all-to-all does.

If local_expert_ids[0] = [5, 2, 0], local_dst[0] = [owner(5), owner(2), owner(0)] = [2, 1, 0].

1import numpy as np
2
3# 4 simulated devices, 8 experts (2 per device), top-k = 1.
4W, E, k = 4, 8, 1
5experts_per_device = E // W
6def owner(expert_id):
7    return expert_id // experts_per_device
8
9# Each device holds 3 tokens of dimension d=4 (its local batch slice).
10d = 4
11rng = np.random.default_rng(0)
12local_x = [rng.standard_normal((3, d)) for _ in range(W)]
13
14# Router weights are replicated on every device.
15W_router = rng.standard_normal((E, d)) * 0.1
16
17# Expert FFNs are SHARDED: each device only stores its own.
18local_experts = {
19    e: (rng.standard_normal((d, d)) * 0.1)
20    for e in range(E) if owner(e) == 0  # device 0's experts
21}
22# (In real code each device would only build its own dict.)
23
24def route(x):
25    """Pick the top-1 expert id for each row of x."""
26    logits = x @ W_router.T            # (T, E)
27    return logits.argmax(axis=1)       # (T,)
28
29# ---- 1. ROUTE (local, no comms) ----
30local_expert_ids = [route(x) for x in local_x]
31local_dst = [
32    np.array([owner(e) for e in ids]) for ids in local_expert_ids
33]
34
35# ---- 2. DISPATCH = all-to-all of tokens ----
36# Build, for each (src, dst) pair, the rows src wants to send to dst.
37send = {(s, t): [] for s in range(W) for t in range(W)}
38for s in range(W):
39    for row, dst in enumerate(local_dst[s]):
40        send[(s, dst)].append((local_x[s][row], local_expert_ids[s][row], row))
41
42# Each device's inbox = concat of what every other device sent IT.
43inbox = [
44    [pkt for s in range(W) for pkt in send[(s, t)]] for t in range(W)
45]
46
47# ---- 3. EXPERT COMPUTE (local on each device) ----
48# We simulate device 0 only — the others run the same code with their experts.
49outputs_dev0 = []
50for x_row, eid, src_row in inbox[0]:
51    W_e = local_experts[eid]           # device 0 owns this expert
52    y = np.maximum(0, W_e @ x_row)
53    outputs_dev0.append((y, src_row, 0))  # 0 is "I am dev 0"
54
55print("inbox[0] size:", len(inbox[0]),
56      "of", sum(len(b) for b in inbox))

Notice the only line that could not exist in real code: the block that builds device 0's expert dict by filtering. In real expert parallelism, that filter happens at process creation — each rank only ever allocates its own experts. There is no global dict, no other device's weights waiting to be filtered out.

Sanity check. Set $W = 1$. Every token's destination is device 0; the send dict has one entry; the all-to-all is a no-op; we recover the single-device MoE from section 5.1. Set $W = E$ (one expert per device): every all-to-all is fully cross-device and there are no local picks — the bandwidth tax is at its maximum.

PyTorch: all_to_all_single in Anger

Production code uses torch.distributed.all_to_all_single with split sizes. The module below is a complete, runnable expert-parallel MoE layer — short enough to read end-to-end, faithful enough to the real DeepSeek implementation that the structure transfers directly.

1import torch
2import torch.distributed as dist
3import torch.nn as nn
4import torch.nn.functional as F
5
6class ExpertParallelMoE(nn.Module):
7    """One expert-parallel MoE layer. World size W must divide num_experts."""
8    def __init__(self, d_model, d_ff, num_experts, k, ep_group):
9        super().__init__()
10        self.k = k
11        self.E = num_experts
12        self.group = ep_group
13        self.W = dist.get_world_size(ep_group)
14        self.rank = dist.get_rank(ep_group)
15        assert num_experts % self.W == 0
16        self.local_E = num_experts // self.W
17
18        # Only LOCAL experts are instantiated on this rank.
19        self.experts = nn.ModuleList([
20            nn.Sequential(nn.Linear(d_model, d_ff), nn.GELU(),
21                          nn.Linear(d_ff, d_model))
22            for _ in range(self.local_E)
23        ])
24        self.router = nn.Linear(d_model, num_experts, bias=False)
25
26    def forward(self, x):  # x: (T, D) — already a flat token slice
27        T, D = x.shape
28        logits = self.router(x)                          # (T, E)
29        topw, topi = logits.topk(self.k, dim=-1)         # (T, k)
30        gates = F.softmax(topw, dim=-1)                  # (T, k)
31
32        # Flatten (token, slot) → one row per (token, chosen expert).
33        flat_x   = x.unsqueeze(1).expand(T, self.k, D).reshape(-1, D)
34        flat_eid = topi.reshape(-1)                      # (T*k,)
35        flat_g   = gates.reshape(-1, 1)                  # (T*k, 1)
36        flat_dst = flat_eid // self.local_E              # owner rank
37
38        # --- Sort rows by destination rank so all-to-all can use counts ---
39        order = flat_dst.argsort()
40        send_x   = flat_x[order]
41        send_eid = flat_eid[order]
42        send_g   = flat_g[order]
43        send_counts = torch.bincount(flat_dst, minlength=self.W)
44
45        # --- Exchange counts so every rank knows how many rows it will receive ---
46        recv_counts = torch.empty_like(send_counts)
47        dist.all_to_all_single(recv_counts, send_counts, group=self.group)
48
49        # --- The two big all-to-alls: activations + their expert ids ---
50        recv_x = torch.empty(
51            (recv_counts.sum().item(), D),
52            dtype=x.dtype, device=x.device,
53        )
54        dist.all_to_all_single(
55            recv_x, send_x,
56            output_split_sizes=recv_counts.tolist(),
57            input_split_sizes=send_counts.tolist(),
58            group=self.group,
59        )
60        recv_eid = torch.empty(recv_counts.sum().item(), dtype=torch.long,
61                               device=x.device)
62        dist.all_to_all_single(
63            recv_eid, send_eid,
64            output_split_sizes=recv_counts.tolist(),
65            input_split_sizes=send_counts.tolist(),
66            group=self.group,
67        )
68
69        # --- LOCAL expert compute: one packed matmul per local expert ---
70        local_out = torch.empty_like(recv_x)
71        local_eid = recv_eid - self.rank * self.local_E  # 0..local_E-1
72        for e in range(self.local_E):
73            mask = local_eid == e
74            if mask.any():
75                local_out[mask] = self.experts[e](recv_x[mask])
76
77        # --- Combine = inverse all-to-all to send results home ---
78        return_x = torch.empty_like(send_x)
79        dist.all_to_all_single(
80            return_x, local_out,
81            output_split_sizes=send_counts.tolist(),
82            input_split_sizes=recv_counts.tolist(),
83            group=self.group,
84        )
85
86        # Unsort + apply gates + scatter-add back to per-token rows.
87        unsort = order.argsort()
88        out_flat = return_x[unsort] * flat_g
89        out = out_flat.view(T, self.k, D).sum(dim=1)
90        return out

Three subtleties that bite first-time EP implementers:

1. The counts exchange is mandatory. all_to_all_single requires the receiver buffer to be pre-sized. You cannot know your inbox size without first exchanging the per-source counts — that is why two collectives appear before the activations move. Skipping this step is the most common cause of segfaults in hand-written EP code.
2. Sorting by destination is the cheap optimization. all_to_all_single moves contiguous blocks. If you do not sort, you fall back to all_to_all (with explicit per-rank tensors) which is dramatically slower because it triggers many small NCCL sends instead of one large one. The argsort + argsort-of-argsort idiom is worth memorizing.
3. Autograd flows through the collective. NCCL's PyTorch bindings have proper backward implementations for all_to_all_single: the gradient of a dispatch is the corresponding combine, and vice versa. You write the forward; autograd produces a correctly-distributed backward at no additional code cost. This is why hand-rolling collectives in raw dist.send / dist.recv is almost never worth it.

What Changes at Massive Scale

At cluster scale, expert parallelism does not stand alone — it is one axis of a 3D or 4D parallelism mesh. DeepSeek-V3's reported training configuration is illustrative:

The cluster mesh is therefore $64 \times 16 \times 32 = 32{,}768$ ranks. The MoE layer's 256 routed experts are sharded across the 32-rank EP group: 8 experts per GPU. Every step, every EP group does its own pair of all-to-alls; the 64 data-parallel replicas do them independently in parallel.

Node-limited routing

Cross-node bandwidth (NVLink + InfiniBand combined) is typically 4–8× slower than intra-node NVLink. DeepSeek-V3 introduces node-limited routing: each token is restricted to experts living on at most $M = 4$ nodes out of the 8 nodes in its EP group. The router's top-$k$ is masked so the chosen experts span no more than $M$ nodes. This halves cross-node traffic with almost no quality loss — a beautifully practical compromise.

Capacity factor and token dropping

Even with auxiliary load balancing, the routed traffic is never perfectly uniform. To bound the receiver buffer at compile time, MoE implementations set a capacity factor $C$: each expert reserves space for at most $C \cdot (\text{average tokens per expert})$ rows per step. Tokens beyond the cap are dropped for that expert — they pass through the shared experts and the residual stream unmodified by the MoE block. DeepSeek typically uses $C \approx 1.0$ in training and bumps it to $C \approx 1.5$ for inference where latency matters more than throughput.

Communication-compute overlap

The bandwidth budget can be partly hidden. While device $r$ is running its local experts on the inbox, it can begin sending output rows back to their source devices for the combine — once any single expert finishes, its rows can start traveling. Real implementations pipeline this overlap explicitly. The shared experts play a second role here: they are running locally during the dispatch all-to-all, hiding 10–30% of the network latency behind useful compute.

Engineering Reality: Capacity, Overlap, and Topology

Expert parallelism is the easiest piece of MoE to implement at toy scale and the easiest to get wrong at production scale. Four failure modes are worth carrying in your head:

1. Imbalanced routing → straggler GPU. If one device owns the most popular experts, the entire EP group waits for it on every step. Symptoms: low SM utilization on most ranks, high on one; all-to-all latency climbing over training. Fix: aggressive load-balancing (chapter 6) and capacity factor > 1.
2. EP group placed across the wrong topology. If your 32-rank EP group spans 8 nodes instead of 4, every all-to-all is cross-node. Modern clusters give you NVLink within a node and InfiniBand across nodes — those are 10× apart in bandwidth. Place EP groups along NVLink-dense axes, not pipeline-dense ones.
3. Token-drop cascade in mixed precision. With $C = 1.0$ and BF16 accumulation, a tiny drift in router logits can flip a token from "just under the cap" to "dropped". Dropped tokens get zero gradient through the MoE for that step. Stack hundreds of such events per batch and you have a silent quality regression that no loss curve will catch — only eval metrics will. Periodically log per-expert drop rate.
4. All-to-all on the wrong stream. If the dispatch collective shares a CUDA stream with the local expert compute, they serialize even though they could overlap. Production EP code puts collectives on dedicated streams and uses CUDA events to synchronize only where strictly required. A 30% throughput swing is on the table.

The one sentence to carry forward: expert parallelism turns MoE's structural sparsity into a memory savings — and pays for it with two all-to-all collectives per layer per step. Everything else in massive MoE engineering — load balancing, capacity factors, node-limited routing, communication-compute overlap — exists to keep that all-to-all cheap.