Senior interview questionAsked at Mistral AI · interview write-up
Explain FSDP and ZeRO stages 1, 2 and 3. How do you overlap computation with communication?
Model answer
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.
Practise the rest
10 MoE & parallelism questions, 492 in all, each with a model answer. Free.
More from Mistral AI interview questions.
More MoE & parallelism interview questions
- Explain mixture of experts, and what expert parallelism is.
- Compare tensor parallelism and pipeline parallelism.
- Which collective operations does distributed training rely on, and where does each show up?
- Estimate the GPU memory a very large matmul needs, then compare pipeline and tensor parallelism for it.
- Design a pretraining pipeline for a mixture-of-experts model that scales.
- Cover quantization, pruning, distillation and kernel fusion, then compare tensor and pipeline parallelism.