Shard a Large Input Embedding Table Across Devices

Embedding Sharding Strategies: Vocabulary vs. Hidden Dimension Problem Restatement You are given an embedding table with shape $$[V, H]$$, where $$V$$ is the vocabulary size and $$H$$ is the hidden dimension. This matrix is too large to fit on a single device. You must partition it across multiple devices and implement a forward pass: given a batch of token IDs, produce the corresponding hidden vectors and feed them into the next layer. Compare two orthogonal partitioning…