Vocabulary-Parallel Embedding and Cross-Entropy
Hard · Solved in PyTorch · Machine learning coding practice
Problem
Megatron shards the vocabulary dimension across P ranks at both ends of the model, so neither the embedding table nor the logits are ever materialized in full on one rank (Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism (Shoeybi et al., 2019)). Rank s owns the contiguous vocabulary range [sV/P,\ (s+1)V/P), where V is the vocabulary size and V/P is the block size.
token_ids holds integer ids of shape (B, S) with 0 \leq \text{id} < V. embedding_shards holds the sharded table of shape (P, V/P, D), so embedding_shards[s] is rank s's rows in local order. logit_shards holds the sharded output logits of shape (P, N, V/P) for N predicted positions, split over the same vocabulary blocks. targets holds integer ids of shape (N,). Do not modify the inputs.
Topics: megatron, tensor-parallel, cross-entropy, numerical-stability, pytorch, numpy
The full statement, worked examples, hints and the test suite are available once you sign in. You can then solve Vocabulary-Parallel Embedding and Cross-Entropy in the browser and run it against the tests.