Sequence-parallel normalization: shard tokens, retain their hidden features

For a token with H hidden coordinates, LayerNorm computes mean and variance over those coordinates, then applies learned per-feature scale and bias. If a rank owns whole token rows, it has the needed statistics locally. Splitting the sequence into token chunks therefore does not require a global reduction of means across different tokens. Combining token means would normalize a different object. RMSNorm has its own statistic and should not be silently substituted for LayerNorm in an equivalence test.

A hidden-feature-sharded token is different: no rank has the whole normalized axis, so a correct implementation needs the appropriate distributed statistics or a redistribution. Merely naming a tensor sequence-parallel does not establish its actual axes. Write the global shape, local shape, and placement at each boundary. With batch, sequence, and hidden dimensions, library defaults about which axis is sequence must match the program’s convention; a transpose can keep element counts unchanged while moving the wrong axis.

LayerNorm scale and bias parameters are typically replicated across sequence-sharded ranks in this design. Each rank’s local tokens contribute different parameter gradients, so those contributions must be summed over the relevant group before the intended optimizer update. Whether a framework performs this in a particular mapping or gradient-finalization phase is an implementation boundary. Do not independently update replicated parameters from different local token gradients and expect them to remain equal. Loss averaging must be applied once according to the global valid-token convention.

Sequence-sharded activations can be gathered when a tensor-parallel linear layer needs more input tokens, and complete output contributions can be reduce-scattered back into token shards. Such transitions explain activation-memory savings without treating every layer as permanently sequence-sharded. Dropout adds a random-number mapping requirement: matching a seed on every rank can accidentally repeat masks rather than reproduce the intended global mask. Start equivalence checks without dropout, then test an explicit consistent global randomness policy.

A token that is padding may still pass through a normalization kernel, but its contribution to the training objective must follow the defined masking rule. Shard sizes that do not divide the sequence require a supported uneven or padded layout and correct global token identities. PyTorch 2.8 marks some uneven DTensor sharding behavior experimental, so do not assume every module supports it automatically. Verify ε, variance convention, dtype, learned parameters, and custom parameter initialization when comparing replicated normalization modules.

The function lab can illustrate sensitivity in one fixed map only. Reset, set a=1 and x=1, predict y=1 and tangent slope 2, then Step. Change x to 0.5 to predict y=0.25 and slope 1. This is not LayerNorm, a variance calculation, or a sequence-parallel kernel. It reminds you to name which input coordinates a statistic depends on; the lesson’s formula and token-by-token example supply the actual normalization semantics.