Google DeepMind Distributed Training & Parallelism interview questions

On a TPU pod you write one program, annotate how each array is sharded across a named mesh, and the compiler inserts and schedules every collective. What that removes from the engineer's job, what the torus changes about layouts, where Pathways fits, and what a compiler-owned schedule costs.

Not the data, not the weights: the gradients, once per step, in a ring that moves almost twice the model's size through every GPU. The derivation, the per-step byte count for an 8B model, and why it still hides behind the backward pass.

Three ways to split a training job, one table, and the rule that places each of them: activations on NVLink, gradients on the fabric, stage boundaries in between. With the byte counts that justify the placement.

The one utilization number that cannot be gamed by recompute: model FLOPs over hardware peak, derived from a step time in four lines, with the band frontier labs land in and an itemized list of where the other 60% goes.

Sixteen bytes per parameter is the bill for mixed-precision Adam. ZeRO pays it in three installments: optimizer state, then gradients, then the weights themselves. The per-rank bytes for a 70B at 8, 16, 32 and 64 GPUs, and what each stage adds to the wire.

Same algorithm, two implementations, one of which now lives inside PyTorch and composes with everything else there. The differences that decide it in 2026: per-parameter sharding, DTensor, torch.compile, and the two features DeepSpeed still owns outright.

Sixteen bytes per parameter says 1.6 TB of state against a 640 GB node, so the model spans nodes before the first token. The arithmetic that sizes the fleet, the layout that puts each axis on the right link, and the two numbers that decide FSDP against pipeline across nodes.

Two phases, n minus one steps each, one slice per step: the derivation that gives 2(n−1)/n times the buffer per rank, the argument that no algorithm can beat it on bytes, and the latency term that makes NCCL switch to trees at scale.

With p stages and m micro-batches the pipeline idles (p−1)/m of its compute time, and the obvious fix, more micro-batches, buys back throughput with activation memory. 1F1B caps that memory at p micro-batches instead of m. The derivation, the numbers for PP16, and the schedules that go further.

Four all-reduces per transformer block, on activations, on the critical path. Inside the NVLink domain they cost 8% of a step; across the NIC they cost more than the step itself. The derivation for a 70B, and the second reason TP stops that has nothing to do with the network.

Mathematically identical for a mean loss and a stateless model, and different in three ways that matter to infrastructure: activation memory, communication frequency, and any layer that looks at the batch. With the byte counts and the DDP call that makes it work.

Full recomputation costs a third more compute to cut activation memory by an order of magnitude. The per-layer arithmetic that says whether you need it at all, the selective variant that costs almost nothing, and the point at which a smaller micro-batch is the cheaper way out.

At 128k tokens one layer's activations are 73 GB per sequence and attention grows with the square of the length. The arithmetic that forces the sequence onto sixteen GPUs, what a ring of KV chunks costs per layer, why the communication hides, and where Ulysses and Megatron sequence parallelism fit around it.

Every MoE layer routes each token to eight experts that live on other GPUs, twice per layer, forward and back. The bytes per rank per layer for DeepSeek-V3, why that all-to-all runs close to one-to-one against compute, and why a single hot expert stalls every rank waiting at the combine.

A router that sends 90% of tokens to 26 of 256 experts makes those ranks do nine times their share while everyone else waits. The arithmetic of the imbalance, what a capacity factor drops on the floor, why the auxiliary loss fights the model, and the bias-only balancing DeepSeek-V3 used instead.

Tensor parallel 8, pipeline 16, data parallel 128: the layout Meta trained Llama 3 405B on, rebuilt from the memory floor, the compute budget and the bandwidth of each link. Why the axes go in that order, what each one costs per step, and the bubble that the schedule has to hide.