Katkat
Back

Shard an LM Head and Combine Distributed Logits

MicrosoftSoftware Engineer · Onsite
HardSoftware EngineerOnsiteML System Design

Last updated

Shard an LM head by vocabulary or hidden width, distinguishing complete logits from partial dot products and exact distributed top-k from full-distribution sampling.

An LM head maps a hidden state through a weight matrix [H,V] to vocabulary logits. Compare sharding the weight by vocabulary and by hidden dimension. Explain what each rank computes, which reductions are required, and whether full-vocabulary logits must be gathered before top-k selection or sampling.

Constraints & Assumptions

State the incoming hidden-state layout and the requested output operation. Exact full-distribution sampling is different from top-k sampling; do not silently replace one with the other.

Clarifying Questions

Is the hidden vector replicated or sharded? Does the consumer need all logits, only the global top k, or an exact categorical sample? What output layout can downstream operations consume?

Follow-up questions

  • Why can local top-k candidates suffice for exact global top k? Why are they insufficient for unrestricted full-softmax sampling? Where does communication move if an earlier layer already outputs hidden shards?