LLM Training

Megatron Internals III: Mixed Precision, Loss Scaling, and Grad Clipping Mixed precision training is not just “turn on fp16.” It is a state machine: fast low-precision tensors do forward and backward, stable fp32 tensors receive optimizer updates, and every rank agrees whether the step is valid before any shard changes. The classic recipe comes from Micikevicius et al., Mixed Precision Training: use lower precision where hardware is fast, keep fp32 master weights for updates, and use loss scaling when fp16 gradients underflow. Megatron wraps that recipe around tensor, pipeline, and data parallelism. ...

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

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

Data Parallelism: From Parameter Server to Ring All-Reduce Data parallelism is the default scaling move because it preserves the model program. Every rank owns the same model, sees different examples, computes gradients, and applies the same update. The algorithm is simple. The system is not. The whole post is about removing one bottleneck: do not push all gradient traffic through one server when every GPU could be moving bytes at the same time. ...

Pipeline Parallelism from First Principles: Why GPipe Split the Batch Pipeline parallelism starts with a useful accident: a Transformer is already a chain. If one GPU cannot hold the whole stack, put consecutive blocks on consecutive GPUs and send activations across stage boundaries. That solves capacity. It does not solve throughput. A naive layer split turns an expensive cluster into a queue where most devices wait. GPipe’s contribution was to make the queue busy by slicing the mini-batch into micro-batches and recomputing activations instead of storing everything (arXiv:1811.06965). ...