System Design · Microsoft · Hard
Requirements This is an open-ended design discussion from an AI infrastructure onsite: shard the input embedding and the output (LM head) layers of a large language model across multiple GPUs. The input embedding matrix has shape $$[V, H]$$, and the vocabulary is so large that the matrix does not fit on a single GPU. Spread the weights across GPU ranks. Support embedding lookup when any given input token may live on any rank. Explain how requests get to the rank that owns…