Row parallelism: sum partial products before adding bias

Split X into feature blocks X_i with shape [M,H/P], and split W into matching row blocks W_i shaped [H/P,O]. Their local product Z_i has shape [M,O], the same shape on every rank. Each Z_i is missing the other ranks’ feature contributions. It is therefore a pending sum rather than a complete replica. Reconstructing Y requires a sum of Z_i, then the appropriate bias. Concatenating these products would multiply the apparent output width instead of restoring the layer.

A preceding column-parallel expansion can leave its intermediate output features already partitioned in the form this layer needs. That pairing permits local intermediate activation followed by the row-parallel multiplication, without gathering that intermediate solely to repartition it. The exact forward collective can yield a replicated complete output or reduce-scattered token shards, depending on the next layer. Preserve the mathematical sum while choosing the placement that avoids unnecessary traffic for the surrounding computation.

In general, applying a nonlinear activation to each Z_i and then adding does not equal activating their sum. A negative partial contribution can cancel a positive one before the activation, while separate activation would remove that cancellation. The same issue appears with normalization across output features: a local incomplete sum is not a valid complete token vector. Make the pending-reduction status explicit so later code cannot accidentally treat a same-shaped local tensor as a fully computed value.

A replicated output bias added before summing would appear P times. The simple safe convention is to add it after reduction. Some fused implementations return bias separately for a later fused operation; others use carefully equivalent distribution rules. Read the actual API instead of adding a second bias because a local output initially appears unbiased. Backward must likewise respect the decomposition: a full upstream output gradient supplies each input shard’s gradient through its local W_i, with further communication determined by how that upstream gradient is placed.

Use test matrices whose local products have different signs and magnitudes. A test with all zeros can hide missing reductions, and a zero bias cannot reveal duplicate bias addition. Compare a known dense output, then compare backward gradients under one declared loss normalization. An implementation can have the right output shape while silently applying a nonlinear operation too early. Keep dtype, parameter initialization, and the exact output-feature order fixed when comparing layouts.

The function lab gives a direct nonlinear counterexample. Reset, set a=1 and x=1, predict y=1, then Step. Change x to 2 and observe y=4; two copies of f(1) sum to 2, while f(1+1)=4. The input control remains in range and edits retain time, which does not change these selected readouts. This one quadratic is not a row-parallel layer, but it demonstrates why a nonlinear map cannot generally move across a sum without changing the function.