Column parallelism: partition output features without summing them

Partition W by its output columns, giving rank i a matrix W_i with shape [H,O/P] for an even P-way split. Every rank needs the input features required by that multiplication. With replicated X, each rank computes its own output-feature slice without summing other ranks’ forward products. Concatenating those slices in the correct global feature order reconstructs Y. Calling an output slice partial can be ambiguous: it is a complete value for selected features, not a pending sum contribution for every feature.

Bias is partitioned along the same output-feature axis. Rank i adds b_i exactly to its own complete output slice. An elementwise activation can remain local because it does not mix different output coordinates. A later operation that needs the entire hidden vector may require gathering or a compatible partitioned algorithm. The useful design often keeps the intermediate sharded until a complementary row-parallel layer consumes it, avoiding a needless full reconstruction between stages.

Each local output gradient produces a local weight gradient using the same X and its local dY_i. The input gradient, however, combines contributions from every output-feature partition. Under a replicated-input layout, the complete dX is the sum of local dY_i W_i-transpose products. Omitting that sum can leave the forward output correct while training the earlier layer incorrectly. If the desired backward layout shards tokens, a reduce-scatter may deliver the needed summed pieces instead of a replicated all-reduce.

The communication schedule therefore depends on input and output placements, not just the phrase column parallel. A sequence-sharded input may be gathered before the multiplication, while an output may optionally be gathered afterward. Autograd mappings determine where reductions occur. Inspect the versioned library contract and actual layout transitions. State exactly which tensor is replicated, sharded, or a pending sum before quoting a forward/backward collective count; otherwise two correct implementations can appear to disagree only because their boundaries differ.

A shard with the right shape can still contain the wrong output columns. Test known matrices that distinguish every feature, then reconstruct in documented rank and feature order. Include nonzero bias so a missing slice or accidental duplicate becomes visible. For training, compare local weight gradients and the assembled input gradient against a dense reference. Numerical tolerance should reflect dtype and reduction order; shape agreement alone is weak evidence of mathematical equivalence.

The matrix lab shows one fixed two-dimensional family rather than arbitrary tensor shards. Reset, set diagonal scale a=1 and shear b=1, predict determinant one, then Step. The selected transform is [[1,1],[0,1]], with readouts describing that final matrix even during interpolation. It does not multiply distributed tensors or reconstruct column shards. Use it only to inspect an explicit linear map while the worked matrix calculation establishes the partitioned layer’s actual values.