Megatron Internals II: Column/Row Parallel Linear and Vocab Parallel Embedding
Megatron Internals II: Column/Row Parallel Linear and Vocab Parallel Embedding Tensor parallelism is not “split every tensor somehow.” In Megatron, it is a small set of layer contracts: which dimension is local, which collective completes the dense math, and which gradient path communicates. The original Megatron-LM paper is still the cleanest starting point: split transformer matrix multiplies so each GPU does useful dense GEMM, then communicate only where the algebra requires it. This post walks the implementation-level contracts behind ColumnParallelLinear, RowParallelLinear, VocabParallelEmbedding, and parallel cross entropy. ...