Vocabulary shards: look up locally and normalize globally

For an even vocabulary split, each rank owns a contiguous or explicitly mapped subset of rows in an embedding table shaped[V,H]. A rank tests whether each global token ID is in its subset, converts owned IDs to valid local indices, and contributes the corresponding H-vector. Nonowners contribute a neutral zero vector in one common reduction strategy. Summing those masked contributions reconstructs the embedding because exactly one rank owns a valid ID. Another supported layout may return sequence shards instead, changing the collective but not the row identity.

Do not index a local table with an unmasked global ID and then try to discard the result afterward; the invalid access has already happened. Repeated occurrences of one token accumulate gradients into the same owned row. Padding IDs, out-of-range IDs, and a vocabulary padded for divisibility need explicit rules. If input and output embeddings share one logical parameter, the implementation must preserve that tying and combine all intended gradient contributions rather than accidentally create two unrelated copies.

For one token’s output logits, first find the global maximum over valid vocabulary entries. Every rank subtracts that maximum from its owned logits, sums their exponentials locally, and combines the local sums globally. The target logit is obtained from its unique owner or an equivalent distributed selection. Cross-entropy then uses the same global denominator. Local softmax would assign too much probability to the local shard because it ignores competing tokens elsewhere.

Subtracting the global maximum improves numerical stability without changing probabilities. A padded vocabulary entry that is not a real class must be excluded consistently; otherwise it adds false denominator mass. Ignored target tokens must follow the loss’s mask and normalization, not trigger a bogus target owner. Gradients for valid logits are their global probabilities minus the target indicator, with whatever loss scaling was defined. Fused loss implementations can use different reduction schedules, so describe the mathematical quantities before claiming a fixed count of all-reduces.

Distributed vocabulary loss can avoid materializing every full logit vector on every rank, but it still needs global normalization information and correct backward ownership. Parameter savings do not imply zero communication or a globally sharded hidden activation at every boundary. Check the selected library’s input/output placements and tied-weight options. For testing, choose logits where the largest value and target live on different shards, as well as targets at shard boundaries and near a padded vocabulary tail.

The sampling lab illustrates a normalized distribution only. Reset, set p=0.5 and n=4, predict binomial masses proportional to 1,4,6,4,1 across five outcomes, then Step. Their common denominator is 16. The model is an exact binomial distribution, not vocabulary softmax, embeddings, logits, or a collective. Its limited point is that probabilities share one normalization across the defined outcome space; the worked example supplies the different exponential denominator used by language-model loss.