Mixture of experts and parallelism interview questions

An MoE layer swaps one feed-forward block for many experts and a small router that sends each token to the top few, often 2.

Parameters grow, but each token uses a slice: Mixtral 8×7B has about 47B parameters and uses about 13B per token. Expert parallelism places experts on different GPUs and ships tokens to them in an all-to-all exchange, so the costs become communication and load balance.

Tensor parallelism splits each weight matrix across GPUs, so every layer needs a collective step: fast only over very fast links, usually inside one machine. Pipeline parallelism gives each GPU a run of consecutive layers and passes activations along: little communication, but GPUs sit idle in the pipeline bubble unless the batch is split into micro-batches. Large runs combine both with data parallelism.

All-reduce sums gradients across data-parallel replicas, and partial results in tensor parallelism. All-gather and reduce-scatter are its two halves: sharded training such as ZeRO or FSDP gathers each layer's weights just before use and reduce-scatters its gradients after. All-to-all routes tokens to experts in MoE. Each costs bytes over links, so a good layout puts the chattiest collectives on the fastest links.

Count every live tensor.

# Y = X @ W; X (m, k), W (k, n); bf16 = 2 bytes
fwd = 2 * (m * k + k * n + m * n)
# m = k = n = 65536: 8 GiB each, 24 GiB total

Training adds gradients, and mixed-precision Adam needs about 16 bytes per parameter with its FP32 copies. Tensor parallelism splits W itself, so each GPU holds a slice and every step needs an all-reduce or all-gather, best over NVLink. Pipeline parallelism gives whole layers to stages and passes only activations, fine across nodes but with bubbles. It can't split a single matmul too big for one GPU; tensor parallelism can.

Combine parallelisms: expert parallelism spreads experts across GPUs with all-to-all token dispatch, tensor parallelism stays inside a node, and pipeline and data parallelism span nodes. The router needs a load-balancing loss and a capacity factor, or a few experts take most tokens; log per-expert load from day one.

All-to-all is the bottleneck, so keep expert groups on the fastest links and overlap dispatch with compute. Add router z-loss for stability, checkpoint often, and validate the data mix on small runs first.

Quantization stores weights or activations in fewer bits. Pruning removes weights, and only structured sparsity such as 2:4 speeds up GPUs. Distillation trains a small student on a big teacher's outputs. Fusion merges kernels to cut memory traffic.

Tensor parallelism shards each layer's matrices, with all-reduces every layer, so it wants NVLink. Pipeline parallelism assigns layers to stages, sends only activations, and suffers bubbles that micro-batches shrink.

A transformer stacks blocks of attention, which mixes tokens, and an MLP, each with a residual. MoE swaps the MLP for many experts and a router picking a few per token.

Training runs whole sequences in parallel under the causal mask; MoE adds a load-balancing loss so experts don't collapse. Inference prefills the prompt, then decodes token by token with a KV cache; every expert stays in memory though each token uses few.

Vocabulary sharding: each device holds a slice of rows and computes logits for its slice only. Softmax never needs the full vector: all-reduce the max, then the sum of exponentials. Hidden sharding: each device holds a slice of the hidden width and produces partial logits over the whole vocabulary, which must be all-reduced, a vocabulary-sized vector per token.

import torch
import torch.distributed as dist

def dist_topk(local, k, offset):
    # local: [B, V_shard] logits on this rank
    v, i = local.topk(k, dim=-1)
    i = i + offset  # to global vocabulary ids
    n = dist.get_world_size()
    vs = [torch.empty_like(v) for _ in range(n)]
    ix = [torch.empty_like(i) for _ in range(n)]
    dist.all_gather(vs, v)
    dist.all_gather(ix, i)
    v, i = torch.cat(vs, -1), torch.cat(ix, -1)
    top, pos = v.topk(k, dim=-1)
    return top, i.gather(-1, pos)

It's exact: every global top-k entry is in its own shard's top k.

ZeRO removes data parallelism's redundancy: every GPU normally holds full parameters, gradients and Adam states. ZeRO-1 shards the optimiser states, ZeRO-2 also shards gradients, using reduce-scatter instead of all-reduce, and ZeRO-3 also shards parameters, all-gathering each layer's weights just before its forward and backward, then freeing them. FSDP is PyTorch's ZeRO-3-style implementation, sharded per wrapped module.

The price is communication, so overlap it: prefetch the next layer's all-gather while this one computes, and reduce-scatter each layer's gradients as soon as its backward finishes.

Pipeline parallelism splits layers into stages and streams microbatches through them, leaving a fill-and-drain bubble. A GRPO step generates a group of rollouts per prompt, then trains on group-normalised rewards. Generation is decode, which pipelines poorly, so it often runs on a separate inference engine.

Overlapping next-step rollouts with this step's training keeps GPUs busy, but those rollouts come from an older policy version. Tag each rollout with its version, cap the lag at a step or two, and rely on the loss's importance ratio and clipping.