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