Column-parallel linear layer splitting output features across tensor-parallel ranks

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. ...

April 19, 2025 · 7 min · Duo An
Megatron process groups carved from a DP, PP, and TP rank mesh

Megatron Internals I: Building the DP / TP / PP Process Groups

Megatron Internals I: Building the DP / TP / PP Process Groups Megatron-LM’s first trick is not tensor-parallel matmul. It is rank bookkeeping. Before the model runs, every process must know the small set of peers it will communicate with for tensor parallelism, pipeline parallelism, data parallelism, embeddings, and optimizer state. The original Megatron-LM paper introduced intra-layer tensor parallelism for transformer training. The later Megatron-LM scaling paper put tensor parallelism, pipeline parallelism, and data parallelism into one training system. This post is the control plane underneath that system: how a flat list of ranks becomes a DP / PP / TP mesh. ...

April 12, 2025 · 8 min · Duo An