Shard an LLM's embedding table and LM head across devices

A language model's two vocabulary-sized layers no longer fit on one device: the input embedding table [V, H] and the output layer (the LM head) [H, V]. Split each across P devices with tensor parallelism. For each, compare splitting along the vocabulary with splitting along the hidden dimension, and count memory, communication and temporary buffers for a concrete B, V, H and P. The interviewer takes the two layers in order.

Part 1: the input embedding table

A model's input embedding table has shape [V, H] (vocabulary by hidden width) and is too large for one device. Compare sharding it along the vocabulary with sharding it along the hidden dimension: how a token lookup works in each, what communication is needed, how load balances across devices, and how the output feeds the next layers.

A strong answer covers:

Vocabulary sharding: each device holds rows [V/P, H]. With replicated token ids, each device looks up the tokens in its range, writes zeros for the rest, and an all-reduce (or reduce-scatter) gives full vectors. With partitioned ids, tokens are routed to the device that owns them (all-to-all) and results routed back.
Hidden sharding: each device holds [V, H/P], every lookup is local, and the output is naturally sharded along H; an all-gather is needed only if the next layer wants full vectors.
Skew: frequent tokens concentrate on one vocabulary shard, creating hot devices under routing; hidden sharding is balanced by construction.
Per-device memory saved versus temporary buffers and communication added, counted for a concrete B, V, H, P.

Part 2: the LM head and its logits

A language model's output layer (the LM head) maps a hidden state h of size H through a weight matrix W of shape [H, V] to logits over a vocabulary of size V. W is split across P devices. Compare splitting it along the vocabulary with splitting it along the hidden dimension: what each device computes, which collective operations combine the results, and whether full-vocabulary logits must be gathered before top-k selection or sampling.

A strong answer covers:

Vocabulary sharding: each device holds [H, V/P] and, with h replicated, computes complete logits for its slice of the vocabulary. No reduction is needed for those logits; gathering them all costs B * V values.
Hidden sharding: each device holds [H/P, V] and computes partial dot products for every token; an all-reduce (or reduce-scatter to a vocabulary-sharded layout) is needed before any logit is final.
Top k from vocabulary shards: each device's local top k is a superset of what it contributes to the global top k, so gathering P * k candidates is exact.
Exact sampling from the full softmax cannot use only local top k: it needs a distributed softmax (global max and sum via all-reduce, then sampling by shard, for example Gumbel-max per shard), never silently replaced by top-k sampling.

Follow-ups

The most frequent 1% of tokens make up half the traffic on vocabulary shards. What do you do?
When can the hidden shards stay distributed into the first transformer block?
Why are local top-k candidates enough for an exact global top k?
Why are they not enough for unrestricted sampling from the full distribution?