Shard a Large Input Embedding Table Across Devices

An LLM's input embedding matrix has shape [V,H] and is too large for one device. Compare sharding along the vocabulary dimension with sharding along the hidden dimension. Explain lookup, communication, load balance, and how the output connects to later layers.

Constraints & Assumptions

V is vocabulary size and H is hidden width. State the number of devices and the incoming token/output layout. A sharding choice saves parameter memory but does not automatically eliminate communication or temporary buffers.

Clarifying Questions

Are token IDs replicated or partitioned across ranks? Does the next layer expect full or sharded hidden states? Is token traffic skewed? Are embedding and output weights tied?